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:
128
v2/pricing_test.go
Normal file
128
v2/pricing_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user