Steve Dudenhoeffer
5b687839b2
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>
129 lines
3.2 KiB
Go
129 lines
3.2 KiB
Go
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)
|
|
}
|
|
}
|