Track and display token usage per request

Add Usage{InputTokens, OutputTokens} to backend.Response. Both
OllamaBackend (prompt_eval_count/eval_count) and OpenAIBackend
(prompt_tokens/completion_tokens) populate it. The loop emits a
'usage' OutputMsg after each Chat call; the UI displays it as
[↑N ↓N tokens].

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Levi Neely 2026-04-01 00:31:03 +02:00
parent 75e00d438c
commit f9ba6c8bce
5 changed files with 37 additions and 7 deletions

View File

@ -13,10 +13,10 @@ import (
type ToolExecutor func(name string, args json.RawMessage) (string, error)
// OutputMsg is emitted by the loop for each visible event: assistant replies,
// tool calls, tool results, and errors.
// tool calls, tool results, errors, and token usage.
type OutputMsg struct {
Role string // "assistant" | "tool" | "error"
Name string // tool name, for Role=="tool"
Role string // "assistant" | "call" | "tool" | "error" | "usage"
Name string // tool name, for Role=="call"/"tool"
Content string
}
@ -97,6 +97,11 @@ func (l *Loop) Run(ctx context.Context, state State) error {
l.emit(OutputMsg{Role: "assistant", Content: resp.Message.Content})
}
l.emit(OutputMsg{
Role: "usage",
Content: fmt.Sprintf("↑%d ↓%d tokens", resp.Usage.InputTokens, resp.Usage.OutputTokens),
})
// 4. Update: append assistant message and tool results to state.
if err := state.Update(resp.Message, results); err != nil {
return fmt.Errorf("step %d update: %w", step, err)

View File

@ -30,10 +30,17 @@ type ToolCall struct {
Arguments json.RawMessage `json:"arguments"` // always a JSON object
}
// Usage holds token counts for a single Chat call.
type Usage struct {
InputTokens int
OutputTokens int
}
// Response is the model's reply for one Chat call.
type Response struct {
Message Message
StopReason string // "stop" | "tool_calls" | "length" | ...
Usage Usage
}
// Backend is the interface all LLM providers must implement.

View File

@ -58,8 +58,10 @@ type ollamaChatRequest struct {
}
type ollamaChatResponse struct {
Message ollamaMessage `json:"message"`
Done bool `json:"done"`
Message ollamaMessage `json:"message"`
Done bool `json:"done"`
PromptEvalCount int `json:"prompt_eval_count"`
EvalCount int `json:"eval_count"`
}
// -- implementation --
@ -134,5 +136,9 @@ func (b *OllamaBackend) Chat(ctx context.Context, model string, messages []Messa
stopReason = "tool_calls"
}
return &Response{Message: msg, StopReason: stopReason}, nil
return &Response{
Message: msg,
StopReason: stopReason,
Usage: Usage{InputTokens: wire.PromptEvalCount, OutputTokens: wire.EvalCount},
}, nil
}

View File

@ -62,8 +62,14 @@ type openAIChatRequest struct {
Stream bool `json:"stream"`
}
type openAIUsage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
}
type openAIChatResponse struct {
Choices []openAIChoice `json:"choices"`
Usage openAIUsage `json:"usage"`
}
type openAIChoice struct {
@ -158,5 +164,9 @@ func (b *OpenAIBackend) Chat(ctx context.Context, model string, messages []Messa
})
}
return &Response{Message: msg, StopReason: choice.FinishReason}, nil
return &Response{
Message: msg,
StopReason: choice.FinishReason,
Usage: Usage{InputTokens: wire.Usage.PromptTokens, OutputTokens: wire.Usage.CompletionTokens},
}, nil
}

View File

@ -278,6 +278,8 @@ func (m model) runLoop(input string) tea.Cmd {
lines = append(lines, fmt.Sprintf("→ %s(%s)", msg.Name, msg.Content))
case "tool":
lines = append(lines, fmt.Sprintf(" = %s", msg.Content))
case "usage":
lines = append(lines, "["+msg.Content+"]")
case "error":
lines = append(lines, "Error: "+msg.Content)
}