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:
parent
75e00d438c
commit
f9ba6c8bce
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
2
main.go
2
main.go
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
Reference in New Issue