diff --git a/.changeset/anthropic-prompt-caching.md b/.changeset/anthropic-prompt-caching.md new file mode 100644 index 000000000..d47a4b962 --- /dev/null +++ b/.changeset/anthropic-prompt-caching.md @@ -0,0 +1,5 @@ +--- +"@nanocollective/nanocoder": minor +--- + +Added Anthropic prompt caching. The system prompt, tool schemas, and conversation history are now marked with cache breakpoints, so multi-turn sessions read the stable prefix back out of cache instead of paying full price for it every turn. Cost reporting is cache-aware throughout: `/usage` and the per-response indicator price cache reads and writes at their own rates and surface the cached token count. Opt out with `"promptCaching": false` on the provider config. Closes #888. diff --git a/source/ai-sdk-client/chat/chat-handler.spec.ts b/source/ai-sdk-client/chat/chat-handler.spec.ts index 727221b48..69a58154b 100644 --- a/source/ai-sdk-client/chat/chat-handler.spec.ts +++ b/source/ai-sdk-client/chat/chat-handler.spec.ts @@ -361,3 +361,121 @@ test('privacy: scrubs outgoing prompts and rehydrates the response at the histor t.is(content, 'Saved real@example.com'); t.false(content.includes('«')); }); + +function capturingModel(captured: {prompt?: unknown}): LanguageModel { + return { + specificationVersion: 'v3', + provider: 'anthropic', + modelId: 'claude-sonnet-4-5', + doStream: async (options: {prompt: unknown}) => { + captured.prompt = options.prompt; + return { + stream: new ReadableStream({ + start(controller) { + controller.enqueue({type: 'text-start', id: '0'}); + controller.enqueue({type: 'text-delta', id: '0', delta: 'ok'}); + controller.enqueue({type: 'text-end', id: '0'}); + controller.enqueue({ + type: 'finish', + finishReason: 'stop', + usage: {inputTokens: 5000, outputTokens: 10, totalTokens: 5010}, + }); + controller.close(); + }, + }), + }; + }, + } as unknown as LanguageModel; +} + +const anthropicConfig: AIProviderConfig = { + name: 'Anthropic', + type: 'anthropic', + sdkProvider: 'anthropic', + models: ['claude-sonnet-4-5'], + config: {apiKey: 'test-key'}, +}; + +const LONG_SYSTEM = 'SYSTEM '.repeat(1000); + +test('handleChat sends the system prompt as a cache-marked message on anthropic', async t => { + const captured: {prompt?: unknown} = {}; + await handleChat({ + model: capturingModel(captured), + currentModel: 'claude-sonnet-4-5', + providerConfig: anthropicConfig, + messages: [ + {role: 'system', content: LONG_SYSTEM}, + {role: 'user', content: 'hello'}, + ], + tools: {}, + callbacks: {}, + maxRetries: 0, + }); + + const prompt = captured.prompt as Array<{ + role: string; + providerOptions?: Record; + }>; + t.is(prompt[0]?.role, 'system'); + t.deepEqual(prompt[0]?.providerOptions, { + anthropic: {cacheControl: {type: 'ephemeral'}}, + }); + t.deepEqual(prompt[prompt.length - 1]?.providerOptions, { + anthropic: {cacheControl: {type: 'ephemeral'}}, + }); +}); + +test('handleChat does not emit the AI SDK system-in-messages warning', async t => { + const warnings: string[] = []; + const originalWarn = console.warn; + console.warn = (...args: unknown[]) => { + warnings.push(args.join(' ')); + }; + try { + await handleChat({ + model: capturingModel({}), + currentModel: 'claude-sonnet-4-5', + providerConfig: anthropicConfig, + messages: [ + {role: 'system', content: LONG_SYSTEM}, + {role: 'user', content: 'hello'}, + ], + tools: {}, + callbacks: {}, + maxRetries: 0, + }); + } finally { + console.warn = originalWarn; + } + t.false(warnings.some(w => w.includes('System messages in the prompt'))); +}); + +test('handleChat keeps the system string for non-anthropic providers', async t => { + const captured: {prompt?: unknown} = {}; + await handleChat({ + model: capturingModel(captured), + currentModel: 'test-model', + providerConfig: { + name: 'TestProvider', + type: 'openai', + models: ['test-model'], + config: {baseURL: 'https://api.test.com'}, + }, + messages: [ + {role: 'system', content: LONG_SYSTEM}, + {role: 'user', content: 'hello'}, + ], + tools: {}, + callbacks: {}, + maxRetries: 0, + }); + + const prompt = captured.prompt as Array<{ + role: string; + providerOptions?: Record; + }>; + t.is(prompt[0]?.role, 'system'); + t.is(prompt[0]?.providerOptions, undefined); + t.is(prompt[prompt.length - 1]?.providerOptions, undefined); +}); diff --git a/source/ai-sdk-client/chat/chat-handler.ts b/source/ai-sdk-client/chat/chat-handler.ts index f6b037f0f..f20d722a8 100644 --- a/source/ai-sdk-client/chat/chat-handler.ts +++ b/source/ai-sdk-client/chat/chat-handler.ts @@ -28,12 +28,18 @@ import { startMetrics, } from '@/utils/logging/performance.js'; import {getSafeMemory} from '@/utils/logging/safe-process.js'; -import {convertToModelMessages} from '../converters/message-converter.js'; +import { + convertToModelMessages, + withCacheBreakpoints, +} from '../converters/message-converter.js'; import {convertAISDKToolCalls} from '../converters/tool-converter.js'; import {extractRootError} from '../error-handling/error-extractor.js'; import {parseAPIError} from '../error-handling/error-parser.js'; import {isToolSupportError} from '../error-handling/tool-error-detector.js'; -import {buildProviderOptions} from './provider-options.js'; +import { + buildProviderOptions, + isPromptCachingEnabled, +} from './provider-options.js'; import { createOnStepFinishHandler, createPrepareStepHandler, @@ -187,7 +193,11 @@ export async function handleChat( } // Convert messages to AI SDK v5 ModelMessage format - const modelMessages = convertToModelMessages(finalNonSystemMessages); + const promptCaching = isPromptCachingEnabled(providerConfig); + const convertedMessages = convertToModelMessages(finalNonSystemMessages); + const modelMessages = promptCaching + ? withCacheBreakpoints(convertedMessages, finalSystemContent) + : convertedMessages; logger.debug('AI SDK request prepared', { messageCount: modelMessages.length, @@ -213,7 +223,10 @@ export async function handleChat( const result = streamText({ model, - ...(finalSystemContent ? {system: finalSystemContent} : {}), + ...(finalSystemContent && !promptCaching + ? {system: finalSystemContent} + : {}), + ...(promptCaching ? {allowSystemInMessages: true} : {}), messages: modelMessages, tools: aiTools, abortSignal: signal, @@ -478,6 +491,8 @@ export async function handleChat( inputTokens: usage.inputTokens, outputTokens: usage.outputTokens, totalTokens: usage.totalTokens, + cacheReadTokens: usage.inputTokenDetails?.cacheReadTokens, + cacheWriteTokens: usage.inputTokenDetails?.cacheWriteTokens, }, }; } catch (error) { diff --git a/source/ai-sdk-client/chat/provider-options.spec.ts b/source/ai-sdk-client/chat/provider-options.spec.ts index 383eed26f..a0111df71 100644 --- a/source/ai-sdk-client/chat/provider-options.spec.ts +++ b/source/ai-sdk-client/chat/provider-options.spec.ts @@ -1,6 +1,9 @@ import test from 'ava'; import type {AIProviderConfig, OpenRouterParameters} from '@/types/index'; -import {buildProviderOptions} from './provider-options.js'; +import { + buildProviderOptions, + isPromptCachingEnabled, +} from './provider-options.js'; function makeProvider( overrides: Partial = {}, @@ -276,3 +279,43 @@ test('OpenRouter extraBody alone is enough to emit providerOptions', t => { const result = buildProviderOptions(provider, '', undefined); t.deepEqual(result, {openrouter: {experimental_flag: true}}); }); + +test('isPromptCachingEnabled is on by default for the anthropic SDK', t => { + t.true( + isPromptCachingEnabled(makeProvider({sdkProvider: 'anthropic'})), + ); +}); + +test('isPromptCachingEnabled honours an explicit opt-out', t => { + t.false( + isPromptCachingEnabled( + makeProvider({sdkProvider: 'anthropic', promptCaching: false}), + ), + ); +}); + +test('isPromptCachingEnabled honours an explicit opt-in', t => { + t.true( + isPromptCachingEnabled( + makeProvider({sdkProvider: 'anthropic', promptCaching: true}), + ), + ); +}); + +test('isPromptCachingEnabled is off for every non-anthropic SDK', t => { + t.false(isPromptCachingEnabled(makeProvider())); + t.false( + isPromptCachingEnabled(makeProvider({sdkProvider: 'openai-compatible'})), + ); + t.false(isPromptCachingEnabled(makeProvider({sdkProvider: 'google'}))); + t.false(isPromptCachingEnabled(makeProvider({sdkProvider: 'chatgpt-codex'}))); + t.false(isPromptCachingEnabled(makeProvider({sdkProvider: 'github-copilot'}))); +}); + +test('isPromptCachingEnabled stays off for non-anthropic providers that opt in', t => { + t.false( + isPromptCachingEnabled( + makeProvider({sdkProvider: 'openai-compatible', promptCaching: true}), + ), + ); +}); diff --git a/source/ai-sdk-client/chat/provider-options.ts b/source/ai-sdk-client/chat/provider-options.ts index bae87a418..d967163ae 100644 --- a/source/ai-sdk-client/chat/provider-options.ts +++ b/source/ai-sdk-client/chat/provider-options.ts @@ -14,6 +14,15 @@ import {isOpenRouterProvider} from '../providers/openrouter.js'; */ export type ProviderOptions = Record>; +export function isPromptCachingEnabled( + providerConfig: AIProviderConfig, +): boolean { + return ( + providerConfig.sdkProvider === 'anthropic' && + providerConfig.promptCaching !== false + ); +} + /** * Build the `providerOptions` value for a streamText/generateText call. * diff --git a/source/ai-sdk-client/converters/message-converter.spec.ts b/source/ai-sdk-client/converters/message-converter.spec.ts index 6e2a10d5c..117ac473b 100644 --- a/source/ai-sdk-client/converters/message-converter.spec.ts +++ b/source/ai-sdk-client/converters/message-converter.spec.ts @@ -1,3 +1,4 @@ +import type {ModelMessage} from 'ai'; import test from 'ava'; import {MAX_TOOL_RESULT_CHARS} from '@/constants'; import type {Message} from '@/types/index'; @@ -5,6 +6,7 @@ import { convertToModelMessages, dropOrphanedToolResults, isEmptyAssistantMessage, + withCacheBreakpoints, } from './message-converter.js'; import type {TestableMessage} from '../types.js'; @@ -398,3 +400,106 @@ test('dropOrphanedToolResults drops a tool result lacking a tool_call_id', t => t.is(result.length, 1); t.is(result[0].role, 'user'); }); + +const BIG = 'S'.repeat(5000); +const BREAKPOINT = {anthropic: {cacheControl: {type: 'ephemeral'}}}; + +test('withCacheBreakpoints folds the system prompt in as the first message', t => { + const result = withCacheBreakpoints( + [{role: 'user', content: 'hi'}], + 'system text', + ); + t.is(result.length, 2); + t.is(result[0]?.role, 'system'); + t.is(result[0]?.content, 'system text'); +}); + +test('withCacheBreakpoints marks the system prompt and the last message', t => { + const result = withCacheBreakpoints( + [ + {role: 'user', content: BIG}, + {role: 'assistant', content: 'a'}, + {role: 'user', content: 'b'}, + ], + 'system text', + ); + t.deepEqual(result[0]?.providerOptions, BREAKPOINT); + t.is(result[1]?.providerOptions, undefined); + t.is(result[2]?.providerOptions, undefined); + t.deepEqual(result[3]?.providerOptions, BREAKPOINT); +}); + +test('withCacheBreakpoints emits at most two breakpoints', t => { + const messages = Array.from({length: 12}, (_, i) => ({ + role: 'user' as const, + content: `${BIG}${i}`, + })); + const marked = withCacheBreakpoints(messages, BIG).filter( + m => m.providerOptions !== undefined, + ); + t.is(marked.length, 2); +}); + +test('withCacheBreakpoints skips breakpoints below the cacheable minimum', t => { + const result = withCacheBreakpoints( + [{role: 'user', content: 'hi'}], + 'short system', + ); + t.is(result.length, 2); + t.is(result[0]?.providerOptions, undefined); + t.is(result[1]?.providerOptions, undefined); +}); + +test('withCacheBreakpoints marks only the last message when there is no system prompt', t => { + const result = withCacheBreakpoints( + [ + {role: 'user', content: BIG}, + {role: 'assistant', content: 'tail'}, + ], + '', + ); + t.is(result.length, 2); + t.is(result[0]?.providerOptions, undefined); + t.deepEqual(result[1]?.providerOptions, BREAKPOINT); +}); + +test('withCacheBreakpoints marks only the system prompt when there are no messages', t => { + const result = withCacheBreakpoints([], BIG); + t.is(result.length, 1); + t.deepEqual(result[0]?.providerOptions, BREAKPOINT); +}); + +test('withCacheBreakpoints returns an empty array for empty input', t => { + t.deepEqual(withCacheBreakpoints([], ''), []); +}); + +test('withCacheBreakpoints counts array content toward the threshold', t => { + const result = withCacheBreakpoints( + [ + { + role: 'tool', + content: [ + { + type: 'tool-result', + toolCallId: '1', + toolName: 'read_file', + output: {type: 'text', value: BIG}, + }, + ], + }, + ], + '', + ); + t.deepEqual(result[0]?.providerOptions, BREAKPOINT); +}); + +test('withCacheBreakpoints does not mutate its inputs', t => { + const messages: ModelMessage[] = [ + {role: 'user', content: BIG}, + {role: 'assistant', content: 'tail'}, + ]; + withCacheBreakpoints(messages, BIG); + t.is(messages.length, 2); + t.is(messages[0]?.providerOptions, undefined); + t.is(messages[1]?.providerOptions, undefined); +}); diff --git a/source/ai-sdk-client/converters/message-converter.ts b/source/ai-sdk-client/converters/message-converter.ts index 809605a96..5153c4c96 100644 --- a/source/ai-sdk-client/converters/message-converter.ts +++ b/source/ai-sdk-client/converters/message-converter.ts @@ -64,6 +64,52 @@ export function dropOrphanedToolResults(messages: Message[]): Message[] { return result; } +const CACHE_BREAKPOINT = { + anthropic: {cacheControl: {type: 'ephemeral'}}, +}; + +const MIN_CACHEABLE_CHARS = 4096; + +function messageChars(message: ModelMessage): number { + if (typeof message.content === 'string') { + return message.content.length; + } + if (!Array.isArray(message.content)) { + return 0; + } + return message.content.reduce((sum, part) => { + if (part.type === 'text') { + return sum + part.text.length; + } + return sum + JSON.stringify(part).length; + }, 0); +} + +function markCacheBreakpoint(message: ModelMessage): ModelMessage { + return {...message, providerOptions: CACHE_BREAKPOINT} as ModelMessage; +} + +export function withCacheBreakpoints( + messages: ModelMessage[], + systemContent: string, +): ModelMessage[] { + const system: ModelMessage[] = systemContent + ? [{role: 'system', content: systemContent}] + : []; + const totalChars = + systemContent.length + + messages.reduce((sum, message) => sum + messageChars(message), 0); + if (totalChars < MIN_CACHEABLE_CHARS) { + return [...system, ...messages]; + } + const marked = system.map(markCacheBreakpoint); + const lastIndex = messages.length - 1; + messages.forEach((message, index) => { + marked.push(index === lastIndex ? markCacheBreakpoint(message) : message); + }); + return marked; +} + /** * Convert our Message format to AI SDK v6 ModelMessage format * diff --git a/source/commands/usage.tsx b/source/commands/usage.tsx index 68a6a6325..9822ef931 100644 --- a/source/commands/usage.tsx +++ b/source/commands/usage.tsx @@ -24,6 +24,7 @@ import { calculateTokenBreakdown, calculateToolDefinitionsTokensFromDefs, } from '@/usage/calculator'; +import {priceTokens} from '@/usage/response-usage'; import {buildSystemPrompt, getLastBuiltPrompt} from '@/utils/prompt-builder'; export const usageCommand: Command = { @@ -199,10 +200,7 @@ export const usageCommand: Command = { snapshot.inputTokens != null && snapshot.outputTokens != null ) { - currentContextCost = - (pricing.input * snapshot.inputTokens + - pricing.output * snapshot.outputTokens) / - 1_000_000; + currentContextCost = priceTokens(pricing, snapshot); } else { currentContextCost = (pricing.input * breakdown.total) / 1_000_000; } @@ -228,23 +226,14 @@ export const usageCommand: Command = { output: NaN, }; - const knownInputCost = - record.inputTokens != null - ? (recordPricing.input * record.inputTokens) / 1_000_000 - : 0; - const knownOutputCost = - record.outputTokens != null - ? (recordPricing.output * record.outputTokens) / 1_000_000 - : 0; - const callCost = record.inputTokens != null && record.outputTokens != null - ? knownInputCost + knownOutputCost + ? priceTokens(recordPricing, record) : record.totalTokens != null ? (((recordPricing.input + recordPricing.output) / 2) * record.totalTokens) / 1_000_000 - : knownInputCost + knownOutputCost; + : priceTokens(recordPricing, record); cumulativeSession += callCost; perProvider[record.provider] = diff --git a/source/hooks/chat-handler/conversation/conversation-loop.tsx b/source/hooks/chat-handler/conversation/conversation-loop.tsx index 3d477b4e8..ca23e263f 100644 --- a/source/hooks/chat-handler/conversation/conversation-loop.tsx +++ b/source/hooks/chat-handler/conversation/conversation-loop.tsx @@ -624,6 +624,8 @@ export const processAssistantResponse = async ( inputTokens: usage.inputTokens, outputTokens: usage.outputTokens, totalTokens: usage.totalTokens, + cacheReadTokens: usage.cacheReadTokens, + cacheWriteTokens: usage.cacheWriteTokens, timestamp: Date.now(), }); } diff --git a/source/types/config.ts b/source/types/config.ts index a499a613b..f6b27d567 100644 --- a/source/types/config.ts +++ b/source/types/config.ts @@ -29,6 +29,8 @@ export interface AIProviderConfig { disableToolModels?: string[]; // List of model names to disable tools for // SDK provider package to use (default: 'openai-compatible') sdkProvider?: SdkProvider; + // Opt out of Anthropic prompt caching (enabled by default on that SDK). + promptCaching?: boolean; // Model mode defaults for this provider tune?: Partial; // OpenRouter-specific request body fields (provider routing, plugins, diff --git a/source/types/core.ts b/source/types/core.ts index ab038d3f4..d3b1754b1 100644 --- a/source/types/core.ts +++ b/source/types/core.ts @@ -153,6 +153,8 @@ export interface ApiUsage { inputTokens?: number; outputTokens?: number; totalTokens?: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; } export interface ApiUsageSnapshot extends ApiUsage { @@ -171,6 +173,8 @@ export interface ApiCallRecord { inputTokens?: number; outputTokens?: number; totalTokens?: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; timestamp: number; } diff --git a/source/types/usage.ts b/source/types/usage.ts index b2576d897..0e78abb43 100644 --- a/source/types/usage.ts +++ b/source/types/usage.ts @@ -19,6 +19,8 @@ export interface ResponseUsage { inputTokens?: number; outputTokens?: number; totalTokens?: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; /** Estimated cost of this API call in USD; omitted when pricing is unknown. */ cost?: number; } diff --git a/source/usage/format.spec.ts b/source/usage/format.spec.ts index ab4c16785..adcbab6bc 100644 --- a/source/usage/format.spec.ts +++ b/source/usage/format.spec.ts @@ -114,3 +114,23 @@ test('formatUsageIndicator returns null without usable token counts', t => { t.is(formatUsageIndicator({}), null); t.is(formatUsageIndicator({cost: 0.5}), null); }); + +test('formatUsageIndicator reports cached tokens when the provider read from cache', t => { + t.is( + formatUsageIndicator({ + inputTokens: 12_000, + outputTokens: 400, + cacheReadTokens: 9800, + cost: 0.02, + }), + 'Tokens: 12.4k | 9.8k cached | ~$0.02', + ); +}); + +test('formatUsageIndicator omits the cached segment when nothing was cached', t => { + t.is(formatUsageIndicator({totalTokens: 4200}), 'Tokens: 4.2k'); + t.is( + formatUsageIndicator({totalTokens: 4200, cacheReadTokens: 0}), + 'Tokens: 4.2k', + ); +}); diff --git a/source/usage/format.ts b/source/usage/format.ts index 058904387..aa42c59cb 100644 --- a/source/usage/format.ts +++ b/source/usage/format.ts @@ -78,6 +78,10 @@ export function formatUsageIndicator(usage: ResponseUsage): string | null { return null; } const parts = [`Tokens: ${formatCompactTokenCount(total)}`]; + const cacheRead = usage.cacheReadTokens; + if (Number.isFinite(cacheRead) && (cacheRead as number) > 0) { + parts.push(`${formatCompactTokenCount(cacheRead as number)} cached`); + } const cost = usage.cost != null ? formatCost(usage.cost) : null; if (cost) { parts.push(cost); diff --git a/source/usage/response-usage.spec.ts b/source/usage/response-usage.spec.ts index ef47bc8d0..af01def8b 100644 --- a/source/usage/response-usage.spec.ts +++ b/source/usage/response-usage.spec.ts @@ -1,5 +1,9 @@ import test from 'ava'; -import {buildResponseUsage, buildResponseUsageBounded} from './response-usage.js'; +import { + buildResponseUsage, + buildResponseUsageBounded, + priceTokens, +} from './response-usage.js'; console.log('\nresponse-usage.spec.ts'); @@ -130,3 +134,114 @@ test('buildResponseUsageBounded returns undefined when the provider reported not undefined, ); }); + +const cachePricing = async () => ({ + input: 3, + output: 15, + cache_read: 0.3, + cache_write: 3.75, +}); + +test('priceTokens matches the plain input/output rate when no cache tokens are reported', t => { + t.is( + priceTokens( + {input: 3, output: 15}, + {inputTokens: 1_000_000, outputTokens: 100_000}, + ), + 4.5, + ); +}); + +test('priceTokens bills cache reads at the discounted rate', t => { + t.is( + priceTokens( + {input: 3, output: 15, cache_read: 0.3, cache_write: 3.75}, + {inputTokens: 1_000_000, outputTokens: 0, cacheReadTokens: 800_000}, + ), + 0.84, + ); +}); + +test('priceTokens bills cache writes at the premium rate', t => { + t.is( + priceTokens( + {input: 3, output: 15, cache_read: 0.3, cache_write: 3.75}, + {inputTokens: 1_000_000, outputTokens: 0, cacheWriteTokens: 800_000}, + ), + 3.6, + ); +}); + +test('priceTokens falls back to the input rate when cache pricing is unknown', t => { + t.is( + priceTokens( + {input: 3, output: 15}, + { + inputTokens: 1_000_000, + outputTokens: 0, + cacheReadTokens: 500_000, + cacheWriteTokens: 250_000, + }, + ), + 3, + ); +}); + +test('priceTokens never charges negative uncached input', t => { + t.is( + priceTokens( + {input: 3, output: 15, cache_read: 0.3}, + {inputTokens: 1000, outputTokens: 0, cacheReadTokens: 5000}, + ), + 0.0015, + ); +}); + +test('priceTokens treats missing token fields as zero', t => { + t.is(priceTokens({input: 3, output: 15}, {}), 0); +}); + +test('buildResponseUsage prices a cache hit below the uncached equivalent', async t => { + const cached = await buildResponseUsage( + { + inputTokens: 1_000_000, + outputTokens: 100_000, + cacheReadTokens: 900_000, + }, + 'model', + cachePricing, + ); + const uncached = await buildResponseUsage( + {inputTokens: 1_000_000, outputTokens: 100_000}, + 'model', + cachePricing, + ); + t.true((cached?.cost as number) < (uncached?.cost as number)); + t.is(cached?.cost, 2.07); +}); + +test('buildResponseUsage surfaces the cache token counts', async t => { + const result = await buildResponseUsage( + { + inputTokens: 5000, + outputTokens: 100, + cacheReadTokens: 4000, + cacheWriteTokens: 500, + }, + 'model', + cachePricing, + ); + t.is(result?.cacheReadTokens, 4000); + t.is(result?.cacheWriteTokens, 500); +}); + +test('buildResponseUsage leaves non-caching reports byte-identical to before', async t => { + const result = await buildResponseUsage( + {inputTokens: 1_000_000, outputTokens: 100_000}, + 'model', + stubPricing, + ); + t.is(result?.cost, 4.5); + t.is(result?.cacheReadTokens, undefined); + t.is(result?.cacheWriteTokens, undefined); +}); diff --git a/source/usage/response-usage.ts b/source/usage/response-usage.ts index 53d8e2154..fd285fc98 100644 --- a/source/usage/response-usage.ts +++ b/source/usage/response-usage.ts @@ -8,9 +8,38 @@ import {getModelPricing} from '@/models/index'; import type {ApiUsage} from '@/types/core'; import type {ResponseUsage} from '@/types/usage'; -type PricingLookup = ( - model: string, -) => Promise<{input: number; output: number} | null>; +export interface TokenPricing { + input: number; + output: number; + cache_read?: number; + cache_write?: number; +} + +type PricingLookup = (model: string) => Promise; + +export function priceTokens( + pricing: TokenPricing, + usage: { + inputTokens?: number; + outputTokens?: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; + }, +): number { + const cacheRead = usage.cacheReadTokens ?? 0; + const cacheWrite = usage.cacheWriteTokens ?? 0; + const uncachedInput = Math.max( + 0, + (usage.inputTokens ?? 0) - cacheRead - cacheWrite, + ); + return ( + (pricing.input * uncachedInput + + (pricing.cache_read ?? pricing.input) * cacheRead + + (pricing.cache_write ?? pricing.input) * cacheWrite + + pricing.output * (usage.outputTokens ?? 0)) / + 1_000_000 + ); +} /** * Extract the provider-reported token counts, or undefined when the report @@ -32,6 +61,8 @@ function toReportedUsage( inputTokens: usage.inputTokens, outputTokens: usage.outputTokens, totalTokens: usage.totalTokens, + cacheReadTokens: usage.cacheReadTokens, + cacheWriteTokens: usage.cacheWriteTokens, }; } @@ -65,10 +96,7 @@ export async function buildResponseUsage( (usage.outputTokens as number) > 0 || !(usage.totalTokens && usage.totalTokens > 0)); if (hasUsableSplit) { - cost = - (pricing.input * (usage.inputTokens as number) + - pricing.output * (usage.outputTokens as number)) / - 1_000_000; + cost = priceTokens(pricing, usage); } else if (Number.isFinite(usage.totalTokens)) { // Lump-sum reports can't be split into input/output, so average // the two rates — same approximation the /usage command uses.