fix(ai): preserve cache_write_tokens in completions stream usage closes #2802
This commit is contained in:
@@ -734,24 +734,34 @@ function parseChunkUsage(
|
||||
rawUsage: {
|
||||
prompt_tokens?: number;
|
||||
completion_tokens?: number;
|
||||
prompt_tokens_details?: { cached_tokens?: number };
|
||||
prompt_tokens_details?: { cached_tokens?: number; cache_write_tokens?: number };
|
||||
completion_tokens_details?: { reasoning_tokens?: number };
|
||||
},
|
||||
model: Model<"openai-completions">,
|
||||
): AssistantMessage["usage"] {
|
||||
const cachedTokens = rawUsage.prompt_tokens_details?.cached_tokens || 0;
|
||||
const promptTokens = rawUsage.prompt_tokens || 0;
|
||||
const reportedCachedTokens = rawUsage.prompt_tokens_details?.cached_tokens || 0;
|
||||
const cacheWriteTokens = rawUsage.prompt_tokens_details?.cache_write_tokens || 0;
|
||||
const reasoningTokens = rawUsage.completion_tokens_details?.reasoning_tokens || 0;
|
||||
// OpenAI includes cached tokens in prompt_tokens, so subtract to get non-cached input
|
||||
const input = (rawUsage.prompt_tokens || 0) - cachedTokens;
|
||||
|
||||
// Normalize to pi-ai semantics:
|
||||
// - cacheRead: hits from cache created by previous requests only
|
||||
// - cacheWrite: tokens written to cache in this request
|
||||
// Some OpenAI-compatible providers (observed on OpenRouter) report cached_tokens
|
||||
// as (previous hits + current writes). In that case, remove cacheWrite from cacheRead.
|
||||
const cacheReadTokens =
|
||||
cacheWriteTokens > 0 ? Math.max(0, reportedCachedTokens - cacheWriteTokens) : reportedCachedTokens;
|
||||
|
||||
const input = Math.max(0, promptTokens - cacheReadTokens - cacheWriteTokens);
|
||||
// Compute totalTokens ourselves since we add reasoning_tokens to output
|
||||
// and some providers (e.g., Groq) don't include them in total_tokens
|
||||
const outputTokens = (rawUsage.completion_tokens || 0) + reasoningTokens;
|
||||
const usage: AssistantMessage["usage"] = {
|
||||
input,
|
||||
output: outputTokens,
|
||||
cacheRead: cachedTokens,
|
||||
cacheWrite: 0,
|
||||
totalTokens: input + outputTokens + cachedTokens,
|
||||
cacheRead: cacheReadTokens,
|
||||
cacheWrite: cacheWriteTokens,
|
||||
totalTokens: input + outputTokens + cacheReadTokens + cacheWriteTokens,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
calculateCost(model, usage);
|
||||
|
||||
Reference in New Issue
Block a user