Implement max block for token scope

README.md

@@ -77,8 +77,9 @@ The building blocks of a priompt prompt are:
 3. `<empty>`: for specifying empty space, useful for reserving tokens for generation.
 4. `<capture>`: capture the output and parse it right within the prompt.
 5. `<isolate>`: isolate a section of the prompt with its own token limit. This is useful for guaranteeing that the start of the prompt will be the same for caching purposes. it would be nice to extend this to allow token limits like `100% - 100`.
-6. `<br/>`: force a token break at a particular location, which is useful for ensuring exact tokenization matches between two parts of a prompt (e.g. when implementing something like speculative edits).
-7. `<config>`: specify a few common configuration properties, such as `stop` token and `maxResponseTokens`, which can make the priompt dump more self-contained and help with evals.
+6. `<max>`: specify a token limit for a section of the prompt, but unlike `<isolate>`, the inner scopes participate in the global priority calculation. This allows you to limit the tokens used by a section while still allowing high-priority content within that section to be prioritized globally.
+7. `<br/>`: force a token break at a particular location, which is useful for ensuring exact tokenization matches between two parts of a prompt (e.g. when implementing something like speculative edits).
+8. `<config>`: specify a few common configuration properties, such as `stop` token and `maxResponseTokens`, which can make the priompt dump more self-contained and help with evals.
 
 You can create components all you want, just like in React. The builtin components are:
 
@@ -97,7 +98,7 @@ You can create components all you want, just like in React. The builtin componen
 
 A few things that would be cool to add:
 
-1. A `<max>` block: specify a `limit` on the number of tokens within a scope, but unlike `<isolate>`, include the inner scopes in the global priority calculation.
+1. ~~A `<max>` block: specify a `limit` on the number of tokens within a scope, but unlike `<isolate>`, include the inner scopes in the global priority calculation.~~ ✅ **IMPLEMENTED**
 2. Performance-optimized rendering of big trees: minimizing time spent tokenizing is part of it, but part of it is also working around JavaScript object allocation, and it is possible that writing the entire rendering engine in Rust, for example, would make it a lot faster.
 
 ## Caveats

examples/max-block-example.tsx

@@ -0,0 +1,33 @@
+// Example demonstrating the <max> block functionality
+
+import { render } from '@anysphere/priompt';
+
+function ExampleMaxBlock() {
+  return (
+    <>
+      <scope p={100}>This is high priority content outside the max block.</scope>
+
+      <max tokenLimit={50}>
+        <scope p={200}>This is the highest priority content inside max block.</scope>
+        <scope p={150}>This is medium priority content inside max block.</scope>
+        <scope p={50}>This is lower priority content inside max block.</scope>
+      </max>
+
+      <scope p={125}>This is medium-high priority content outside the max block.</scope>
+    </>
+  );
+}
+
+// How this works:
+// 1. The <max> block limits its content to 50 tokens
+// 2. BUT the priorities inside the max block (200, 150, 50) participate in the global priority calculation
+// 3. So if we have limited global tokens, the content with priority 200 will be included first (highest globally)
+// 4. Then priority 150, then priority 125 (outside max), then priority 100 (outside max), then priority 50 (inside max)
+// 5. However, if the max block's 50 token limit is reached, lower priority content within it will be excluded
+
+// This is different from <isolate> which would:
+// 1. Render its contents separately with its own token limit
+// 2. NOT participate in global priority calculation
+// 3. Be included/excluded as a whole unit based on the isolate block's own priority
+
+export default ExampleMaxBlock;

priompt/src/base.test.tsx

@@ -272,3 +272,119 @@ describe("config", () => {
     expect(rendered.config.maxResponseTokens).toBe("tokensReserved");
   });
 });
