diff --git a/internal/entity/models/groq.go b/internal/entity/models/groq.go index 9ea214324b..83b8174127 100644 --- 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 { diff --git a/internal/entity/models/groq_test.go b/internal/entity/models/groq_test.go index b6cc65d1f4..064186416c 100644 --- 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) {