diff --git a/internal/llm/client.go b/internal/llm/client.go index f034f1a..a2f3b1b 100644 --- a/internal/llm/client.go +++ b/internal/llm/client.go @@ -93,13 +93,6 @@ func extractBlockText(block ContentBlock) string { return sb.String() } -// Usage holds token usage statistics from a response. -type Usage struct { - PromptTokens int64 `json:"prompt_tokens"` - CompletionTokens int64 `json:"completion_tokens"` - CacheCreationInputToken int64 `json:"cache_creation_input_tokens,omitempty"` - CacheReadInputTokens int64 `json:"cache_read_input_tokens,omitempty"` -} // Choice holds a single choice from the response. type Choice struct { @@ -133,7 +126,6 @@ type ChatResponse struct { ID string `json:"-"` Model string `json:"-"` Choices []Choice `json:"-"` - Usage *Usage `json:"-"` Headers http.Header `json:"-"` // Raw response headers (may contain session IDs, etc.) } @@ -478,7 +470,6 @@ func (c *Client) doRequest(model string, req ChatRequest) (*ChatResponse, error) ID string `json:"id"` Model string `json:"model"` Choices []Choice `json:"choices"` - Usage *Usage `json:"usage,omitempty"` } if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { return nil, fmt.Errorf("decode response: %w", err) @@ -488,7 +479,6 @@ func (c *Client) doRequest(model string, req ChatRequest) (*ChatResponse, error) ID: apiResp.ID, Model: apiResp.Model, Choices: apiResp.Choices, - Usage: apiResp.Usage, Headers: resp.Header, }, nil } diff --git a/internal/session/history.go b/internal/session/history.go index e1be385..1bdaf29 100644 --- a/internal/session/history.go +++ b/internal/session/history.go @@ -52,12 +52,21 @@ type TaskRecord struct { Error string } +// TokenUsage holds estimated token usage for a single LLM request/response cycle. +// Calculated locally using tiktoken so it works with any model provider. +type TokenUsage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + CacheReadTokens int `json:"cache_read_tokens,omitempty"` + CacheWriteTokens int `json:"cache_write_tokens,omitempty"` +} + // ResponseRecord holds the parsed LLM response. type ResponseRecord struct { Content string ToolCalls []llm.ToolCall Model string - Usage *llm.Usage + Usage *TokenUsage } // ToolResultRecord records the result of a tool call executed after the LLM response. @@ -200,6 +209,7 @@ func (fs *FileSession) AppendTaskRecord(taskType TaskType, messages []llm.Messag } // SetResponse records the LLM response in the most recent TaskRecord of the given type. +// It computes estimated token usage locally using tiktoken, independent of any API format. func (tr *TaskRecord) SetResponse(resp *llm.ChatResponse, duration time.Duration) { if resp == nil || len(resp.Choices) == 0 { tr.Error = "empty response" @@ -211,11 +221,24 @@ func (tr *TaskRecord) SetResponse(resp *llm.ChatResponse, duration time.Duration if choice.Message.Content != nil { content = *choice.Message.Content } + + // Estimate tokens using tiktoken — works with any model provider. + var promptTokens int + for _, m := range tr.RequestMessages { + promptTokens += llm.CountTokens(m.ExtractText()) + } + completionTokens := llm.CountTokens(content) + + usage := &TokenUsage{ + PromptTokens: promptTokens, + CompletionTokens: completionTokens, + } + tr.Response = &ResponseRecord{ Content: content, ToolCalls: choice.Message.ToolCalls, Model: resp.Model, - Usage: resp.Usage, + Usage: usage, } tr.Duration = duration }