diff --git a/open-sse/utils/usageTracking.ts b/open-sse/utils/usageTracking.ts index 8e3e969c4e..f1758b7030 100644 --- a/open-sse/utils/usageTracking.ts +++ b/open-sse/utils/usageTracking.ts @@ -402,6 +402,18 @@ export function extractUsage(chunk) { }); } + // Ollama NDJSON format (raw from provider, before translation) + // Ollama sends: { "model": "...", "done": true, "prompt_eval_count": N, "eval_count": M } + if (chunk.done === true && typeof chunk.prompt_eval_count === "number") { + const promptEvalCount = chunk.prompt_eval_count || 0; + const evalCount = chunk.eval_count || 0; + return normalizeUsage({ + prompt_tokens: promptEvalCount, + completion_tokens: evalCount, + total_tokens: promptEvalCount + evalCount, + }); + } + return null; } diff --git a/tests/unit/usage-extractor.test.ts b/tests/unit/usage-extractor.test.ts index 6c7bdf5273..e156a989af 100644 --- a/tests/unit/usage-extractor.test.ts +++ b/tests/unit/usage-extractor.test.ts @@ -332,3 +332,45 @@ test("extractUsage reads flat cached_tokens and reasoning_tokens from streaming assert.equal(usage.cached_tokens, 192); assert.equal(usage.reasoning_tokens, 49); }); + +// ── Ollama raw NDJSON streaming usage ── +// Ollama sends a final NDJSON line { done: true, prompt_eval_count, eval_count } +// (raw from the provider, before any OpenAI translation). Without a dedicated +// branch, extractUsage returns null and Ollama streaming usage is dropped. + +test("extractUsage reads Ollama raw NDJSON final chunk (done + prompt_eval_count/eval_count)", () => { + const usage = extractUsage({ + model: "llama3.1", + done: true, + prompt_eval_count: 26, + eval_count: 298, + }); + + assert.ok(usage, "expected usage to be extracted from the Ollama final chunk"); + assert.equal(usage.prompt_tokens, 26); + assert.equal(usage.completion_tokens, 298); + assert.equal(usage.total_tokens, 324); +}); + +test("extractUsage defaults missing Ollama eval counts to zero", () => { + const usage = extractUsage({ + model: "llama3.1", + done: true, + prompt_eval_count: 12, + }); + + assert.ok(usage, "expected usage to be extracted even with only prompt_eval_count"); + assert.equal(usage.prompt_tokens, 12); + assert.equal(usage.completion_tokens, 0); + assert.equal(usage.total_tokens, 12); +}); + +test("extractUsage ignores non-final Ollama NDJSON chunks (done=false)", () => { + const usage = extractUsage({ + model: "llama3.1", + done: false, + response: "partial", + }); + + assert.equal(usage, null); +});