diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index efc4386d8..7c1ed3af6 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -11,6 +11,10 @@ - Added `durationMs` to `AssistantMessage`: `AssistantMessageEventStream` measures each response with a monotonic clock from the start of the request to its final message, for every API implementation, including direct calls. Deferred results fetched later stay untimed ([#10549](https://github.com/earendil-works/pi/issues/10549)) - Added an optional `durationMs` to `ToolResultMessage` for the execution time of the tool ([#10549](https://github.com/earendil-works/pi/issues/10549)) +### Changed + +- The faux provider's prompt-cache usage estimate compares the previous and current prompt message by message and compares characters only from the first differing message; the usage numbers are unchanged. + ### Fixed - Fixed `server_busy` and `servers are currently busy` provider errors ending the turn instead of being retried ([#10543](https://github.com/earendil-works/pi/issues/10543)) diff --git a/packages/ai/src/providers/faux.ts b/packages/ai/src/providers/faux.ts index a38c18774..c3771cc2c 100644 --- a/packages/ai/src/providers/faux.ts +++ b/packages/ai/src/providers/faux.ts @@ -214,8 +214,23 @@ function messageToText(message: Message): string { return toolResultToText(message); } -function serializeContext(context: TranscriptContext): string { - return context.messages.map((message) => `${message.role}:${messageToText(message)}`).join("\n\n"); +/** Length of the prompt text that joins `messages` with blank lines. */ +function joinedLength(messages: readonly string[], count = messages.length): number { + let length = count > 0 ? (count - 1) * 2 : 0; + for (let index = 0; index < count; index++) length += messages[index]!.length; + return length; +} + +/** + * Length of the common prefix of the two joined prompts. Equal messages are compared whole; characters are compared + * only from the first message that differs. + */ +function commonPromptPrefixLength(previous: readonly string[], current: readonly string[]): number { + let index = 0; + while (index < previous.length && index < current.length && previous[index] === current[index]) index++; + const rest = (messages: readonly string[]) => + index === messages.length ? "" : (index > 0 ? "\n\n" : "") + messages.slice(index).join("\n\n"); + return joinedLength(previous, index) + commonPrefixLength(rest(previous), rest(current)); } function commonPrefixLength(a: string, b: string): number { @@ -231,10 +246,12 @@ function withUsageEstimate( message: AssistantMessage, context: TranscriptContext, options: StreamOptions | undefined, - promptCache: Map, + promptCache: Map, ): AssistantMessage { - const promptText = serializeContext(context); - const promptTokens = estimateTokens(promptText); + // One text per message; the whole prompt joins them with blank lines. + const prompt = context.messages.map((message) => `${message.role}:${messageToText(message)}`); + const promptLength = joinedLength(prompt); + const promptTokens = Math.ceil(promptLength / 4); const outputTokens = estimateTokens(assistantContentToText(message.content)); let input = promptTokens; let cacheRead = 0; @@ -244,14 +261,14 @@ function withUsageEstimate( if (sessionId && options?.cacheRetention !== "none") { const previousPrompt = promptCache.get(sessionId); if (previousPrompt) { - const cachedChars = commonPrefixLength(previousPrompt, promptText); - cacheRead = estimateTokens(previousPrompt.slice(0, cachedChars)); - cacheWrite = estimateTokens(promptText.slice(cachedChars)); + const cachedChars = commonPromptPrefixLength(previousPrompt, prompt); + cacheRead = Math.ceil(cachedChars / 4); + cacheWrite = Math.ceil((promptLength - cachedChars) / 4); input = Math.max(0, promptTokens - cacheRead); } else { cacheWrite = promptTokens; } - promptCache.set(sessionId, promptText); + promptCache.set(sessionId, prompt); } return { @@ -445,7 +462,7 @@ export function createFauxCore(options: RegisterFauxProviderOptions) { let pendingResponses: FauxResponseStep[] = []; const tokensPerSecond = options.tokensPerSecond; const state: FauxProviderState = { callCount: 0, deferredFetchCount: 0, cancelledDeferred: [] }; - const promptCache = new Map(); + const promptCache = new Map(); const deferredResponses = new Map< string, { diff --git a/packages/ai/test/faux-provider.test.ts b/packages/ai/test/faux-provider.test.ts index 557885f23..6f4178111 100644 --- a/packages/ai/test/faux-provider.test.ts +++ b/packages/ai/test/faux-provider.test.ts @@ -324,6 +324,27 @@ describe("faux provider", () => { expect(second.usage.input + second.usage.cacheRead).toBeGreaterThan(second.usage.input); }); + it("counts cached characters up to the first difference in the joined prompt", async () => { + const registration = registerFauxProvider(); + registrations.push(registration); + registration.setResponses([fauxAssistantMessage("a"), fauxAssistantMessage("b"), fauxAssistantMessage("c")]); + const options = { sessionId: "session-1", cacheRetention: "short" } as const; + const user = (content: string) => ({ role: "user" as const, content, timestamp: 1 }); + + // Prompt texts: "user:hello world" (16 characters), then 16 + 2 + "user:next" (9) = 27. + await complete(registration.getModel(), { messages: [user("hello world")] }, options); + const extended = await complete( + registration.getModel(), + { messages: [user("hello world"), user("next")] }, + options, + ); + expect(extended.usage).toMatchObject({ input: 3, cacheRead: 4, cacheWrite: 3 }); + + // The first message now differs after "user:hello w" (12 characters): "user:hello wide" + 2 + 9 = 26. + const edited = await complete(registration.getModel(), { messages: [user("hello wide"), user("next")] }, options); + expect(edited.usage).toMatchObject({ input: 4, cacheRead: 3, cacheWrite: 4 }); + }); + it("does not simulate caching when cacheRetention is none", async () => { const registration = registerFauxProvider(); registrations.push(registration);