mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-14 06:17:41 +00:00
fix: OpenAI-compat SSE streaming drops usage data (tokens always 0) (#52)
When stream_options.include_usage is enabled, the OpenAI API sends the final usage chunk with an empty choices array. The streaming parser skipped chunks with no choices before checking for usage data, causing token counts to always be 0 in traces. Fix: extract usage data before the empty-choices skip. Also fixes two related SSE parsing issues: - SSE data prefix now handles both "data: " and "data:" (some providers like Kimi omit the space after the colon, per SSE spec) - Added scanner.Err() check to surface stream read errors (timeout, connection reset) instead of silently returning partial results Co-authored-by: Luvu182 <208665161+Luvu182@users.noreply.github.com>
This commit is contained in:
1 parent
78abdec887
commit
ab9f582de6
1 file changed
+21
-13
@@ -121,6 +121,23 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun
|
||||
continue
|
||||
}
|
||||
|
||||
// Usage chunk often has empty choices — extract usage before skipping.
|
||||
// When stream_options.include_usage is true, the final chunk contains
|
||||
// usage data but choices is typically an empty array.
|
||||
if chunk.Usage != nil {
|
||||
result.Usage = &Usage{
|
||||
PromptTokens: chunk.Usage.PromptTokens,
|
||||
CompletionTokens: chunk.Usage.CompletionTokens,
|
||||
TotalTokens: chunk.Usage.TotalTokens,
|
||||
}
|
||||
if chunk.Usage.PromptTokensDetails != nil {
|
||||
result.Usage.CacheReadTokens = chunk.Usage.PromptTokensDetails.CachedTokens
|
||||
}
|
||||
if chunk.Usage.CompletionTokensDetails != nil && chunk.Usage.CompletionTokensDetails.ReasoningTokens > 0 {
|
||||
result.Usage.ThinkingTokens = chunk.Usage.CompletionTokensDetails.ReasoningTokens
|
||||
}
|
||||
}
|
||||
|
||||
if len(chunk.Choices) == 0 {
|
||||
continue
|
||||
}
|
||||
@@ -161,20 +178,11 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun
|
||||
result.FinishReason = chunk.Choices[0].FinishReason
|
||||
}
|
||||
|
||||
if chunk.Usage != nil {
|
||||
result.Usage = &Usage{
|
||||
PromptTokens: chunk.Usage.PromptTokens,
|
||||
CompletionTokens: chunk.Usage.CompletionTokens,
|
||||
TotalTokens: chunk.Usage.TotalTokens,
|
||||
}
|
||||
if chunk.Usage.PromptTokensDetails != nil {
|
||||
result.Usage.CacheReadTokens = chunk.Usage.PromptTokensDetails.CachedTokens
|
||||
}
|
||||
if chunk.Usage.CompletionTokensDetails != nil && chunk.Usage.CompletionTokensDetails.ReasoningTokens > 0 {
|
||||
result.Usage.ThinkingTokens = chunk.Usage.CompletionTokensDetails.ReasoningTokens
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check for scanner errors (timeout, connection reset, etc.)
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil, fmt.Errorf("%s: stream read error: %w", p.name, err)
|
||||
}
|
||||
|
||||
// Check for scanner errors (timeout, connection reset, etc.)
|
||||
|
||||
Reference in new issue
Block a user