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

View File

@@ -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"])
}
}