Refactor(go-models): add token usage for groq (#17574)

### Summary Add token usage reporting for the Groq model provider. Related to #17284. Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-08-01 13:33:48 +08:00 · 2026-07-31 11:37:51 +08:00
parent c6f946a562
commit e4c658c046
2 changed files with 98 additions and 42 deletions
--- a/internal/entity/models/groq.go
+++ b/internal/entity/models/groq.go
@@ -88,9 +88,21 @@ type groqChatChoice struct {
 }

 type groqChatResponse struct {
-	Choices      []groqChatChoice `json:"choices"`
-	Error        interface{}      `json:"error"`
-	FinishReason string           `json:"finish_reason"`
+	ID                string           `json:"id"`
+	Choices           []groqChatChoice `json:"choices"`
+	Error             interface{}      `json:"error"`
+	Usage             *groqChatUsage   `json:"usage"`
+	Created           int64            `json:"created"`
+	Model             string           `json:"model"`
+	Object            string           `json:"object"`
+	SystemFingerprint string           `json:"system_fingerprint"`
+	FinishReason      string           `json:"finish_reason"`
+}
+
+type groqChatUsage struct {
+	PromptTokens     int `json:"prompt_tokens"`
+	CompletionTokens int `json:"completion_tokens"`
+	TotalTokens      int `json:"total_tokens"`
 }

 func (g *GroqModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
@@ -140,27 +152,38 @@ func (g *GroqModel) ChatWithMessages(ctx context.Context, modelName string, mess
 		return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
 	}

-	var result groqChatResponse
-	if err = json.Unmarshal(body, &result); err != nil {
-		return nil, fmt.Errorf("failed to parse response: %w", err)
-	}
-	if result.Error != nil {
-		return nil, fmt.Errorf("groq: upstream error: %v", result.Error)
-	}
-	if len(result.Choices) == 0 {
-		return nil, fmt.Errorf("no choices in response")
-	}
+	return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, _ *ChatConfig) (chatResponseParts, error) {
+		var result groqChatResponse
+		if err := json.Unmarshal(body, &result); err != nil {
+			return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err)
+		}
+		if result.Error != nil {
+			return chatResponseParts{}, fmt.Errorf("groq: upstream error: %v", result.Error)
+		}
+		if len(result.Choices) == 0 {
+			return chatResponseParts{}, fmt.Errorf("no choices in response")
+		}

-	content := result.Choices[0].Message.Content
-	reasonContent := result.Choices[0].Message.ReasoningContent
-	if reasonContent == "" {
-		reasonContent = result.Choices[0].Message.Reasoning
-	}
+		content := result.Choices[0].Message.Content
+		reasonContent := result.Choices[0].Message.ReasoningContent
+		if reasonContent == "" {
+			reasonContent = result.Choices[0].Message.Reasoning
+		}

-	return &ChatResponse{
-		Answer:        &content,
-		ReasonContent: &reasonContent,
-	}, nil
+		parts := chatResponseParts{
+			RequestID:     result.ID,
+			Content:       &content,
+			ReasonContent: &reasonContent,
+		}
+		if result.Usage != nil {
+			parts.Usage = &TokenUsage{
+				PromptTokens:     result.Usage.PromptTokens,
+				CompletionTokens: result.Usage.CompletionTokens,
+				TotalTokens:      result.Usage.TotalTokens,
+			}
+		}
+		return parts, nil
+	})
 }

 func (g *GroqModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
--- a/internal/entity/models/groq_test.go
+++ b/internal/entity/models/groq_test.go
@@ -43,6 +43,40 @@ func newGroqForTest(baseURL string) *GroqModel {
 	)
 }

