mirror of
https://github.com/mauriceboe/TREK.git
synced 2026-10-23 07:26:30 +00:00
feat(assistant): read the conversation from Anthropic's prompt cache
The Anthropic request marked only the instructions and the trip card for the prompt cache, so every step of a turn paid the whole history again. It now also marks the end of the history, which the next step starts from, and the end of what came before the user's question, which every step of the turn sends unchanged, within the four marks a request takes. Every provider now reports how much of the input it read from its cache (Anthropic cache_read_input_tokens, OpenAI cached_tokens), the turn adds it up, and the usage log line says it, so the saving can be measured.
This commit is contained in:
@@ -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 };
|
||||
|
||||
|
||||
@@ -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<void> {
|
||||
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,
|
||||
|
||||
@@ -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<string, AiFinishReason> = { 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<AiChatEvent> {
|
||||
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;
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -116,7 +116,7 @@ export class OpenAiResponsesClient implements AiChatClient {
|
||||
const events = await this.open(body, request.cacheKey, signal);
|
||||
|
||||
const items: Record<string, unknown>[] = [];
|
||||
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));
|
||||
|
||||
@@ -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<string>;
|
||||
/** 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<StepAnswer | Age
|
||||
} else if (event.type === 'tool-call') answer.received.push(event.call);
|
||||
else if (event.type === 'usage') {
|
||||
turn.usage.inputTokens += event.inputTokens;
|
||||
turn.usage.cachedInputTokens += event.cachedInputTokens;
|
||||
turn.usage.outputTokens += event.outputTokens;
|
||||
} else if (event.type === 'provider-state') answer.providerState = event.state;
|
||||
else if (event.type === 'finish') answer.cutOff = event.reason === 'length';
|
||||
@@ -423,7 +426,7 @@ export function runAgentLoop(
|
||||
w,
|
||||
signal,
|
||||
history: [...input.history],
|
||||
usage: { inputTokens: 0, outputTokens: 0 },
|
||||
usage: { inputTokens: 0, cachedInputTokens: 0, outputTokens: 0 },
|
||||
knownPlaceIds,
|
||||
nudged: false,
|
||||
};
|
||||
|
||||
@@ -868,6 +868,7 @@ export class AssistantService {
|
||||
}
|
||||
await this.usage.record(spec.userId, spec.connection.paidBy, {
|
||||
inputTokens: outcome.inputTokens,
|
||||
cachedInputTokens: outcome.cachedInputTokens,
|
||||
outputTokens: outcome.outputTokens,
|
||||
});
|
||||
});
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
import type { AiChatEvent, AiChatRequest } from '../../../../src/nest/ai-connections/ai-chat-client.interface';
|
||||
import {
|
||||
AnthropicMessagesClient,
|
||||
markHistoryForCache,
|
||||
toAnthropicMessages,
|
||||
} from '../../../../src/nest/ai-connections/providers/anthropic-messages.client';
|
||||
import { OpenAiChatClient, toChatMessages } from '../../../../src/nest/ai-connections/providers/openai-chat.client';
|
||||
@@ -194,6 +195,44 @@ describe('AnthropicMessagesClient', () => {
|
||||
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' },
|
||||
]);
|
||||
|
||||
@@ -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({
|
||||
|
||||
Reference in New Issue
Block a user