diff --git a/server/src/nest/ai-connections/ai-chat-client.interface.ts b/server/src/nest/ai-connections/ai-chat-client.interface.ts index 94caba6a1..a9c80b91d 100644 --- a/server/src/nest/ai-connections/ai-chat-client.interface.ts +++ b/server/src/nest/ai-connections/ai-chat-client.interface.ts @@ -51,7 +51,8 @@ export type AiFinishReason = 'stop' | 'tool-calls' | 'length'; export type AiChatEvent = | { type: 'text-delta'; delta: string } | { type: 'tool-call'; call: AiToolCall } - | { type: 'usage'; inputTokens: number; outputTokens: number } + /** `cachedInputTokens` is the part of `inputTokens` the provider read from its prompt cache. */ + | { type: 'usage'; inputTokens: number; cachedInputTokens: number; outputTokens: number } | { type: 'provider-state'; state: AiProviderState } | { type: 'finish'; reason: AiFinishReason }; diff --git a/server/src/nest/ai-connections/ai-usage.service.ts b/server/src/nest/ai-connections/ai-usage.service.ts index ca75ad79d..88277ac26 100644 --- a/server/src/nest/ai-connections/ai-usage.service.ts +++ b/server/src/nest/ai-connections/ai-usage.service.ts @@ -1,5 +1,6 @@ import { AssistantUsage } from '../../db/entities/AssistantUsage.entity'; import type { AssistantUsageRepository } from '../../db/repositories/AssistantUsage.repository'; +import { logDebug } from '../audit/audit-log.logger'; import { InjectRepository } from '@mikro-orm/nestjs'; import { Injectable } from '@nestjs/common'; import { todayUtc, type AiConnectionScope } from '@trek/shared'; @@ -13,12 +14,22 @@ import { todayUtc, type AiConnectionScope } from '@trek/shared'; export class AiUsageService { constructor(@InjectRepository(AssistantUsage) private readonly usage: AssistantUsageRepository) {} - /** Adds a turn's usage to today's row; the turn itself is counted once, when it starts. */ + /** + * Adds a turn's usage to today's row; the turn itself is counted once, when + * it starts. How much of the input came from the provider's prompt cache is + * only logged, to see what the cache saves. + */ async record( userId: number, paidBy: AiConnectionScope, - delta: { turns?: number; inputTokens?: number; outputTokens?: number }, + delta: { turns?: number; inputTokens?: number; cachedInputTokens?: number; outputTokens?: number }, ): Promise { + if (delta.inputTokens) { + logDebug( + `[assistant] turn of user ${userId}: ${delta.inputTokens} input tokens, ` + + `${delta.cachedInputTokens ?? 0} of them from the prompt cache, ${delta.outputTokens ?? 0} output tokens`, + ); + } await this.usage.add(userId, todayUtc(), paidBy, { turns: delta.turns ?? 0, inputTokens: delta.inputTokens ?? 0, diff --git a/server/src/nest/ai-connections/providers/anthropic-messages.client.ts b/server/src/nest/ai-connections/providers/anthropic-messages.client.ts index b8dfb6dd2..d81f7ca11 100644 --- a/server/src/nest/ai-connections/providers/anthropic-messages.client.ts +++ b/server/src/nest/ai-connections/providers/anthropic-messages.client.ts @@ -67,12 +67,38 @@ export function toAnthropicMessages(messages: AiChatMessage[]): AnthropicMessage return out; } +/** The most cache breakpoints Anthropic takes in one request; the system blocks use the first ones. */ +const MAX_CACHE_BREAKPOINTS = 4; +const CACHED = { type: 'ephemeral' } as const; + +/** + * Marks the history for the prompt cache, with what the system blocks left of + * the breakpoints: the end of the whole history, which the next step of the + * turn starts from, and the end of what came before the user's question of + * this turn, which every step of the turn sends unchanged. The second one + * still hits when a step added more blocks (many tool calls and their + * results) than the cache looks back over from the first. The turns are the + * request's own copies, so they are marked in place. + */ +export function markHistoryForCache(messages: AnthropicMessage[], breakpoints: number): void { + const blocks = messages.flatMap((message) => message.content.map((block) => ({ role: message.role, block }))); + const last = blocks.at(-1); + if (!last || breakpoints < 1) return; + last.block.cache_control = CACHED; + let question = blocks.length - 1; + while (question >= 0 && !(blocks[question]?.role === 'user' && blocks[question]?.block.type === 'text')) question--; + const before = question > 0 ? blocks[question - 1] : undefined; + if (breakpoints >= 2 && before && before !== last) before.block.cache_control = CACHED; +} + const STOP_REASONS: Record = { tool_use: 'tool-calls', max_tokens: 'length' }; /** * Anthropic's Messages API, streamed. The instructions and the trip card are * two system blocks, each marked for the prompt cache, so the tools and the - * instructions are read from cache on every later step of a conversation. + * instructions are read from cache on every later step of a conversation; the + * history is marked too (markHistoryForCache), so a step reads the turns + * before it from the cache instead of paying for them again. */ export class AnthropicMessagesClient implements AiChatClient { readonly kind = 'anthropic' as const; @@ -86,13 +112,15 @@ export class AnthropicMessagesClient implements AiChatClient { async *stream(request: AiChatRequest, signal: AbortSignal): AsyncGenerator { const system = [request.system, request.context] .filter((text): text is string => !!text) - .map((text) => ({ type: 'text', text, cache_control: { type: 'ephemeral' } })); + .map((text) => ({ type: 'text', text, cache_control: CACHED })); + const messages = toAnthropicMessages(request.messages); + markHistoryForCache(messages, MAX_CACHE_BREAKPOINTS - system.length); const body = { model: this.model, max_tokens: MAX_TOKENS, stream: true, system, - messages: toAnthropicMessages(request.messages), + messages, ...(request.tools.length > 0 ? { tools: request.tools.map((tool) => ({ @@ -129,7 +157,12 @@ export class AnthropicMessagesClient implements AiChatClient { if (out) yield out; } if (!turn.finished) throw new AiProviderError('provider_unavailable', 0, true, undefined, 'stream_cut'); - yield { type: 'usage', inputTokens: turn.inputTokens, outputTokens: turn.outputTokens }; + yield { + type: 'usage', + inputTokens: turn.inputTokens, + cachedInputTokens: turn.cachedInputTokens, + outputTokens: turn.outputTokens, + }; yield { type: 'finish', reason: turn.reason }; } } @@ -137,6 +170,7 @@ export class AnthropicMessagesClient implements AiChatClient { /** What one streamed Messages API turn has said so far, fed one event at a time. */ class AnthropicTurn { inputTokens = 0; + cachedInputTokens = 0; outputTokens = 0; reason: AiFinishReason = 'stop'; finished = false; @@ -152,6 +186,7 @@ class AnthropicTurn { asCount(usage.input_tokens) + asCount(usage.cache_creation_input_tokens) + asCount(usage.cache_read_input_tokens); + this.cachedInputTokens = asCount(usage.cache_read_input_tokens); this.outputTokens = asCount(usage.output_tokens); return null; } diff --git a/server/src/nest/ai-connections/providers/openai-chat.client.ts b/server/src/nest/ai-connections/providers/openai-chat.client.ts index 227e9bd9f..a0b7526e6 100644 --- a/server/src/nest/ai-connections/providers/openai-chat.client.ts +++ b/server/src/nest/ai-connections/providers/openai-chat.client.ts @@ -115,7 +115,12 @@ export class OpenAiChatClient implements AiChatClient { } if (!turn.done && !turn.reason) throw new AiProviderError('provider_unavailable', 0, true, undefined, 'stream_cut'); for (const call of turn.calls()) yield { type: 'tool-call', call }; - yield { type: 'usage', inputTokens: turn.inputTokens, outputTokens: turn.outputTokens }; + yield { + type: 'usage', + inputTokens: turn.inputTokens, + cachedInputTokens: turn.cachedInputTokens, + outputTokens: turn.outputTokens, + }; yield { type: 'finish', reason: turn.finishReason() }; } } @@ -127,6 +132,7 @@ export class OpenAiChatClient implements AiChatClient { */ class ChatTurn { inputTokens = 0; + cachedInputTokens = 0; outputTokens = 0; reason: string | null = null; done = false; @@ -137,6 +143,7 @@ class ChatTurn { const usage = asRecord(chunk.usage); if (chunk.usage) { this.inputTokens = asCount(usage.prompt_tokens); + this.cachedInputTokens = asCount(asRecord(usage.prompt_tokens_details).cached_tokens); this.outputTokens = asCount(usage.completion_tokens); } const choice = asRecord(Array.isArray(chunk.choices) ? chunk.choices[0] : undefined); diff --git a/server/src/nest/ai-connections/providers/openai-responses.client.ts b/server/src/nest/ai-connections/providers/openai-responses.client.ts index bbbca9672..3790b9db6 100644 --- a/server/src/nest/ai-connections/providers/openai-responses.client.ts +++ b/server/src/nest/ai-connections/providers/openai-responses.client.ts @@ -116,7 +116,7 @@ export class OpenAiResponsesClient implements AiChatClient { const events = await this.open(body, request.cacheKey, signal); const items: Record[] = []; - let usage = { inputTokens: 0, outputTokens: 0 }; + let usage = { inputTokens: 0, cachedInputTokens: 0, outputTokens: 0 }; let completed = false; for await (const sse of readSseEvents(events)) { const event = parseEventData(sse.data); @@ -134,7 +134,11 @@ export class OpenAiResponsesClient implements AiChatClient { } } else if (event.type === 'response.completed') { const reported = asRecord(asRecord(event.response).usage); - usage = { inputTokens: asCount(reported.input_tokens), outputTokens: asCount(reported.output_tokens) }; + usage = { + inputTokens: asCount(reported.input_tokens), + cachedInputTokens: asCount(asRecord(reported.input_tokens_details).cached_tokens), + outputTokens: asCount(reported.output_tokens), + }; completed = true; } else if (event.type === 'error' || event.type === 'response.failed' || event.type === 'response.incomplete') { throw providerFailure(this.kind, 0, errorCodeOf(event) ?? errorCodeOf(event.response) ?? asText(event.type)); diff --git a/server/src/nest/assistant/assistant-agent.ts b/server/src/nest/assistant/assistant-agent.ts index 86d63316c..7eb805fa5 100644 --- a/server/src/nest/assistant/assistant-agent.ts +++ b/server/src/nest/assistant/assistant-agent.ts @@ -82,6 +82,8 @@ interface AgentOutcome { /** What failed (a provider refusal, the step limit), null for a turn that ended well or was stopped. */ code: AssistantErrorCode | null; inputTokens: number; + /** The part of inputTokens the provider read from its prompt cache. */ + cachedInputTokens: number; outputTokens: number; } @@ -179,7 +181,7 @@ interface Turn { w: UiStreamWriter; signal: AbortSignal; history: AiChatMessage[]; - usage: { inputTokens: number; outputTokens: number }; + usage: { inputTokens: number; cachedInputTokens: number; outputTokens: number }; knownPlaceIds: Set; /** Whether the turn used its one push (an empty, cut-off or announcing answer). */ nudged: boolean; @@ -224,6 +226,7 @@ async function streamAnswer(turn: Turn, rowId: number): Promise { expect([0, 1, 2, 3].map((i) => dig(body, 'messages', i, 'role'))).toEqual(['user', 'assistant', 'user', undefined]); }); + it('ANTH-006: marks the end of the history and the end of what came before the question for the cache', async () => { + const http = fetchAnswering(streamResponse(STREAM)); + await run(new AnthropicMessagesClient('m', 'k', http.fetcher).stream(REQUEST, signal())); + const marked = JSON.stringify(http.body()).match(/cache_control/g) ?? []; + expect(marked).toHaveLength(4); + // The question "Thanks" shares its turn with the tool result before it: both carry a mark. + expect(dig(http.body(), 'messages', 2, 'content')).toEqual([ + { type: 'tool_result', tool_use_id: 'call_1', content: '{"id":7}', cache_control: { type: 'ephemeral' } }, + { type: 'text', text: 'Thanks', cache_control: { type: 'ephemeral' } }, + ]); + expect(dig(http.body(), 'messages', 0, 'content', 0, 'cache_control')).toBeUndefined(); + }); + + it('ANTH-007: later steps of a turn keep the mark before the question; never more marks than allowed', () => { + const step = toAnthropicMessages([ + ...REQUEST.messages, + { role: 'assistant', text: '', toolCalls: [{ id: 'call_2', name: 'list_places', args: {} }] }, + { role: 'tool', toolCallId: 'call_2', name: 'list_places', output: '[]' }, + ]); + markHistoryForCache(step, 2); + const marks = step.flatMap((turn, i) => + turn.content.flatMap((block, j) => (block.cache_control ? [`${i}.${j}`] : [])), + ); + expect(marks).toEqual(['2.0', '4.0']); + + const first = toAnthropicMessages([{ role: 'user', text: 'Hi' }]); + markHistoryForCache(first, 2); + expect(first[0]?.content).toEqual([{ type: 'text', text: 'Hi', cache_control: { type: 'ephemeral' } }]); + + const none = toAnthropicMessages(REQUEST.messages); + markHistoryForCache(none, 0); + expect(JSON.stringify(none)).not.toContain('cache_control'); + const one = toAnthropicMessages(REQUEST.messages); + markHistoryForCache(one, 1); + expect(JSON.stringify(one).match(/cache_control/g)).toHaveLength(1); + markHistoryForCache([], 2); + }); + it('ANTH-002: a tool result and the next user text share one user turn', () => { expect(toAnthropicMessages(REQUEST.messages)[2]).toEqual({ role: 'user', @@ -216,7 +255,7 @@ describe('AnthropicMessagesClient', () => { expect(events).toEqual([ { type: 'text-delta', delta: 'On it' }, { type: 'tool-call', call: { id: 'toolu_01', name: 'create_place', args: { name: 'Ginkaku-ji' } } }, - { type: 'usage', inputTokens: 100, outputTokens: 42 }, + { type: 'usage', inputTokens: 100, cachedInputTokens: 90, outputTokens: 42 }, { type: 'finish', reason: 'tool-calls' }, ]); }); @@ -273,7 +312,7 @@ describe('OpenAiChatClient', () => { { choices: [{ delta: { tool_calls: [{ index: 0, function: { arguments: '{"name":' } }] } }] }, { choices: [{ delta: { tool_calls: [{ index: 0, function: { arguments: '"Nijo"}' } }] } }] }, { choices: [{ delta: {}, finish_reason: 'tool_calls' }] }, - { choices: [], usage: { prompt_tokens: 120, completion_tokens: 30 } }, + { choices: [], usage: { prompt_tokens: 120, prompt_tokens_details: { cached_tokens: 64 }, completion_tokens: 30 } }, '[DONE]', ]; @@ -314,7 +353,7 @@ describe('OpenAiChatClient', () => { { type: 'text-delta', delta: 'Sure' }, { type: 'text-delta', delta: ', adding.' }, { type: 'tool-call', call: { id: 'call_a', name: 'create_place', args: { name: 'Nijo' } } }, - { type: 'usage', inputTokens: 120, outputTokens: 30 }, + { type: 'usage', inputTokens: 120, cachedInputTokens: 64, outputTokens: 30 }, { type: 'finish', reason: 'tool-calls' }, ]); }); @@ -392,7 +431,10 @@ describe('OpenAiResponsesClient', () => { { type: 'response.output_text.delta', delta: 'Adding ' }, { type: 'response.output_text.delta', delta: 'it.' }, { type: 'response.output_item.done', item: CALL }, - { type: 'response.completed', response: { usage: { input_tokens: 500, output_tokens: 60 } } }, + { + type: 'response.completed', + response: { usage: { input_tokens: 500, input_tokens_details: { cached_tokens: 400 }, output_tokens: 60 } }, + }, ]; const tokens = (...values: string[]) => vi.fn(async (_rejected?: string) => values.shift() ?? 'none'); @@ -423,7 +465,7 @@ describe('OpenAiResponsesClient', () => { { type: 'text-delta', delta: 'Adding ' }, { type: 'text-delta', delta: 'it.' }, { type: 'tool-call', call: { id: 'call_x', name: 'create_place', args: { name: 'Ryoan-ji' } } }, - { type: 'usage', inputTokens: 500, outputTokens: 60 }, + { type: 'usage', inputTokens: 500, cachedInputTokens: 400, outputTokens: 60 }, { type: 'provider-state', state: { adapter: 'openai-responses', items: [REASONING, CALL] } }, { type: 'finish', reason: 'tool-calls' }, ]); diff --git a/server/tests/unit/nest/assistant/assistant-agent.test.ts b/server/tests/unit/nest/assistant/assistant-agent.test.ts index c44c170d5..0df6980a6 100644 --- a/server/tests/unit/nest/assistant/assistant-agent.test.ts +++ b/server/tests/unit/nest/assistant/assistant-agent.test.ts @@ -158,12 +158,12 @@ describe('runAgentLoop', () => { [ { type: 'text-delta', delta: 'Hello ' }, { type: 'text-delta', delta: 'there' }, - { type: 'usage', inputTokens: 10, outputTokens: 3 }, + { type: 'usage', inputTokens: 10, cachedInputTokens: 0, outputTokens: 3 }, { type: 'finish', reason: 'stop' }, ], ]); const outcome = await runAgentLoop(t.deps, t.input, t.writer, new AbortController().signal); - expect(outcome).toEqual({ end: 'stop', code: null, inputTokens: 10, outputTokens: 3 }); + expect(outcome).toEqual({ end: 'stop', code: null, inputTokens: 10, cachedInputTokens: 0, outputTokens: 3 }); expect(t.memory.saved).toEqual([ { rowId: 100, step: { parts: [{ type: 'text', text: 'Hello there' }], providerState: null }, status: 'complete' }, ]); @@ -184,16 +184,16 @@ describe('runAgentLoop', () => { [ { type: 'tool-call', call: { id: 'toolu_1', name: 'list_places', args: { tripId: 1 } } }, { type: 'provider-state', state: { adapter: 'openai-responses', items: [{ type: 'function_call' }] } }, - { type: 'usage', inputTokens: 5, outputTokens: 1 }, + { type: 'usage', inputTokens: 5, cachedInputTokens: 0, outputTokens: 1 }, { type: 'finish', reason: 'tool-calls' }, ], [ { type: 'text-delta', delta: 'Found it' }, - { type: 'usage', inputTokens: 7, outputTokens: 2 }, + { type: 'usage', inputTokens: 7, cachedInputTokens: 5, outputTokens: 2 }, ], ]); const outcome = await runAgentLoop(t.deps, t.input, t.writer, new AbortController().signal); - expect(outcome).toEqual({ end: 'stop', code: null, inputTokens: 12, outputTokens: 3 }); + expect(outcome).toEqual({ end: 'stop', code: null, inputTokens: 12, cachedInputTokens: 5, outputTokens: 3 }); expect(t.tools.call).toHaveBeenCalledWith('list_places', {}); const second = t.requests[1].messages; expect(second[1]).toEqual({