Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions open-sse/utils/usageTracking.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Comment on lines +408 to +409

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To prevent potential string concatenation or unexpected type coercion (for example, if eval_count is returned as a string or non-numeric value by an Ollama-compatible provider, which would result in 26 + "298" = "26298" total tokens), we should defensively ensure that eval_count is a number. Since prompt_eval_count is already verified to be a number in the if condition, we can also simplify its assignment.

Suggested change
const promptEvalCount = chunk.prompt_eval_count || 0;
const evalCount = chunk.eval_count || 0;
const promptEvalCount = chunk.prompt_eval_count;
const evalCount = typeof chunk.eval_count === "number" ? chunk.eval_count : 0;

return normalizeUsage({
prompt_tokens: promptEvalCount,
completion_tokens: evalCount,
total_tokens: promptEvalCount + evalCount,
});
}

return null;
}

Expand Down
42 changes: 42 additions & 0 deletions tests/unit/usage-extractor.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
Loading