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:
@@ -29,15 +29,6 @@ func (m *mockProvider) Stream(ctx context.Context, req provider.Request, events
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockProvider) lastRequest() provider.Request {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if len(m.requests) == 0 {
|
||||
return provider.Request{}
|
||||
}
|
||||
return m.requests[len(m.requests)-1]
|
||||
}
|
||||
|
||||
func newMockModel(fn func(ctx context.Context, req provider.Request) (provider.Response, error)) *llm.Model {
|
||||
mp := &mockProvider{completeFunc: fn}
|
||||
return llm.NewClient(mp).Model("mock-model")
|
||||
@@ -53,7 +44,7 @@ func TestAgent_Run(t *testing.T) {
|
||||
model := newSimpleMockModel("Hello from agent!")
|
||||
a := New(model, "You are a helpful assistant.")
|
||||
|
||||
result, err := a.Run(context.Background(), "Say hello")
|
||||
result, _, err := a.Run(context.Background(), "Say hello")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
@@ -83,7 +74,7 @@ func TestAgent_Run_WithTools(t *testing.T) {
|
||||
})
|
||||
|
||||
a := New(model, "You are helpful.", WithTools(llm.NewToolBox(tool)))
|
||||
result, err := a.Run(context.Background(), "Use the greet tool")
|
||||
result, _, err := a.Run(context.Background(), "Use the greet tool")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
@@ -147,7 +138,7 @@ func TestAgent_AsTool_ParentChild(t *testing.T) {
|
||||
)),
|
||||
)
|
||||
|
||||
result, err := parent.Run(context.Background(), "Tell me about Go generics")
|
||||
result, _, err := parent.Run(context.Background(), "Tell me about Go generics")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
@@ -169,7 +160,7 @@ func TestAgent_RunMessages(t *testing.T) {
|
||||
llm.UserMessage("Follow up"),
|
||||
}
|
||||
|
||||
result, err := a.RunMessages(context.Background(), messages)
|
||||
result, _, err := a.RunMessages(context.Background(), messages)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
@@ -187,7 +178,7 @@ func TestAgent_ContextCancellation(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel() // Cancel immediately
|
||||
|
||||
_, err := a.Run(ctx, "This should fail")
|
||||
_, _, err := a.Run(ctx, "This should fail")
|
||||
if err == nil {
|
||||
t.Fatal("expected error from cancelled context")
|
||||
}
|
||||
@@ -204,7 +195,7 @@ func TestAgent_WithRequestOptions(t *testing.T) {
|
||||
WithRequestOptions(llm.WithTemperature(0.3), llm.WithMaxTokens(100)),
|
||||
)
|
||||
|
||||
_, err := a.Run(context.Background(), "test")
|
||||
_, _, err := a.Run(context.Background(), "test")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
@@ -224,7 +215,7 @@ func TestAgent_Run_Error(t *testing.T) {
|
||||
})
|
||||
a := New(model, "You are helpful.")
|
||||
|
||||
_, err := a.Run(context.Background(), "test")
|
||||
_, _, err := a.Run(context.Background(), "test")
|
||||
if err == nil {
|
||||
t.Fatal("expected error, got nil")
|
||||
}
|
||||
@@ -234,7 +225,7 @@ func TestAgent_EmptySystem(t *testing.T) {
|
||||
model := newSimpleMockModel("no system prompt")
|
||||
a := New(model, "") // Empty system prompt
|
||||
|
||||
result, err := a.Run(context.Background(), "test")
|
||||
result, _, err := a.Run(context.Background(), "test")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
@@ -242,3 +233,34 @@ func TestAgent_EmptySystem(t *testing.T) {
|
||||
t.Errorf("unexpected result: %q", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgent_Run_ReturnsUsage(t *testing.T) {
|
||||
model := newMockModel(func(ctx context.Context, req provider.Request) (provider.Response, error) {
|
||||
return provider.Response{
|
||||
Text: "result",
|
||||
Usage: &provider.Usage{
|
||||
InputTokens: 100,
|
||||
OutputTokens: 50,
|
||||
TotalTokens: 150,
|
||||
},
|
||||
}, nil
|
||||
})
|
||||
|
||||
a := New(model, "You are helpful.")
|
||||
result, usage, err := a.Run(context.Background(), "test")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if result != "result" {
|
||||
t.Errorf("expected 'result', got %q", result)
|
||||
}
|
||||
if usage == nil {
|
||||
t.Fatal("expected usage, got nil")
|
||||
}
|
||||
if usage.InputTokens != 100 {
|
||||
t.Errorf("expected input 100, got %d", usage.InputTokens)
|
||||
}
|
||||
if usage.OutputTokens != 50 {
|
||||
t.Errorf("expected output 50, got %d", usage.OutputTokens)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user