Skip to content

Commit a1fb74d

Browse files
committed
1 parent ba46d1f commit a1fb74d

2 files changed

Lines changed: 51 additions & 1 deletion

File tree

‎src/api/providers/__tests__/openai-usage-tracking.spec.ts‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,31 @@ describe("OpenAiHandler with usage tracking fix", () => {
181181
})
182182
})
183183

184+
it("should report OpenAI-compatible cached prompt tokens", async () => {
185+
mockCreate.mockImplementationOnce(async () =>
186+
asyncStreamFrom([
187+
{
188+
choices: [{ delta: { content: "Cached response" }, index: 0 }],
189+
usage: {
190+
prompt_tokens: 5_053,
191+
completion_tokens: 16,
192+
total_tokens: 5_069,
193+
prompt_tokens_details: { cached_tokens: 4_864 },
194+
},
195+
},
196+
]),
197+
)
198+
199+
const chunks = await collectStream(handler.createMessage(systemPrompt, messages))
200+
201+
expect(chunks).toContainEqual({
202+
type: "usage",
203+
inputTokens: 5_053,
204+
outputTokens: 16,
205+
cacheReadTokens: 4_864,
206+
})
207+
})
208+
184209
it("should handle case where no usage is provided", async () => {
185210
// Override the mock for this specific test
186211
mockCreate.mockImplementationOnce(async (options) => {
@@ -212,4 +237,27 @@ describe("OpenAiHandler with usage tracking fix", () => {
212237
expect(usageChunks).toHaveLength(0)
213238
})
214239
})
240+
241+
it("should report cached prompt tokens from a non-streaming response", async () => {
242+
const nonStreamingHandler = new OpenAiHandler({ ...mockOptions, openAiStreamingEnabled: false })
243+
mockCreate.mockImplementationOnce(async () => ({
244+
id: "test-completion",
245+
choices: [{ message: { role: "assistant", content: "Cached response" } }],
246+
usage: {
247+
prompt_tokens: 4_621,
248+
completion_tokens: 16,
249+
total_tokens: 4_637,
250+
prompt_tokens_details: { cached_tokens: 4_608 },
251+
},
252+
}))
253+
254+
const chunks = await collectStream(nonStreamingHandler.createMessage("system prompt", []))
255+
256+
expect(chunks).toContainEqual({
257+
type: "usage",
258+
inputTokens: 4_621,
259+
outputTokens: 16,
260+
cacheReadTokens: 4_608,
261+
})
262+
})
215263
})

‎src/api/providers/openai.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -282,12 +282,14 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
282282
}
283283

284284
protected processUsageMetrics(usage: any, _modelInfo?: ModelInfo): ApiStreamUsageChunk {
285+
const cacheReadTokens = usage?.cache_read_input_tokens ?? usage?.prompt_tokens_details?.cached_tokens
286+
285287
return {
286288
type: "usage",
287289
inputTokens: usage?.prompt_tokens || 0,
288290
outputTokens: usage?.completion_tokens || 0,
289291
cacheWriteTokens: usage?.cache_creation_input_tokens || undefined,
290-
cacheReadTokens: usage?.cache_read_input_tokens || undefined,
292+
cacheReadTokens: cacheReadTokens || undefined,
291293
}
292294
}
293295

0 commit comments

Comments
 (0)