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

128
v2/pricing_test.go Normal file
View File

@@ -0,0 +1,128 @@
package llm
import (
"math"
"testing"
)
func TestModelPricing_Cost(t *testing.T) {
pricing := ModelPricing{
InputPricePerToken: 0.000003, // $3/MTok
OutputPricePerToken: 0.000015, // $15/MTok
}
usage := &Usage{
InputTokens: 1000,
OutputTokens: 500,
TotalTokens: 1500,
}
cost := pricing.Cost(usage)
expected := 1000*0.000003 + 500*0.000015
if math.Abs(cost-expected) > 1e-10 {
t.Errorf("expected cost %f, got %f", expected, cost)
}
}
func TestModelPricing_Cost_WithCachedTokens(t *testing.T) {
pricing := ModelPricing{
InputPricePerToken: 0.000003, // $3/MTok
OutputPricePerToken: 0.000015, // $15/MTok
CachedInputPricePerToken: 0.0000015, // $1.50/MTok (50% discount)
}
usage := &Usage{
InputTokens: 1000,
OutputTokens: 500,
TotalTokens: 1500,
Details: map[string]int{
UsageDetailCachedInputTokens: 400,
},
}
cost := pricing.Cost(usage)
// 600 regular input tokens + 400 cached tokens + 500 output tokens
expected := 600*0.000003 + 400*0.0000015 + 500*0.000015
if math.Abs(cost-expected) > 1e-10 {
t.Errorf("expected cost %f, got %f", expected, cost)
}
}
func TestModelPricing_Cost_NilUsage(t *testing.T) {
pricing := ModelPricing{
InputPricePerToken: 0.000003,
OutputPricePerToken: 0.000015,
}
cost := pricing.Cost(nil)
if cost != 0 {
t.Errorf("expected 0 for nil usage, got %f", cost)
}
}
func TestModelPricing_Cost_NoCachedPrice(t *testing.T) {
// When CachedInputPricePerToken is 0, all input tokens use InputPricePerToken
pricing := ModelPricing{
InputPricePerToken: 0.000003,
OutputPricePerToken: 0.000015,
}
usage := &Usage{
InputTokens: 1000,
OutputTokens: 500,
TotalTokens: 1500,
Details: map[string]int{
UsageDetailCachedInputTokens: 400,
},
}
cost := pricing.Cost(usage)
expected := 1000*0.000003 + 500*0.000015
if math.Abs(cost-expected) > 1e-10 {
t.Errorf("expected cost %f, got %f", expected, cost)
}
}
func TestPricingRegistry(t *testing.T) {
registry := NewPricingRegistry()
registry.Set("gpt-4o", ModelPricing{
InputPricePerToken: 0.0000025,
OutputPricePerToken: 0.00001,
})
if !registry.Has("gpt-4o") {
t.Error("expected Has('gpt-4o') to be true")
}
if registry.Has("gpt-3.5-turbo") {
t.Error("expected Has('gpt-3.5-turbo') to be false")
}
usage := &Usage{InputTokens: 1000, OutputTokens: 200, TotalTokens: 1200}
cost := registry.Cost("gpt-4o", usage)
expected := 1000*0.0000025 + 200*0.00001
if math.Abs(cost-expected) > 1e-10 {
t.Errorf("expected cost %f, got %f", expected, cost)
}
// Unknown model returns 0
cost = registry.Cost("unknown-model", usage)
if cost != 0 {
t.Errorf("expected 0 for unknown model, got %f", cost)
}
}
func TestPricingRegistry_Override(t *testing.T) {
registry := NewPricingRegistry()
registry.Set("model-a", ModelPricing{InputPricePerToken: 0.001, OutputPricePerToken: 0.002})
registry.Set("model-a", ModelPricing{InputPricePerToken: 0.003, OutputPricePerToken: 0.004})
usage := &Usage{InputTokens: 100, OutputTokens: 50, TotalTokens: 150}
cost := registry.Cost("model-a", usage)
expected := 100*0.003 + 50*0.004
if math.Abs(cost-expected) > 1e-10 {
t.Errorf("expected overridden cost %f, got %f", expected, cost)
}
}