+// groqChatPayload returns a minimal valid Groq chat completion response.
+// Tests override the fields they care about.
+func groqChatPayload(content, reasoning string, usage *groqUsage) map[string]any {
+	p := map[string]any{
+		"id":      "chatcmpl-test",
+		"object":  "chat.completion",
+		"created": 1730241104,
+		"model":   "llama-3.3-70b-versatile",
+		"choices": []map[string]any{{
+			"index": 0,
+			"message": map[string]any{
+				"role":              "assistant",
+				"content":           content,
+				"reasoning_content": reasoning,
+			},
+			"finish_reason": "stop",
+		}},
+	}
+	if usage != nil {
+		p["usage"] = map[string]any{
+			"prompt_tokens":     usage.promptTokens,
+			"completion_tokens": usage.completionTokens,
+			"total_tokens":      usage.totalTokens,
+		}
+	}
+	return p
+}
+
+type groqUsage struct {
+	promptTokens     int
+	completionTokens int
+	totalTokens      int
+}
+
 func TestGroqName(t *testing.T) {
 	if got := newGroqForTest("http://unused").Name(); got != "groq" {
 		t.Errorf("Name()=%q", got)
@@ -108,14 +142,9 @@ func TestGroqChatHappyPath(t *testing.T) {
 		if _, ok := body["reasoning_effort"]; ok {
 			t.Errorf("reasoning_effort should not be sent: %v", body["reasoning_effort"])
 		}
-		_ = json.NewEncoder(w).Encode(map[string]interface{}{
-			"choices": []map[string]interface{}{{
-				"message": map[string]interface{}{
-					"content":           "pong",
-					"reasoning_content": "thinking",
-				},
-			}},
-		})
+		_ = json.NewEncoder(w).Encode(groqChatPayload("pong", "thinking", &groqUsage{
+			promptTokens: 18, completionTokens: 556, totalTokens: 574,
+		}))
 	})
 	defer srv.Close()

@@ -143,6 +172,12 @@ func TestGroqChatHappyPath(t *testing.T) {
 	if resp.ReasonContent == nil || *resp.ReasonContent != "thinking" {
 		t.Fatalf("ReasonContent=%v, want thinking", resp.ReasonContent)
 	}
+	if resp.Usage == nil {
+		t.Fatalf("Usage=nil, want populated")
+	}
+	if resp.Usage.PromptTokens != 18 || resp.Usage.CompletionTokens != 556 || resp.Usage.TotalTokens != 574 {
+		t.Fatalf("Usage=%#v, want prompt=18 completion=556 total=574", resp.Usage)
+	}
 }

 func TestGroqChatRequiresModelName(t *testing.T) {
@@ -340,16 +375,12 @@ func TestGroqBaseURLTrimsTrailingSlash(t *testing.T) {
 		if r.URL.Path != "/chat/completions" {
 			t.Errorf("path=%s", r.URL.Path)
 		}
-		_ = json.NewEncoder(w).Encode(map[string]interface{}{
-			"choices": []map[string]interface{}{{
-				"message": map[string]interface{}{"content": "ok"},
-			}},
-		})
+		_ = json.NewEncoder(w).Encode(groqChatPayload("ok", "", nil))
 	})
 	defer srv.Close()

 	apiKey := "test-key"
-	_, err := newGroqForTest(srv.URL+"/").ChatWithMessages(
+	resp, err := newGroqForTest(srv.URL+"/").ChatWithMessages(
 		ctx,
 		"llama-3.3-70b-versatile",
 		[]Message{{Role: "user", Content: "x"}},
@@ -360,6 +391,9 @@ func TestGroqBaseURLTrimsTrailingSlash(t *testing.T) {
 	if err != nil {
 		t.Fatalf("ChatWithMessages: %v", err)
 	}
+	if resp.Usage != nil {
+		t.Fatalf("Usage=%#v, want nil for usage-less response", resp.Usage)
+	}
 }

 func TestGroqUsesEmptyRegionCustomBaseURL(t *testing.T) {
@@ -368,11 +402,7 @@ func TestGroqUsesEmptyRegionCustomBaseURL(t *testing.T) {
 		if r.URL.Path != "/chat/completions" {
 			t.Errorf("path=%s", r.URL.Path)
 		}
-		_ = json.NewEncoder(w).Encode(map[string]interface{}{
-			"choices": []map[string]interface{}{{
-				"message": map[string]interface{}{"content": "ok"},
-			}},
-		})
+		_ = json.NewEncoder(w).Encode(groqChatPayload("ok", "", nil))
 	})
 	defer srv.Close()

@@ -382,7 +412,7 @@ func TestGroqUsesEmptyRegionCustomBaseURL(t *testing.T) {
 		map[string]string{"": srv.URL},
 		URLSuffix{Chat: "chat/completions", Models: "models"},
 	)
-	_, err := model.ChatWithMessages(
+	resp, err := model.ChatWithMessages(
 		ctx,
 		"llama-3.3-70b-versatile",
 		[]Message{{Role: "user", Content: "x"}},
@@ -393,6 +423,9 @@ func TestGroqUsesEmptyRegionCustomBaseURL(t *testing.T) {
 	if err != nil {
 		t.Fatalf("ChatWithMessages: %v", err)
 	}
+	if resp.Usage != nil {
+		t.Fatalf("Usage=%#v, want nil for usage-less response", resp.Usage)
+	}
 }

 func TestGroqUnsupportedMethods(t *testing.T) {