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:
2026-08-18 18:14:37 +08:00
parent 98c08d51c4
commit 6f1c806591
2 changed files with 56 additions and 5 deletions

View File

@ -53,11 +53,14 @@ type RespMessage struct {
}
type ChatChunk struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []ChunkChoice `json:"choices"`
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
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 {
@ -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,
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")
if flusher != nil {
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,
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")
if flusher != nil {
flusher.Flush()

View File

@ -85,6 +85,28 @@ type TokenUsage struct {
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 {
B64JSON string `json:"b64_json,omitempty"`
URL string `json:"url,omitempty"`