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:
@@ -25,7 +25,7 @@ func TestGenerate(t *testing.T) {
|
||||
})
|
||||
model := newMockModel(mp)
|
||||
|
||||
result, err := Generate[testPerson](context.Background(), model, "Tell me about Alice")
|
||||
result, _, err := Generate[testPerson](context.Background(), model, "Tell me about Alice")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
@@ -63,7 +63,7 @@ func TestGenerateWith(t *testing.T) {
|
||||
UserMessage("Tell me about Bob"),
|
||||
}
|
||||
|
||||
result, err := GenerateWith[testPerson](context.Background(), model, messages)
|
||||
result, _, err := GenerateWith[testPerson](context.Background(), model, messages)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
@@ -90,7 +90,7 @@ func TestGenerate_NoToolCall(t *testing.T) {
|
||||
})
|
||||
model := newMockModel(mp)
|
||||
|
||||
_, err := Generate[testPerson](context.Background(), model, "Tell me about someone")
|
||||
_, _, err := Generate[testPerson](context.Background(), model, "Tell me about someone")
|
||||
if err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
}
|
||||
@@ -111,7 +111,7 @@ func TestGenerate_InvalidJSON(t *testing.T) {
|
||||
})
|
||||
model := newMockModel(mp)
|
||||
|
||||
_, err := Generate[testPerson](context.Background(), model, "Tell me about someone")
|
||||
_, _, err := Generate[testPerson](context.Background(), model, "Tell me about someone")
|
||||
if err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
}
|
||||
@@ -143,7 +143,7 @@ func TestGenerate_NestedStruct(t *testing.T) {
|
||||
})
|
||||
model := newMockModel(mp)
|
||||
|
||||
result, err := Generate[testPersonWithAddress](context.Background(), model, "Tell me about Carol")
|
||||
result, _, err := Generate[testPersonWithAddress](context.Background(), model, "Tell me about Carol")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
@@ -170,7 +170,7 @@ func TestGenerate_WithOptions(t *testing.T) {
|
||||
})
|
||||
model := newMockModel(mp)
|
||||
|
||||
_, err := Generate[testPerson](context.Background(), model, "Tell me about Dave",
|
||||
_, _, err := Generate[testPerson](context.Background(), model, "Tell me about Dave",
|
||||
WithTemperature(0.5),
|
||||
WithMaxTokens(200),
|
||||
)
|
||||
@@ -207,7 +207,7 @@ func TestGenerate_WithMiddleware(t *testing.T) {
|
||||
})
|
||||
model := newMockModel(mp).WithMiddleware(mw)
|
||||
|
||||
result, err := Generate[testPerson](context.Background(), model, "Tell me about Eve")
|
||||
result, _, err := Generate[testPerson](context.Background(), model, "Tell me about Eve")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
@@ -231,7 +231,7 @@ func TestGenerate_WrongToolName(t *testing.T) {
|
||||
})
|
||||
model := newMockModel(mp)
|
||||
|
||||
_, err := Generate[testPerson](context.Background(), model, "Tell me about Frank")
|
||||
_, _, err := Generate[testPerson](context.Background(), model, "Tell me about Frank")
|
||||
if err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
}
|
||||
@@ -239,3 +239,44 @@ func TestGenerate_WrongToolName(t *testing.T) {
|
||||
t.Errorf("expected ErrNoStructuredOutput, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerate_ReturnsUsage(t *testing.T) {
|
||||
mp := newMockProvider(provider.Response{
|
||||
ToolCalls: []provider.ToolCall{
|
||||
{
|
||||
ID: "call_1",
|
||||
Name: "structured_output",
|
||||
Arguments: `{"name":"Grace","age":22}`,
|
||||
},
|
||||
},
|
||||
Usage: &provider.Usage{
|
||||
InputTokens: 50,
|
||||
OutputTokens: 20,
|
||||
TotalTokens: 70,
|
||||
Details: map[string]int{
|
||||
"reasoning_tokens": 5,
|
||||
},
|
||||
},
|
||||
})
|
||||
model := newMockModel(mp)
|
||||
|
||||
result, usage, err := Generate[testPerson](context.Background(), model, "Tell me about Grace")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if result.Name != "Grace" {
|
||||
t.Errorf("expected name 'Grace', got %q", result.Name)
|
||||
}
|
||||
if usage == nil {
|
||||
t.Fatal("expected usage, got nil")
|
||||
}
|
||||
if usage.InputTokens != 50 {
|
||||
t.Errorf("expected input 50, got %d", usage.InputTokens)
|
||||
}
|
||||
if usage.OutputTokens != 20 {
|
||||
t.Errorf("expected output 20, got %d", usage.OutputTokens)
|
||||
}
|
||||
if usage.Details["reasoning_tokens"] != 5 {
|
||||
t.Errorf("expected reasoning_tokens=5, got %d", usage.Details["reasoning_tokens"])
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user