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:
@@ -82,6 +82,7 @@ type UsageTracker struct {
|
||||
TotalInput int64
|
||||
TotalOutput int64
|
||||
TotalRequests int64
|
||||
TotalDetails map[string]int64
|
||||
}
|
||||
|
||||
// Add records usage from a single request.
|
||||
@@ -94,6 +95,14 @@ func (ut *UsageTracker) Add(u *Usage) {
|
||||
ut.TotalInput += int64(u.InputTokens)
|
||||
ut.TotalOutput += int64(u.OutputTokens)
|
||||
ut.TotalRequests++
|
||||
if len(u.Details) > 0 {
|
||||
if ut.TotalDetails == nil {
|
||||
ut.TotalDetails = make(map[string]int64)
|
||||
}
|
||||
for k, v := range u.Details {
|
||||
ut.TotalDetails[k] += int64(v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Summary returns the accumulated totals.
|
||||
@@ -103,6 +112,20 @@ func (ut *UsageTracker) Summary() (input, output, requests int64) {
|
||||
return ut.TotalInput, ut.TotalOutput, ut.TotalRequests
|
||||
}
|
||||
|
||||
// Details returns a copy of the accumulated detail totals.
|
||||
func (ut *UsageTracker) Details() map[string]int64 {
|
||||
ut.mu.Lock()
|
||||
defer ut.mu.Unlock()
|
||||
if ut.TotalDetails == nil {
|
||||
return nil
|
||||
}
|
||||
cp := make(map[string]int64, len(ut.TotalDetails))
|
||||
for k, v := range ut.TotalDetails {
|
||||
cp[k] = v
|
||||
}
|
||||
return cp
|
||||
}
|
||||
|
||||
// WithUsageTracking returns middleware that accumulates token usage across calls.
|
||||
func WithUsageTracking(tracker *UsageTracker) Middleware {
|
||||
return func(next CompletionFunc) CompletionFunc {
|
||||
|
||||
Reference in New Issue
Block a user