jianf/ModelRouter
1
0
Fork 0
mirror of https://gitcode.com/JianFeeeee/ModelRouter.git synced 2026-09-20 08:57:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): pass through upstream token usage in streams for all adapters

Browse Source 
The prior usage-passthrough fix only covered openai/opencode; the same
empty-choices+usage drop bug remained in the 5 sibling OpenAI-compatible
adapters, and non-OpenAI providers (anthropic/gemini/ollama) never surfaced
streaming usage at all.

- deepseek/github/groq/kimicode/mistral: preserve usage on empty-choices
  chunks and attach it to normal chunks (same pattern as openai.lua)
- anthropic: emit usage from message_start (prompt) and message_delta
  (completion); gateway merges split usage additively
- gemini: read usageMetadata in the stream path
- ollama: fix non-streaming key (usage -> token_usage, matches
  UnifiedResponse json tag) and read prompt_eval_count/eval_count;
  surface counts from the done stream chunk
- gateway: mergeUsage combines usage across chunks (non-zero fields win,
  total recomputed from prompt+completion) so split usage doesn't lose
  the prompt half; single-chunk case (OpenAI) preserved exactly
- usage-only chunks: done=false (no redundant terminal stop), matching
  the Go fallback standardSSEChunk
This commit is contained in:
jianf
2026-08-18 23:24:01 +08:00
parent cb9c025952
commit 3069cfce4e
11 changed files with 230 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
internal/gateway/chat.go
View File

@ -580,6 +580,32 @@ func (g *Gateway) writeRec(rec *Req) {
g.stats.Record(*rec)
}
// mergeUsage combines token usage across stream chunks additively. Some
// providers split usage across chunks (e.g. Anthropic reports prompt tokens
// in message_start and the final completion tokens in message_delta); a plain
// "last non-nil wins" would discard the prompt half. Non-zero fields from cur
// override prev; total is recomputed from the merged parts so a partial later
// chunk can't shrink it. For the common single-chunk case (OpenAI's terminal
// empty-choices+usage chunk) upstream totals are preserved exactly.
func mergeUsage(prev, cur *types.TokenUsage) *types.TokenUsage {
if prev == nil {
u := *cur
if u.Total == 0 && (u.Prompt > 0 || u.Completion > 0) {
u.Total = u.Prompt + u.Completion
}
return &u
}
out := *prev
if cur.Prompt > 0 {
out.Prompt = cur.Prompt
}
if cur.Completion > 0 {
out.Completion = cur.Completion
}
out.Total = out.Prompt + out.Completion
return &out
}
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
rec.LatMs = 0
t0 := time.Now()
@ -632,7 +658,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
var lastUsage *types.TokenUsage
for ck := range chunks {
if ck.Usage != nil {
lastUsage = ck.Usage
lastUsage = mergeUsage(lastUsage, ck.Usage)
}
chunk := ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
@ -802,7 +828,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
var lastUsage *types.TokenUsage
for ck := range chunks {
if ck.Usage != nil {
lastUsage = ck.Usage
lastUsage = mergeUsage(lastUsage, ck.Usage)
}
chunk := ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,

32
internal/lua/adapters/anthropic.lua
View File

@ -99,9 +99,37 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if chunk.type == "message_start" then return "" end
if chunk.type == "message_start" then
local uses = nil
if chunk.message and type(chunk.message.usage) == "table" then
local u = chunk.message.usage
local p = u.input_tokens or 0
local c = u.output_tokens or 0
if p > 0 or c > 0 then
uses = { prompt = p, completion = c, total = p + c }
end
end
if uses ~= nil then
return json.encode({ usage = uses, done = false })
end
return ""
end
if chunk.type == "message_delta" then
return json.encode({ content = "", done = (chunk.delta and chunk.delta.stop_reason ~= nil) })
local uses = nil
if type(chunk.usage) == "table" then
local u = chunk.usage
local p = u.input_tokens or 0
local c = u.output_tokens or 0
if p > 0 or c > 0 then
uses = { prompt = p, completion = c, total = p + c }
end
end
local done = (chunk.delta and chunk.delta.stop_reason ~= nil)
if uses ~= nil then
-- completion is final here; prompt is merged from message_start
return json.encode({ content = "", done = done, usage = uses })
end
return json.encode({ content = "", done = done })
end
if chunk.type == "content_block_start" and chunk.content_block
and chunk.content_block.type == "tool_use" then

25
internal/lua/adapters/deepseek.lua
View File

@ -82,7 +82,27 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 then return "" end
-- OpenAI-style streams may attach usage to a chunk with empty choices
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
local uses = nil
if type(chunk.usage) == "table" then
uses = {
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
total = chunk.usage.total_tokens or chunk.usage.total or 0,
}
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
@ -90,6 +110,9 @@ function adapter.transform_stream_chunk(raw_chunk)
content = delta.content or "",
done = (fr ~= nil)
}
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end

22
internal/lua/adapters/gemini.lua
View File

@ -91,7 +91,24 @@ function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.candidates or #chunk.candidates == 0 then return "" end
-- Gemini attaches usageMetadata to the final chunk (alongside or after
-- candidates). Keys map to Go's TokenUsage json tags (prompt/...).
local uses = nil
if type(chunk.usageMetadata) == "table" then
local p = chunk.usageMetadata.promptTokenCount or 0
local c = chunk.usageMetadata.candidatesTokenCount or 0
local t = chunk.usageMetadata.totalTokenCount or 0
if p > 0 or c > 0 or t > 0 then
uses = { prompt = p, completion = c, total = t }
end
end
if not chunk.candidates or #chunk.candidates == 0 then
if uses ~= nil then
return json.encode({ usage = uses, done = false })
end
return ""
end
local cand = chunk.candidates[1]
local unified = { content = "", done = (cand.finishReason ~= nil) }
local reasoning = ""
@ -117,6 +134,9 @@ function adapter.transform_stream_chunk(raw_chunk)
end
if reasoning ~= "" then unified.reasoning_content = reasoning end
if #tools > 0 then unified.tool_calls = tools end
if uses ~= nil then
unified.usage = uses
end
return json.encode(unified)
end

25
internal/lua/adapters/github.lua
View File

@ -63,7 +63,27 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 then return "" end
-- OpenAI-style streams may attach usage to a chunk with empty choices
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
local uses = nil
if type(chunk.usage) == "table" then
uses = {
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
total = chunk.usage.total_tokens or chunk.usage.total or 0,
}
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
@ -71,6 +91,9 @@ function adapter.transform_stream_chunk(raw_chunk)
content = delta.content or "",
done = (fr ~= nil)
}
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end