+
+describe("max", () => {
+  function Max(
+    props: PromptProps<{ useMax: boolean; tokenLimit: number }>
+  ): PromptElement {
+    if (props.useMax) {
+      return (
+        <>
+          <max p={props.p} prel={props.prel} tokenLimit={props.tokenLimit}>
+            {props.children}
+          </max>
+        </>
+      );
+    } else {
+      return (
+        <>
+          <scope p={props.p} prel={props.prel}>
+            {props.children}
+          </scope>
+        </>
+      );
+    }
+  }
+
+  function TestMaxVsScope(props: PromptProps<{ useMax: boolean }>): PromptElement {
+    return (
+      <>
+        This is the start of the prompt.
+        <Max tokenLimit={50} useMax={props.useMax}>
+          <scope prel={-100}>High priority content that should always be included.</scope>
+          <scope prel={-200}>Medium priority content that might be excluded.</scope>
+          <scope prel={-300}>Low priority content that should be excluded when token limit is reached.</scope>
+        </Max>
+        <scope p={1000}>This high priority content outside max block should be included.</scope>
+      </>
+    );
+  }
+
+  it("should enforce token limit while participating in global priority calculation", async () => {
+    const tokenizer = await getTokenizerByName_ONLY_FOR_OPENAI_TOKENIZERS("gpt-4");
+
+    // Test with max block
+    const renderedWithMax = await render(
+      <TestMaxVsScope useMax={true} />,
+      {
+        tokenLimit: 1000,
+        tokenizer,
+      }
+    );
+
+    // Test without max block (regular scope)
+    const renderedWithScope = await render(
+      <TestMaxVsScope useMax={false} />,
+      {
+        tokenLimit: 1000,
+        tokenizer,
+      }
+    );
+
+    expect(isPlainPrompt(renderedWithMax.prompt)).toBe(true);
+    expect(isPlainPrompt(renderedWithScope.prompt)).toBe(true);
+
+    const maxPromptText = renderedWithMax.prompt as string;
+    const scopePromptText = renderedWithScope.prompt as string;
+
+    // Both should include the high priority external content
+    expect(maxPromptText).toContain("This high priority content outside max block should be included");
+    expect(scopePromptText).toContain("This high priority content outside max block should be included");
+
+    // Both should include the high priority content within the max/scope block
+    expect(maxPromptText).toContain("High priority content that should always be included");
+    expect(scopePromptText).toContain("High priority content that should always be included");
+
+    // The max block should enforce its token limit and potentially exclude lower priority content
+    // while the scope block might include more content if the global token limit allows
+    console.log("Max block rendered:", maxPromptText);
+    console.log("Scope block rendered:", scopePromptText);
+  });
+
+  function TestGlobalPriorityParticipation(): PromptElement {
+    return (
+      <>
+        <scope p={100}>Very high priority content outside max.</scope>
+        <max tokenLimit={30}>
+          <scope p={200}>Highest priority content inside max.</scope>
+          <scope p={50}>Lower priority content inside max.</scope>
+        </max>
+        <scope p={150}>High priority content outside max.</scope>
+      </>
+    );
+  }
+
+  it("should participate in global priority calculation", async () => {
+    const tokenizer = await getTokenizerByName_ONLY_FOR_OPENAI_TOKENIZERS("gpt-4");
+
+    const rendered = await render(
+      <TestGlobalPriorityParticipation />,
+      {
+        tokenLimit: 100, // Limited global token limit
+        tokenizer,
+      }
+    );
+
+    expect(isPlainPrompt(rendered.prompt)).toBe(true);
+    const promptText = rendered.prompt as string;
+
+    // The content with priority 200 should be included (highest priority globally)
+    expect(promptText).toContain("Highest priority content inside max");
+
+    // The content with priority 150 should be included
+    expect(promptText).toContain("High priority content outside max");
+
+    console.log("Global priority test rendered:", promptText);
+  });
+});
+

priompt/src/lib.ts

