feat: comprehensive token usage tracking for V2
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:
@@ -59,12 +59,32 @@ func (p *Provider) Stream(ctx context.Context, req provider.Request, events chan
|
||||
|
||||
var fullText strings.Builder
|
||||
var toolCalls []provider.ToolCall
|
||||
var usage *provider.Usage
|
||||
|
||||
for resp, err := range cl.Models.GenerateContentStream(ctx, req.Model, contents, cfg) {
|
||||
if err != nil {
|
||||
return fmt.Errorf("google stream error: %w", err)
|
||||
}
|
||||
|
||||
// Track usage from the last chunk (final chunk has cumulative counts)
|
||||
if resp.UsageMetadata != nil {
|
||||
usage = &provider.Usage{
|
||||
InputTokens: int(resp.UsageMetadata.PromptTokenCount),
|
||||
OutputTokens: int(resp.UsageMetadata.CandidatesTokenCount),
|
||||
TotalTokens: int(resp.UsageMetadata.TotalTokenCount),
|
||||
}
|
||||
details := map[string]int{}
|
||||
if resp.UsageMetadata.CachedContentTokenCount > 0 {
|
||||
details[provider.UsageDetailCachedInputTokens] = int(resp.UsageMetadata.CachedContentTokenCount)
|
||||
}
|
||||
if resp.UsageMetadata.ThoughtsTokenCount > 0 {
|
||||
details[provider.UsageDetailThoughtsTokens] = int(resp.UsageMetadata.ThoughtsTokenCount)
|
||||
}
|
||||
if len(details) > 0 {
|
||||
usage.Details = details
|
||||
}
|
||||
}
|
||||
|
||||
for _, c := range resp.Candidates {
|
||||
if c.Content == nil {
|
||||
continue
|
||||
@@ -105,6 +125,7 @@ func (p *Provider) Stream(ctx context.Context, req provider.Request, events chan
|
||||
Response: &provider.Response{
|
||||
Text: fullText.String(),
|
||||
ToolCalls: toolCalls,
|
||||
Usage: usage,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -284,6 +305,16 @@ func (p *Provider) convertResponse(resp *genai.GenerateContentResponse) (provide
|
||||
OutputTokens: int(resp.UsageMetadata.CandidatesTokenCount),
|
||||
TotalTokens: int(resp.UsageMetadata.TotalTokenCount),
|
||||
}
|
||||
details := map[string]int{}
|
||||
if resp.UsageMetadata.CachedContentTokenCount > 0 {
|
||||
details[provider.UsageDetailCachedInputTokens] = int(resp.UsageMetadata.CachedContentTokenCount)
|
||||
}
|
||||
if resp.UsageMetadata.ThoughtsTokenCount > 0 {
|
||||
details[provider.UsageDetailThoughtsTokens] = int(resp.UsageMetadata.ThoughtsTokenCount)
|
||||
}
|
||||
if len(details) > 0 {
|
||||
res.Usage.Details = details
|
||||
}
|
||||
}
|
||||
|
||||
return res, nil
|
||||
|
||||
Reference in New Issue
Block a user