25
internal/lua/adapters/groq.lua
View File

@ -62,7 +62,27 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 then return "" end
-- OpenAI-style streams may attach usage to a chunk with empty choices
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
local uses = nil
if type(chunk.usage) == "table" then
uses = {
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
total = chunk.usage.total_tokens or chunk.usage.total or 0,
}
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
@ -70,6 +90,9 @@ function adapter.transform_stream_chunk(raw_chunk)
content = delta.content or "",
done = (fr ~= nil)
}
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end

25
internal/lua/adapters/kimicode.lua
View File

@ -95,7 +95,27 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 then return "" end
-- OpenAI-style streams may attach usage to a chunk with empty choices
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
local uses = nil
if type(chunk.usage) == "table" then
uses = {
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
total = chunk.usage.total_tokens or chunk.usage.total or 0,
}
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
@ -103,6 +123,9 @@ function adapter.transform_stream_chunk(raw_chunk)
content = delta.content or "",
done = (fr ~= nil)
}
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end

25
internal/lua/adapters/mistral.lua
View File

@ -62,7 +62,27 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 then return "" end
-- OpenAI-style streams may attach usage to a chunk with empty choices
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
local uses = nil
if type(chunk.usage) == "table" then
uses = {
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
total = chunk.usage.total_tokens or chunk.usage.total or 0,
}
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
@ -70,6 +90,9 @@ function adapter.transform_stream_chunk(raw_chunk)
content = delta.content or "",
done = (fr ~= nil)
}
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end

27
internal/lua/adapters/ollama.lua
View File

@ -53,11 +53,14 @@ function adapter.transform_response(raw_body)
local ok, resp = pcall(json.decode, raw_body)
if not ok then return raw_body end
local p = resp.prompt_eval_count or 0
local c = resp.eval_count or 0
local unified = {
content = "",
finish_reason = resp.done_reason or "",
tool_calls = {},
usage = { prompt = 0, completion = 0, total = 0 }
-- key must be token_usage to match Go's UnifiedResponse json tag
token_usage = { prompt = p, completion = c, total = p + c }
}
if resp.message then
@ -70,12 +73,32 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.message then return "" end
-- Ollama's terminal chunk (done=true) carries token counts but may omit
-- message; pass them through so the gateway emits real usage.
local uses = nil
if chunk.done then
local p = chunk.prompt_eval_count or 0
local c = chunk.eval_count or 0
if p > 0 or c > 0 then
uses = { prompt = p, completion = c, total = p + c }
end
end
if not chunk.message then
if uses ~= nil then
return json.encode({ content = "", done = true, usage = uses })
end
return ""
end
local unified = {
content = chunk.message.content or "",
done = chunk.done or false
}
if uses ~= nil then
unified.usage = uses
end
if chunk.message.reasoning_content then
unified.reasoning_content = chunk.message.reasoning_content
end

4
internal/lua/adapters/openai.lua
View File

@ -81,7 +81,9 @@ function adapter.transform_stream_chunk(raw_chunk)
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
return json.encode({ usage = uses, done = (chunk.usage ~= nil) })
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end

4
internal/lua/adapters/opencode.lua
View File

@ -127,7 +127,9 @@ function adapter.transform_stream_chunk(raw_chunk)
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
return json.encode({ usage = uses, done = (chunk.usage ~= nil) })
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end