@@ -391,6 +391,29 @@ export function createElement(tag: ((props: BaseProps & Record<string, unknown>)
 					absolutePriority: (typeof props.p === 'number') ? props.p : undefined,
 					relativePriority: (typeof props.prel === 'number') ? props.prel : undefined,
 					name: (props !== null && typeof props.name === 'string') ? props.name : undefined,
+					onEject: props.onEject,
+					onInclude: props.onInclude,
+				};
+			}
+		case 'max':
+			{
+				// must have tokenLimit
+				if (!props || typeof props.tokenLimit !== 'number') {
+					throw new Error(`max tag must have a tokenLimit prop, got ${props}`);
+				}
+
+				return {
+					type: 'scope',
+					children: [{
+						type: 'max',
+						tokenLimit: props.tokenLimit,
+						children: children.flat(),
+					}],
+					absolutePriority: (typeof props.p === 'number') ? props.p : undefined,
+					relativePriority: (typeof props.prel === 'number') ? props.prel : undefined,
+					name: (props !== null && typeof props.name === 'string') ? props.name : undefined,
+					onEject: props.onEject,
+					onInclude: props.onInclude,
 				};
 			}
 		case 'capture':
@@ -1157,6 +1180,7 @@ function normalizePrompt(elem: PromptElement): NormalizedPromptElement {
 				case 'config':
 				case 'capture':
 				case 'isolate':
+				case 'max':
 				case 'breaktoken':
 				case 'image':
 				case 'empty': {
@@ -1408,6 +1432,43 @@ async function renderWithLevelAndCountTokens(elem: NormalizedNode[] | Normalized
 				config: emptyConfig(),
 			}
 		}
+		case 'max': {
+			// Unlike isolate, max blocks participate in global priority calculation
+			// but still enforce a local token limit by rendering children with the smaller limit
+			const maxResult = await renderWithLevelAndCountTokens(elem.children, level, tokenizer);
+
+			// If the result exceeds our token limit, we need to find a higher priority cutoff
+			// that keeps us within the limit
+			if (maxResult.tokenCount > elem.tokenLimit) {
+				// Find the minimum priority level that keeps us under the token limit
+				// We'll use binary search approach similar to the main renderer
+				const childPriorityLevels = new Set<number>();
+				computePriorityLevels(elem.children, level, childPriorityLevels);
+				const sortedChildLevels = Array.from(childPriorityLevels).sort((a, b) => b - a);
+
+				// Binary search for the right priority cutoff
+				let left = 0;
+				let right = sortedChildLevels.length;
+				let bestResult = maxResult;
+
+				while (left < right) {
+					const mid = Math.floor((left + right) / 2);
+					const testLevel = sortedChildLevels[mid];
+					const testResult = await renderWithLevelAndCountTokens(elem.children, testLevel, tokenizer);
+
+					if (testResult.tokenCount <= elem.tokenLimit) {
+						bestResult = testResult;
+						right = mid;
+					} else {
+						left = mid + 1;
+					}
+				}
+
+				return bestResult;
+			}
+
+			return maxResult;
+		}
 		case 'chat': {
 			const p = await renderWithLevelAndCountTokens(elem.children, level, tokenizer);
 			if (isChatPrompt(p.prompt)) {
@@ -1667,6 +1728,11 @@ function renderWithLevelAndEarlyExitWithTokenEstimation(elem: PromptElement, lev
 				emptyTokenCount += elem.cachedRenderOutput.tokensReserved;
 				return;
 			}
+			case 'max': {
+				// For max blocks, we render children normally but enforce the token limit locally
+				renderInPlace(elem.children);
+				return;
+			}
 			case 'scope': {
 				if (elem.absolutePriority === undefined) {
 					throw new Error(`BUG!! computePriorityLevels should have set absolutePriority for all scopes`);
@@ -1825,6 +1891,7 @@ function hydrateEmptyTokenCount(elem: PromptElement, tokenizer: PriomptTokenizer
 		case 'capture':
 		case 'image':
 		case 'isolate':
+		case 'max':
 		case 'breaktoken':
 		case 'config':
 		case 'toolDefinition':
@@ -1892,6 +1959,11 @@ function hydrateIsolates(elem: PromptElement, tokenizer: PriomptTokenizer, shoul
 			}
 			return;
 		}
+		case 'max': {
+			// Max blocks don't need special hydration like isolates
+			// They participate in normal rendering flow
+			return hydrateIsolates(elem.children, tokenizer, shouldBuildSourceMap);
+		}
 		case 'chat': {
 			return hydrateIsolates(elem.children, tokenizer, shouldBuildSourceMap);
 		}
@@ -2047,6 +2119,20 @@ function renderWithLevel(
 				result.streamResponseObjectHandlers.push(...elem.cachedRenderOutput.streamResponseObjectHandlers);
 				return elem.cachedRenderOutput.sourceMap;
 			}
+			case 'max': {
+				// For max blocks, we render children normally but enforce the token limit locally
+				// This will be handled by the recursive call to renderWithLevelInPlace
+				const sourceMap = renderWithLevelInPlace(elem.children, sourceInfo !== undefined ? {
+					name: `max(${elem.tokenLimit})`,
+					isLast: sourceInfo.isLast,
+				} : undefined);
+				return (sourceMap === undefined || sourceInfo === undefined) ? undefined : {
+					name: sourceInfo.name,
+					children: [sourceMap],
+					start: 0,
+					end: sourceMap.end
+				};
+			}
 			case 'scope': {
 				if (elem.absolutePriority === undefined) {
 					throw new Error(`BUG!! computePriorityLevels should have set absolutePriority for all scopes`);
@@ -2358,6 +2444,7 @@ function validateNoUnhandledTypes(elem: PromptElement): void {
 			return;
 		}
 		case 'isolate':
+		case 'max':
 		case 'breaktoken':
 		case 'config':
 		case 'capture':
@@ -2391,6 +2478,7 @@ function validateNotBothAbsoluteAndRelativePriority(elem: PromptElement): void {
 	switch (elem.type) {
 		case 'chat':
 		case 'isolate':
+		case 'max':
 		case 'first': {
 			for (const child of elem.children) {
 				validateNotBothAbsoluteAndRelativePriority(child);
@@ -2450,6 +2538,13 @@ function validateNoChildrenHigherPriorityThanParent(elem: PromptElement, parentP
 			validateNoChildrenHigherPriorityThanParent(elem.children);
 			return;
 		}
+		case 'max': {
+			// max blocks participate in global priority calculation, so we do send the parent priority
+			for (const child of elem.children) {
+				validateNoChildrenHigherPriorityThanParent(child, parentPriority);
+			}
+			return;
+		}
 		case 'capture':
 		case 'image':
 		case 'breaktoken':
@@ -2521,6 +2616,14 @@ function computePriorityLevels(elem: AnyNode[] | AnyNode, parentPriority: number
 			// nothing happens because we fully re-render
 			return;
 		}
+		case 'max': {
+			// Unlike isolate, max blocks participate in global priority calculation
+			// We compute priorities for children and include them in global levels
+			for (const child of elem.children) {
+				computePriorityLevels(child, parentPriority, levels);
+			}
+			return;
+		}
 		case 'scope': {
 			// compute the priority of this scope
 			// the absolutePriority takes precedence over the relativePriority
@@ -2610,6 +2713,7 @@ function computePriorityLevelsTokensMapping(elem: NormalizedNode[] | NormalizedN
 			return;
 		}
 		case 'isolate':
+		case 'max':
 		case 'breaktoken':
 		case 'capture':
 		case 'config':
@@ -3126,6 +3230,7 @@ export function getPromptElementNodeCount(elem: PromptElement): number {
 			return 1;
 		case 'first':
 		case 'isolate':
+		case 'max':
 		case 'scope':
 		case 'chat':
 			return 1 + getPromptElementNodeCount(elem.children);