chore: add legacy prioritization for migration

package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@vscode/prompt-tsx",
-	"version": "0.3.0-alpha",
+	"version": "0.3.0-alpha.1",
 	"description": "Declare LLM prompts with TSX",
 	"main": "./dist/base/index.js",
 	"types": "./dist/base/index.d.ts",

src/base/materialized.ts

@@ -29,6 +29,7 @@ export class MaterializedContainer implements IMaterializedNode {
 		public readonly priority: number,
 		public readonly children: MaterializedNode[],
 		public readonly metadata: PromptMetadata[],
+		public readonly isLegacyPrioritization = false,
 	) { }
 
 	/** @inheritdoc */
@@ -75,7 +76,11 @@ export class MaterializedContainer implements IMaterializedNode {
 
 	/** Removes the node in the tree with the lowest priority. */
 	removeLowestPriorityChild(): void {
-		removeLowestPriorityChild(this.children);
+		if (this.isLegacyPrioritization) {
+			removeLowestPriorityLegacy(this);
+		} else {
+			removeLowestPriorityChild(this.children);
+		}
 	}
 }
 
@@ -92,8 +97,8 @@ export class MaterializedChatMessageTextChunk {
 		return this._upperBound(tokenizer);
 	}
 
-	private readonly _upperBound = once((tokenizer: ITokenizer) => {
-		return tokenizer.tokenLength(this.text);
+	private readonly _upperBound = once(async (tokenizer: ITokenizer) => {
+		return await tokenizer.tokenLength(this.text) + (this.lineBreakBefore ? 1 : 0);
 	});
 }
 
@@ -131,7 +136,7 @@ export class MaterializedChatMessage implements IMaterializedNode {
 		this.onChunksChange();
 	}
 
-	private onChunksChange() {
+	onChunksChange() {
 		this._tokenCount.clear();
 		this._upperBound.clear();
 		this._text.clear();
@@ -145,7 +150,7 @@ export class MaterializedChatMessage implements IMaterializedNode {
 		let total = await this._baseMessageTokenCount(tokenizer)
 		await Promise.all(this.children.map(async (chunk) => {
 			const amt = await chunk.upperBoundTokenCount(tokenizer);
-			total += amt + (chunk instanceof MaterializedChatMessageTextChunk && chunk.lineBreakBefore ? 1 : 0);
+			total += amt;
 		}));
 		return total;
 	});
@@ -231,6 +236,55 @@ function* textChunks(node: MaterializedNode): Generator<MaterializedChatMessageT
 	}
 }
 
+function removeLowestPriorityLegacy(root: MaterializedNode) {
+	let lowest: undefined | {
+		chain: (MaterializedContainer | MaterializedChatMessage)[],
+		node: MaterializedChatMessageTextChunk;
+	};
+
+	function findLowestInTree(node: MaterializedNode, chain: (MaterializedContainer | MaterializedChatMessage)[]) {
+		if (node instanceof MaterializedChatMessageTextChunk) {
+			if (!lowest || node.priority < lowest.node.priority) {
+				lowest = { chain: chain.slice(), node };
+			}
+		} else {
+			chain.push(node);
+			for (const child of node.children) {
+				findLowestInTree(child, chain);
+			}
+			chain.pop();
+		}
+	}
+
+	findLowestInTree(root, []);
+
+	if (!lowest) {
+		throw new Error('No lowest priority node found');
+	}
+
+	let needle: MaterializedNode = lowest.node;
+	let i = lowest.chain.length - 1;
+	for (; i >= 0; i--) {
+		const node = lowest.chain[i];
+		node.children.splice(node.children.indexOf(needle), 1);
+		if (node instanceof MaterializedChatMessage) {
+			node.onChunksChange();
+		}
+		if (node.children.length > 0) {
+			break;
+		}
+
+		needle = node;
+	}
+
+	for (; i >= 0; i--) {
+		const node = lowest.chain[i];
+		if (node instanceof MaterializedChatMessage) {
+			node.onChunksChange();
+		}
+	}
+}
+
 function removeLowestPriorityChild(children: MaterializedNode[]) {
 	if (!children.length) {
 		return;

src/base/promptElements.tsx

@@ -270,3 +270,16 @@ export class ToolResult extends PromptElement<IToolResultProps> {
 		}
 	}
 }
+
+/**
+ * Marker element that uses the legacy global prioritization algorithm (0.2.x
+ * if this library) for pruning child elements. This will be removed in
+ * the future.
+ *
+ * @deprecated
+ */
+export class LegacyPrioritization extends PromptElement {
+	render() {
+		return <>{this.props.children}</>;
+	}
+}

src/base/promptRenderer.ts

@@ -8,7 +8,7 @@ import { PromptNodeType } from './jsonTypes';
 import { MaterializedChatMessage, MaterializedChatMessageTextChunk, MaterializedContainer } from './materialized';
 import { ChatMessage } from "./openai";
 import { PromptElement } from "./promptElement";
-import { AssistantMessage, BaseChatMessage, ChatMessagePromptElement, TextChunk, ToolMessage, isChatMessagePromptElement } from "./promptElements";
+import { AssistantMessage, BaseChatMessage, ChatMessagePromptElement, LegacyPrioritization, TextChunk, ToolMessage, isChatMessagePromptElement } from "./promptElements";
 import { PromptMetadata, PromptReference } from "./results";
 import { ITokenizer } from "./tokenizer/tokenizer";
 import { ITracer } from './tracer';
@@ -650,6 +650,7 @@ class PromptTreeElement {
 				this._obj?.props.priority || 0,
 				this._children.map(child => child.materialize()),
 				this._metadata,
+				this._obj instanceof LegacyPrioritization,
 			);
 		}
 	}

src/base/test/renderer.test.tsx

@@ -8,6 +8,7 @@ import { BaseTokensPerCompletion, ChatMessage, ChatRole } from '../openai';
 import { PromptElement } from '../promptElement';
 import {
 	AssistantMessage,
+	LegacyPrioritization,
 	PrioritizedList,
 	SystemMessage,
 	TextChunk,
@@ -407,6 +408,34 @@ suite('PromptRenderer', () => {
 				</SystemMessage>
 			</>, ['a', 'b', 'c', 'd']);
 		});
+
+		test('uses legacy prioritization', async () => {
+			class Wrap1 extends PromptElement {
+				render() {
+					return <>
+						<TextChunk priority={1}>a</TextChunk>
+						<TextChunk priority={10}>b</TextChunk>
+					</>
+				}
+			}
+			class Wrap2 extends PromptElement {
+				render() {
+					return <>
+						<TextChunk priority={2}>c</TextChunk>
+						<TextChunk priority={15}>d</TextChunk>
+					</>
+				}
+			}
+			await assertPruningOrder(<LegacyPrioritization>
+				<UserMessage>
+					<Wrap1 priority={1} />
+					<Wrap2 priority={2} />
+				</UserMessage>
+				<UserMessage>
+					<TextChunk priority={5}>e</TextChunk>
+				</UserMessage>
+			</LegacyPrioritization>, ['a', 'c', 'e', 'b', 'd']);
+		});
 	});
 
 	suite('truncates tokens exceeding token budget', async () => {