diff --git a/main.go b/main.go index 0e5c8e1..ee77df1 100644 --- a/main.go +++ b/main.go @@ -764,6 +764,7 @@ type Usage struct { CachedTokens int `json:"cached_tokens"` } `json:"prompt_tokens_details"` CachedTokens int `json:"cached_tokens"` + Model string `json:"-"` } func (u Usage) Cached() int { @@ -799,12 +800,19 @@ func formatUsage(u Usage, cw int) string { if cached > 0 { uncached := u.PromptTokens - cached if uncached < 0 { uncached = 0 } - return fmt.Sprintf("[tokens: %d prompt (%d cached, %d uncached) + %d completion | context: %d/%d (%.1f%%)]", u.PromptTokens, cached, uncached, u.CompletionTokens, u.PromptTokens, cw, pct) + if u.Model != "" { + return fmt.Sprintf("[%s: %d prompt (%d cached, %d uncached) + %d completion | context: %d/%d (%.1f%%)]", u.Model, u.PromptTokens, cached, uncached, u.CompletionTokens, u.PromptTokens, cw, pct) + } + return fmt.Sprintf("[%d prompt (%d cached, %d uncached) + %d completion | context: %d/%d (%.1f%%)]", u.PromptTokens, cached, uncached, u.CompletionTokens, u.PromptTokens, cw, pct) } - return fmt.Sprintf("[tokens: %d prompt + %d completion | context: %d/%d (%.1f%%)]", u.PromptTokens, u.CompletionTokens, u.PromptTokens, cw, pct) + if u.Model != "" { + return fmt.Sprintf("[%s: %d prompt + %d completion | context: %d/%d (%.1f%%)]", u.Model, u.PromptTokens, u.CompletionTokens, u.PromptTokens, cw, pct) + } + return fmt.Sprintf("[%d prompt + %d completion | context: %d/%d (%.1f%%)]", u.PromptTokens, u.CompletionTokens, u.PromptTokens, cw, pct) } type streamDelta struct { + Model string `json:"model"` Choices []struct { Delta struct { ReasoningContent string `json:"reasoning_content"` @@ -946,6 +954,7 @@ func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any) if COL { fmt.Print("\r\033[K") } if !cfg.Stream { var cr struct { + Model string `json:"model"` Choices []struct { Message struct { Message @@ -964,6 +973,7 @@ func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any) if m.ReasoningContent == "" { m.ReasoningContent = cr.Choices[0].Message.Reasoning } if m.ReasoningContent == "" { m.ReasoningContent = cr.Choices[0].Message.Thought } u := cr.Usage + if cr.Model != "" { u.Model = cr.Model } if u.PromptTokens == 0 { u.PromptTokens = estTokens(msgs) u.CompletionTokens = estTokens([]Message{m}) @@ -1010,6 +1020,7 @@ func parseStream(ctx context.Context, r io.Reader) (Message, Usage, error) { if d.Usage != nil && (d.Usage.PromptTokens > 0 || d.Usage.TotalTokens > 0) { lastUsage = *d.Usage } + if d.Model != "" { lastUsage.Model = d.Model } if len(d.Choices) == 0 { continue } dl := d.Choices[0].Delta rc := dl.ReasoningContent @@ -1209,6 +1220,7 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message) ([]Message, Usage, error) turnUsage.CompletionTokens += u.CompletionTokens turnUsage.TotalTokens += u.TotalTokens if u.Cached() > 0 { turnUsage.CachedTokens = u.Cached() } + if u.Model != "" { turnUsage.Model = u.Model } for j := range m.ToolCalls { tc := &m.ToolCalls[j] tc.Function.Arguments = filterText(tc.Function.Arguments) diff --git a/main_test.go b/main_test.go index 017422b..02cf06a 100644 --- a/main_test.go +++ b/main_test.go @@ -1865,14 +1865,14 @@ func TestFormatUsage(t *testing.T) { // With cached tokens u1 := Usage{PromptTokens: 1000, CompletionTokens: 200, TotalTokens: 1200, CachedTokens: 800} s1 := formatUsage(u1, 200000) - if s1 != "[tokens: 1000 prompt (800 cached, 200 uncached) + 200 completion | context: 1000/200000 (0.5%)]" { + if s1 != "[1000 prompt (800 cached, 200 uncached) + 200 completion | context: 1000/200000 (0.5%)]" { t.Errorf("formatUsage u1 = %q", s1) } // Without cached tokens u2 := Usage{PromptTokens: 120000, CompletionTokens: 500, TotalTokens: 120500} s2 := formatUsage(u2, 200000) - if s2 != "[tokens: 120000 prompt + 500 completion | context: 120000/200000 (60.0%)]" { + if s2 != "[120000 prompt + 500 completion | context: 120000/200000 (60.0%)]" { t.Errorf("formatUsage u2 = %q", s2) } }