feat: comprehensive token usage tracking for V2
All checks were successful
CI / Lint (pull_request) Successful in 10m18s
CI / Root Module (pull_request) Successful in 11m4s
CI / V2 Module (pull_request) Successful in 11m5s

Add provider-specific usage details, fix streaming usage, and return
usage from all high-level APIs (Chat.Send, Generate[T], Agent.Run).

Breaking changes:
- Chat.Send/SendMessage/SendWithImages now return (string, *Usage, error)
- Generate[T]/GenerateWith[T] now return (T, *Usage, error)
- Agent.Run/RunMessages now return (string, *Usage, error)

New features:
- Usage.Details map for provider-specific token breakdowns
  (reasoning, cached, audio, thoughts tokens)
- OpenAI streaming now captures usage via StreamOptions.IncludeUsage
- Google streaming now captures UsageMetadata from final chunk
- UsageTracker.Details() for accumulated detail totals
- ModelPricing and PricingRegistry for cost computation

Closes #2

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-03-02 04:33:18 +00:00
parent 7e1705c385
commit 5b687839b2
17 changed files with 684 additions and 61 deletions

View File

@@ -58,15 +58,30 @@ func (p *Provider) Stream(ctx context.Context, req provider.Request, events chan
cl := openai.NewClient(opts...)
oaiReq := p.buildRequest(req)
oaiReq.StreamOptions = openai.ChatCompletionStreamOptionsParam{
IncludeUsage: openai.Bool(true),
}
stream := cl.Chat.Completions.NewStreaming(ctx, oaiReq)
var fullText strings.Builder
var toolCalls []provider.ToolCall
toolCallArgs := map[int]*strings.Builder{}
var usage *provider.Usage
for stream.Next() {
chunk := stream.Current()
// Capture usage from the final chunk (present when StreamOptions.IncludeUsage is true)
if chunk.Usage.TotalTokens > 0 {
usage = &provider.Usage{
InputTokens: int(chunk.Usage.PromptTokens),
OutputTokens: int(chunk.Usage.CompletionTokens),
TotalTokens: int(chunk.Usage.TotalTokens),
Details: extractUsageDetails(chunk.Usage),
}
}
for _, choice := range chunk.Choices {
// Text delta
if choice.Delta.Content != "" {
@@ -138,6 +153,7 @@ func (p *Provider) Stream(ctx context.Context, req provider.Request, events chan
Response: &provider.Response{
Text: fullText.String(),
ToolCalls: toolCalls,
Usage: usage,
},
}
@@ -363,6 +379,7 @@ func (p *Provider) convertResponse(resp *openai.ChatCompletion) provider.Respons
OutputTokens: int(resp.Usage.CompletionTokens),
TotalTokens: int(resp.Usage.TotalTokens),
}
res.Usage.Details = extractUsageDetails(resp.Usage)
}
return res
@@ -381,6 +398,27 @@ func audioFormat(contentType string) string {
}
}
// extractUsageDetails extracts provider-specific detail tokens from an OpenAI CompletionUsage.
func extractUsageDetails(usage openai.CompletionUsage) map[string]int {
details := map[string]int{}
if usage.CompletionTokensDetails.ReasoningTokens > 0 {
details[provider.UsageDetailReasoningTokens] = int(usage.CompletionTokensDetails.ReasoningTokens)
}
if usage.CompletionTokensDetails.AudioTokens > 0 {
details[provider.UsageDetailAudioOutputTokens] = int(usage.CompletionTokensDetails.AudioTokens)
}
if usage.PromptTokensDetails.CachedTokens > 0 {
details[provider.UsageDetailCachedInputTokens] = int(usage.PromptTokensDetails.CachedTokens)
}
if usage.PromptTokensDetails.AudioTokens > 0 {
details[provider.UsageDetailAudioInputTokens] = int(usage.PromptTokensDetails.AudioTokens)
}
if len(details) == 0 {
return nil
}
return details
}
// audioFormatFromURL guesses the audio format from a URL's file extension.
func audioFormatFromURL(u string) string {
ext := strings.ToLower(path.Ext(u))