collect usage

llm/claude/v2/chat.go

@@ -111,6 +111,9 @@ func (c *Claude) Chat(ctx context.Context, chatReq llm.ChatRequest) (*llm.ChatRe
 		return nil, err
 	}
 
+	if c.tokenUsage.Model == "" && resp.Model != "" {
+		c.tokenUsage.Model = string(resp.Model)
+	}
 	c.tokenUsage.PromptTokens += resp.Usage.InputTokens
 	c.tokenUsage.CompletionTokens += resp.Usage.OutputTokens
 	c.tokenUsage.TotalTokens += resp.Usage.InputTokens + resp.Usage.OutputTokens

llm/google/v2/chat.go

@@ -110,6 +110,9 @@ func (g *Gemini) Chat(ctx context.Context, chatReq llm.ChatRequest) (*llm.ChatRe
 		return nil, err
 	}
 
+	if g.tokenUsage.Model == "" && response.ModelVersion != "" {
+		g.tokenUsage.Model = response.ModelVersion
+	}
 	if response.UsageMetadata != nil {
 		g.tokenUsage.PromptTokens += int(response.UsageMetadata.PromptTokenCount)
 		g.tokenUsage.CompletionTokens += int(response.UsageMetadata.CandidatesTokenCount)

llm/openai_compatible/chat.go

@@ -119,6 +119,9 @@ func (o *OpenAICompatible) Chat(ctx context.Context, chatReq llm.ChatRequest) (*
 		return nil, err
 	}
 
+	if o.tokenUsage.Model == "" && chatResp.Model != "" {
+		o.tokenUsage.Model = chatResp.Model
+	}
 	o.tokenUsage.PromptTokens = chatResp.Usage.PromptTokens
 	o.tokenUsage.CompletionTokens = chatResp.Usage.CompletionTokens
 	o.tokenUsage.TotalTokens = chatResp.Usage.TotalTokens

llm/types.go

@@ -15,10 +15,11 @@ type StreamChatResponse openai.ChatCompletionStreamResponse
 type ChatMessage openai.ChatCompletionMessage
 
 type TokenUsage struct {
-	PromptTokens     int `json:"prompt_tokens"`
-	CompletionTokens int `json:"completion_tokens"`
-	ToolsTokens      int `json:"total_tokens"`
-	TotalTokens      int `json:"total_tokens"`
+	Model            string `json:"model"`
+	PromptTokens     int    `json:"prompt_tokens"`
+	CompletionTokens int    `json:"completion_tokens"`
+	ToolsTokens      int    `json:"tools_tokens"`
+	TotalTokens      int    `json:"total_tokens"`
 }
 
 type ErrorResponse struct {