From a2d2712e377e47bd514c686efc016bf5b78b9b87 Mon Sep 17 00:00:00 2001 From: Abhishek Divekar Date: Tue, 23 Jun 2026 01:44:37 +0530 Subject: [PATCH] fix(stream): estimate input tokens when upstream reports prompt_tokens=0 (#4615) Integrated into release/v3.8.34 (rebuilt onto tip) --- open-sse/utils/stream.ts | 20 +++ ...stream-prompt-tokens-zero-upstream.test.ts | 114 ++++++++++++++++++ 2 files changed, 134 insertions(+) create mode 100644 tests/unit/stream-prompt-tokens-zero-upstream.test.ts diff --git a/open-sse/utils/stream.ts b/open-sse/utils/stream.ts index 3e51ea4e91..6d54886840 100644 --- a/open-sse/utils/stream.ts +++ b/open-sse/utils/stream.ts @@ -1545,6 +1545,26 @@ export function createSSEStream(options: StreamOptions = {}) { if (Array.isArray(parsed.choices) && parsed.choices.length === 0) { const emptyChoicesUsage = extractUsage(parsed) ?? parsed.usage; if (hasValidUsage(emptyChoicesUsage)) { + // Some upstreams (e.g. Ollama Cloud) emit prompt_tokens: 0 + // even when input was sent — they simply don't count input + // tokens. When we have a non-zero output but zero input, + // estimate the real input token count from the request body. + if ( + emptyChoicesUsage && + typeof emptyChoicesUsage === "object" && + !Array.isArray(emptyChoicesUsage) && + emptyChoicesUsage.completion_tokens > 0 + ) { + const pt = emptyChoicesUsage.prompt_tokens ?? 0; + if (pt === 0) { + const estimated = estimateUsage(body, totalContentLength, FORMATS.OPENAI); + if (estimated?.prompt_tokens > 0) { + emptyChoicesUsage.prompt_tokens = estimated.prompt_tokens; + emptyChoicesUsage.total_tokens = + (emptyChoicesUsage.total_tokens ?? 0) + estimated.prompt_tokens; + } + } + } usage = emptyChoicesUsage; output = `data: ${JSON.stringify(parsed)}\n\n`; injectedUsage = true; diff --git a/tests/unit/stream-prompt-tokens-zero-upstream.test.ts b/tests/unit/stream-prompt-tokens-zero-upstream.test.ts new file mode 100644 index 0000000000..dc579dc2b0 --- /dev/null +++ b/tests/unit/stream-prompt-tokens-zero-upstream.test.ts @@ -0,0 +1,114 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-stream-zero-prompt-")); +process.env.DATA_DIR = TEST_DATA_DIR; +const core = await import("../../src/lib/db/core.ts"); + +const { createSSEStream } = await import("../../open-sse/utils/stream.ts"); +const { FORMATS } = await import("../../open-sse/translator/formats.ts"); + +const textEncoder = new TextEncoder(); + +async function readTransformed(chunks, options) { + const source = new ReadableStream({ + start(controller) { + for (const chunk of chunks) { + controller.enqueue(textEncoder.encode(chunk)); + } + controller.close(); + }, + }); + + return new Response(source.pipeThrough(createSSEStream(options))).text(); +} + +test("createSSEStream passthrough estimates input tokens when upstream reports prompt_tokens=0 (Ollama Cloud pattern)", async () => { + let onCompletePayload = null; + const body = { messages: [{ role: "user", content: "Write a 500-line research report" }] }; + const text = await readTransformed( + [ + `data: ${JSON.stringify({ + id: "chatcmpl_ollama", + object: "chat.completion.chunk", + created: 1, + model: "minimax-m3", + system_fingerprint: "fp_ollama", + choices: [{ index: 0, delta: { content: "Here" } }], + })}\n\n`, + `data: ${JSON.stringify({ + id: "chatcmpl_ollama", + object: "chat.completion.chunk", + created: 1, + model: "minimax-m3", + system_fingerprint: "fp_ollama", + choices: [{ index: 0, delta: { content: " is" } }], + })}\n\n`, + `data: ${JSON.stringify({ + id: "chatcmpl_ollama", + object: "chat.completion.chunk", + created: 1, + model: "minimax-m3", + system_fingerprint: "fp_ollama", + choices: [{ index: 0, delta: { content: " the" } }], + })}\n\n`, + `data: ${JSON.stringify({ + id: "chatcmpl_ollama", + object: "chat.completion.chunk", + created: 1, + model: "minimax-m3", + system_fingerprint: "fp_ollama", + choices: [{ index: 0, delta: { content: " response" } }], + })}\n\n`, + // Ollama Cloud returns usage with prompt_tokens: 0 — it doesn't count input tokens + `data: ${JSON.stringify({ + id: "chatcmpl_ollama", + object: "chat.completion.chunk", + created: 1, + model: "minimax-m3", + system_fingerprint: "fp_ollama", + choices: [], + usage: { prompt_tokens: 0, completion_tokens: 10, total_tokens: 0 }, + })}\n\n`, + `data: [DONE]\n\n`, + ], + { + mode: "passthrough", + sourceFormat: FORMATS.OPENAI, + provider: "ollamacloud", + model: "minimax-m3", + body, + onComplete(payload) { + onCompletePayload = payload; + }, + } + ); + + // The fix should estimate input tokens from the request body instead of showing 0 + assert.equal( + onCompletePayload.responseBody.usage.completion_tokens, + 10, + "completion tokens should be unchanged" + ); + assert.ok( + onCompletePayload.responseBody.usage.prompt_tokens > 0, + `prompt_tokens should be estimated (> 0), got ${onCompletePayload.responseBody.usage.prompt_tokens}` + ); + assert.equal( + onCompletePayload.responseBody.usage.total_tokens, + onCompletePayload.responseBody.usage.prompt_tokens + 10, + "total_tokens should equal prompt_tokens + completion_tokens" + ); +}); + +test.after(() => { + core.resetDbInstance(); + if (fs.existsSync(TEST_DATA_DIR)) { + for (const entry of fs.readdirSync(TEST_DATA_DIR)) { + fs.rmSync(path.join(TEST_DATA_DIR, entry), { recursive: true, force: true }); + } + } +});