diff --git a/package-lock.json b/package-lock.json index df540f5..d543768 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "@vscode/prompt-tsx", - "version": "0.2.5-alpha", + "version": "0.2.6-alpha", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@vscode/prompt-tsx", - "version": "0.2.5-alpha", + "version": "0.2.6-alpha", "license": "SEE LICENSE IN LICENSE", "devDependencies": { "@microsoft/tiktokenizer": "^1.0.6", diff --git a/package.json b/package.json index c5703bb..674759b 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@vscode/prompt-tsx", - "version": "0.2.5-alpha", + "version": "0.2.6-alpha", "description": "Declare LLM prompts with TSX", "main": "./dist/base/index.js", "types": "./dist/base/index.d.ts", diff --git a/src/base/promptRenderer.ts b/src/base/promptRenderer.ts index deb01ba..d29f182 100644 --- a/src/base/promptRenderer.ts +++ b/src/base/promptRenderer.ts @@ -21,6 +21,11 @@ export interface RenderPromptResult { * The references that survived prioritization in the rendered {@link RenderPromptResult.messages messages}. */ readonly references: PromptReference[]; + + /** + * The references attached to chat message chunks that did not survive prioritization. + */ + readonly omittedReferences: PromptReference[]; } export type QueueItem = { @@ -212,6 +217,7 @@ export class PromptRenderer

{ prioritizedChunks.sort((a, b) => cmp(things[a.index], things[b.index])); let remainingBudget = this._endpoint.modelMaxPromptTokens; + const omittedChunks: T[] = []; while (prioritizedChunks.length > 0) { const prioritizedChunk = prioritizedChunks.shift()!; const index = prioritizedChunk.index; @@ -224,6 +230,7 @@ export class PromptRenderer

{ } if (tokenCount > remainingBudget) { // Wouldn't fit anymore + omittedChunks.push(chunk); break; } chunkResult[index] = chunk; @@ -233,7 +240,13 @@ export class PromptRenderer

{ remainingBudget -= tokenCount; } - return { result: coalesce(chunkResult), tokenCount: this._endpoint.modelMaxPromptTokens - remainingBudget }; + for (const omittedChunk of prioritizedChunks) { + const index = omittedChunk.index; + const chunk = things[index]; + omittedChunks.push(chunk); + } + + return { result: coalesce(chunkResult), tokenCount: this._endpoint.modelMaxPromptTokens - remainingBudget, omittedChunks }; } /** @@ -256,7 +269,7 @@ export class PromptRenderer

{ // First pass: sort message chunks by priority. Note that this can yield an imprecise result due to token boundaries within a single chat message // so we also need to do a second pass over the full chat messages later const chunkMessages = new Set(); - const { result: prioritizedChunks } = await this._prioritize( + const { result: prioritizedChunks, omittedChunks } = await this._prioritize( resultChunks, (a, b) => MaterializedChatMessageTextChunk.cmp(a, b), async (chunk) => { @@ -298,7 +311,7 @@ export class PromptRenderer

{ const messageResult = prioritizedMaterializedChatMessages.map(message => message?.toChatMessage()); // Remove undefined and duplicate references - const { references } = prioritizedMaterializedChatMessages.reduce<{ references: PromptReference[], names: Set }>((acc, message) => { + const { references, names } = prioritizedMaterializedChatMessages.reduce<{ references: PromptReference[], names: Set }>((acc, message) => { [...this._references, ...message.references].forEach((ref) => { const isVariableName = 'variableName' in ref.anchor; if (isVariableName && !acc.names.has(ref.anchor.variableName)) { @@ -311,7 +324,21 @@ export class PromptRenderer

{ return acc; }, { references: [], names: new Set() }); - return { messages: messageResult, hasIgnoredFiles: this._ignoredFiles.length > 0, tokenCount, references: coalesce(references) }; + // Collect the references for chat message chunks that did not survive prioritization + const { references: omittedReferences } = omittedChunks.reduce<{ references: PromptReference[] }>((acc, message) => { + message.references.forEach((ref) => { + const isVariableName = 'variableName' in ref.anchor; + if (isVariableName && !names.has(ref.anchor.variableName)) { + acc.references.push(ref); + names.add(ref.anchor.variableName); + } else if (!isVariableName) { + acc.references.push(ref); + } + }); + return acc; + }, { references: [] }); + + return { messages: messageResult, hasIgnoredFiles: this._ignoredFiles.length > 0, tokenCount, references: coalesce(references), omittedReferences: coalesce(omittedReferences) }; } private _handlePromptChildren(element: QueueItem, P>, pieces: ProcessedPromptPiece[], sizing: PromptSizingContext, progress: Progress | undefined, token: CancellationToken | undefined) { diff --git a/src/base/results.ts b/src/base/results.ts index 0dce560..67440fe 100644 --- a/src/base/results.ts +++ b/src/base/results.ts @@ -14,12 +14,19 @@ export abstract class PromptMetadata { } } +export enum ChatResponseReferencePartStatusKind { + Complete = 1, + Partial = 2, + Omitted = 3 +} + /** * A reference used for creating the prompt. */ export class PromptReference { constructor( readonly anchor: Uri | Location | { variableName: string; value?: Uri | Location }, - readonly iconPath?: Uri | ThemeIcon | { light: Uri; dark: Uri } + readonly iconPath?: Uri | ThemeIcon | { light: Uri; dark: Uri }, + readonly options?: { status?: { description: string; kind: ChatResponseReferencePartStatusKind } } ) { } } diff --git a/src/base/test/renderer.test.tsx b/src/base/test/renderer.test.tsx index fc18a3b..816fda9 100644 --- a/src/base/test/renderer.test.tsx +++ b/src/base/test/renderer.test.tsx @@ -893,7 +893,7 @@ LOW MED 00 01 02 03 04 05 06 07 08 09 const res = await inst.render(undefined, undefined); assert.equal(res.messages.length, 1); assert.equal(res.references.length, 0); - + assert.equal(res.omittedReferences.length, 1); }); test('reports references under nested extrinsics', async () => {