steve/go-llm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: comprehensive token usage tracking for V2
All checks were successful
CI / Lint (pull_request) Successful in 10m18s
Details
CI / Root Module (pull_request) Successful in 11m4s
Details
CI / V2 Module (pull_request) Successful in 11m5s
Details

Browse Source 
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:
Steve Dudenhoeffer
2026-03-02 04:33:18 +00:00
parent 7e1705c385
commit 5b687839b2
17 changed files with 684 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

132
v2/chat_test.go
View File

@@ -14,7 +14,7 @@ func TestChat_Send(t *testing.T) {
model := newMockModel(mp)
chat := NewChat(model)
text, err := chat.Send(context.Background(), "Hi")
text, _, err := chat.Send(context.Background(), "Hi")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
@@ -28,7 +28,7 @@ func TestChat_SendMessage(t *testing.T) {
model := newMockModel(mp)
chat := NewChat(model)
_, err := chat.SendMessage(context.Background(), UserMessage("msg1"))
_, _, err := chat.SendMessage(context.Background(), UserMessage("msg1"))
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
@@ -79,7 +79,7 @@ func TestChat_SetSystem(t *testing.T) {
}
// System message stays first even after adding other messages
_, _ = chat.Send(context.Background(), "Hi")
_, _, _ = chat.Send(context.Background(), "Hi")
chat.SetSystem("New system")
msgs = chat.Messages()
if msgs[0].Role != RoleSystem {
@@ -113,7 +113,7 @@ func TestChat_ToolCallLoop(t *testing.T) {
})
chat.SetTools(NewToolBox(tool))
text, err := chat.Send(context.Background(), "test")
text, _, err := chat.Send(context.Background(), "test")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
@@ -158,7 +158,7 @@ func TestChat_ToolCallLoop_NoTools(t *testing.T) {
model := newMockModel(mp)
chat := NewChat(model)
_, err := chat.Send(context.Background(), "test")
_, _, err := chat.Send(context.Background(), "test")
if !errors.Is(err, ErrNoToolsConfigured) {
t.Errorf("expected ErrNoToolsConfigured, got %v", err)
}
@@ -248,7 +248,7 @@ func TestChat_Messages(t *testing.T) {
model := newMockModel(mp)
chat := NewChat(model)
_, _ = chat.Send(context.Background(), "test")
_, _, _ = chat.Send(context.Background(), "test")
msgs := chat.Messages()
// Verify it's a copy — modifying returned slice shouldn't affect chat
@@ -265,7 +265,7 @@ func TestChat_Reset(t *testing.T) {
model := newMockModel(mp)
chat := NewChat(model)
_, _ = chat.Send(context.Background(), "test")
_, _, _ = chat.Send(context.Background(), "test")
if len(chat.Messages()) == 0 {
t.Fatal("expected messages before reset")
}
@@ -281,7 +281,7 @@ func TestChat_Fork(t *testing.T) {
model := newMockModel(mp)
chat := NewChat(model)
_, _ = chat.Send(context.Background(), "msg1")
_, _, _ = chat.Send(context.Background(), "msg1")
fork := chat.Fork()
@@ -291,14 +291,14 @@ func TestChat_Fork(t *testing.T) {
}
// Adding to fork should not affect original
_, _ = fork.Send(context.Background(), "msg2")
_, _, _ = fork.Send(context.Background(), "msg2")
if len(fork.Messages()) == len(chat.Messages()) {
t.Error("fork messages should be independent of original")
}
// Adding to original should not affect fork
originalLen := len(chat.Messages())
_, _ = chat.Send(context.Background(), "msg3")
_, _, _ = chat.Send(context.Background(), "msg3")
if len(chat.Messages()) == originalLen {
t.Error("original should have more messages after send")
}
@@ -310,7 +310,7 @@ func TestChat_SendWithImages(t *testing.T) {
chat := NewChat(model)
img := Image{URL: "https://example.com/image.png"}
text, err := chat.SendWithImages(context.Background(), "What's in this image?", img)
text, _, err := chat.SendWithImages(context.Background(), "What's in this image?", img)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
@@ -355,7 +355,7 @@ func TestChat_MultipleToolCallRounds(t *testing.T) {
})
chat.SetTools(NewToolBox(tool))
text, err := chat.Send(context.Background(), "count three times")
text, _, err := chat.Send(context.Background(), "count three times")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
@@ -378,7 +378,7 @@ func TestChat_SendError(t *testing.T) {
model := newMockModel(mp)
chat := NewChat(model)
_, err := chat.Send(context.Background(), "test")
_, _, err := chat.Send(context.Background(), "test")
if err == nil {
t.Fatal("expected error, got nil")
}
@@ -392,7 +392,7 @@ func TestChat_WithRequestOptions(t *testing.T) {
model := newMockModel(mp)
chat := NewChat(model, WithTemperature(0.5), WithMaxTokens(200))
_, err := chat.Send(context.Background(), "test")
_, _, err := chat.Send(context.Background(), "test")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
@@ -405,3 +405,107 @@ func TestChat_WithRequestOptions(t *testing.T) {
t.Errorf("expected maxTokens 200, got %v", req.MaxTokens)
}
}
func TestChat_Send_UsageAccumulation(t *testing.T) {
var callCount int32
mp := newMockProviderFunc(func(ctx context.Context, req provider.Request) (provider.Response, error) {
n := atomic.AddInt32(&callCount, 1)
if n == 1 {
return provider.Response{
ToolCalls: []provider.ToolCall{
{ID: "tc1", Name: "greet", Arguments: "{}"},
},
Usage: &provider.Usage{InputTokens: 10, OutputTokens: 5, TotalTokens: 15},
}, nil
}
return provider.Response{
Text: "done",
Usage: &provider.Usage{InputTokens: 20, OutputTokens: 8, TotalTokens: 28},
}, nil
})
model := newMockModel(mp)
chat := NewChat(model)
tool := DefineSimple("greet", "Says hello", func(ctx context.Context) (string, error) {
return "hello!", nil
})
chat.SetTools(NewToolBox(tool))
text, usage, err := chat.Send(context.Background(), "test")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if text != "done" {
t.Errorf("expected 'done', got %q", text)
}
if usage == nil {
t.Fatal("expected usage, got nil")
}
if usage.InputTokens != 30 {
t.Errorf("expected accumulated input 30, got %d", usage.InputTokens)
}
if usage.OutputTokens != 13 {
t.Errorf("expected accumulated output 13, got %d", usage.OutputTokens)
}
if usage.TotalTokens != 43 {
t.Errorf("expected accumulated total 43, got %d", usage.TotalTokens)
}
}
func TestChat_Send_UsageWithDetails(t *testing.T) {
var callCount int32
mp := newMockProviderFunc(func(ctx context.Context, req provider.Request) (provider.Response, error) {
n := atomic.AddInt32(&callCount, 1)
if n == 1 {
return provider.Response{
ToolCalls: []provider.ToolCall{
{ID: "tc1", Name: "greet", Arguments: "{}"},
},
Usage: &provider.Usage{
InputTokens: 10,
OutputTokens: 5,
TotalTokens: 15,
Details: map[string]int{
"cached_input_tokens": 3,
},
},
}, nil
}
return provider.Response{
Text: "done",
Usage: &provider.Usage{
InputTokens: 20,
OutputTokens: 8,
TotalTokens: 28,
Details: map[string]int{
"cached_input_tokens": 7,
"reasoning_tokens": 2,
},
},
}, nil
})
model := newMockModel(mp)
chat := NewChat(model)
tool := DefineSimple("greet", "Says hello", func(ctx context.Context) (string, error) {
return "hello!", nil
})
chat.SetTools(NewToolBox(tool))
_, usage, err := chat.Send(context.Background(), "test")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if usage == nil {
t.Fatal("expected usage, got nil")
}
if usage.Details == nil {
t.Fatal("expected usage details, got nil")
}
if usage.Details["cached_input_tokens"] != 10 {
t.Errorf("expected cached_input_tokens=10, got %d", usage.Details["cached_input_tokens"])
}
if usage.Details["reasoning_tokens"] != 2 {
t.Errorf("expected reasoning_tokens=2, got %d", usage.Details["reasoning_tokens"])
}
}