mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 08:57:57 +00:00
fix(gateway): emit OpenAI-standard token usage in responses
- TokenUsage.MarshalJSON now emits both standard (prompt_tokens, completion_tokens, total_tokens) and legacy (prompt, completion, total) keys, so OpenAI-compatible clients (DSH, DevEco Code, etc.) can read token usage. - ChatChunk gains an optional Usage field; stream responses now send a final usage chunk (empty choices) before [DONE]. Refs: usage not visible in clients because the gateway serialized only the internal short keys and never emitted a streaming usage chunk.
This commit is contained in:
@ -53,11 +53,14 @@ type RespMessage struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type ChatChunk struct {
|
type ChatChunk struct {
|
||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
Object string `json:"object"`
|
Object string `json:"object"`
|
||||||
Created int64 `json:"created"`
|
Created int64 `json:"created"`
|
||||||
Model string `json:"model"`
|
Model string `json:"model"`
|
||||||
Choices []ChunkChoice `json:"choices"`
|
Choices []ChunkChoice `json:"choices"`
|
||||||
|
// Usage is sent in the final chunk of a stream (empty choices) so
|
||||||
|
// OpenAI-compatible clients can read token usage.
|
||||||
|
Usage *types.TokenUsage `json:"usage,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
type ChunkChoice struct {
|
type ChunkChoice struct {
|
||||||
@ -653,6 +656,19 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
|||||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
||||||
})
|
})
|
||||||
|
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
|
||||||
|
if rec.Prompt+rec.Compl > 0 {
|
||||||
|
usage := types.TokenUsage{
|
||||||
|
Prompt: int(rec.Prompt),
|
||||||
|
Completion: int(rec.Compl),
|
||||||
|
Total: int(rec.Prompt + rec.Compl),
|
||||||
|
}
|
||||||
|
send(ChatChunk{
|
||||||
|
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||||
|
Choices: []ChunkChoice{},
|
||||||
|
Usage: &usage,
|
||||||
|
})
|
||||||
|
}
|
||||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||||
if flusher != nil {
|
if flusher != nil {
|
||||||
flusher.Flush()
|
flusher.Flush()
|
||||||
@ -798,6 +814,19 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
|||||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
||||||
})
|
})
|
||||||
|
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
|
||||||
|
if rec.Prompt+rec.Compl > 0 {
|
||||||
|
usage := types.TokenUsage{
|
||||||
|
Prompt: int(rec.Prompt),
|
||||||
|
Completion: int(rec.Compl),
|
||||||
|
Total: int(rec.Prompt + rec.Compl),
|
||||||
|
}
|
||||||
|
send(ChatChunk{
|
||||||
|
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||||
|
Choices: []ChunkChoice{},
|
||||||
|
Usage: &usage,
|
||||||
|
})
|
||||||
|
}
|
||||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||||
if flusher != nil {
|
if flusher != nil {
|
||||||
flusher.Flush()
|
flusher.Flush()
|
||||||
|
|||||||
@ -85,6 +85,28 @@ type TokenUsage struct {
|
|||||||
Total int `json:"total"`
|
Total int `json:"total"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// MarshalJSON emits both the legacy short keys (prompt/completion/total, used
|
||||||
|
// by the internal unified representation and older clients) and the OpenAI
|
||||||
|
// standard keys (prompt_tokens/completion_tokens/total_tokens). Standard
|
||||||
|
// clients such as DSH and DevEco Code read the *_tokens fields.
|
||||||
|
func (t TokenUsage) MarshalJSON() ([]byte, error) {
|
||||||
|
return json.Marshal(struct {
|
||||||
|
PromptTokens int `json:"prompt_tokens"`
|
||||||
|
CompletionTokens int `json:"completion_tokens"`
|
||||||
|
TotalTokens int `json:"total_tokens"`
|
||||||
|
Prompt int `json:"prompt"`
|
||||||
|
Completion int `json:"completion"`
|
||||||
|
Total int `json:"total"`
|
||||||
|
}{
|
||||||
|
PromptTokens: t.Prompt,
|
||||||
|
CompletionTokens: t.Completion,
|
||||||
|
TotalTokens: t.Total,
|
||||||
|
Prompt: t.Prompt,
|
||||||
|
Completion: t.Completion,
|
||||||
|
Total: t.Total,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
type ImageData struct {
|
type ImageData struct {
|
||||||
B64JSON string `json:"b64_json,omitempty"`
|
B64JSON string `json:"b64_json,omitempty"`
|
||||||
URL string `json:"url,omitempty"`
|
URL string `json:"url,omitempty"`
|
||||||
|
|||||||
Reference in New Issue
Block a user