diff --git a/internal/entity/models/chat_usage_test.go b/internal/entity/models/chat_usage_test.go new file mode 100644 index 0000000000..b07a4094b5 --- /dev/null +++ b/internal/entity/models/chat_usage_test.go @@ -0,0 +1,85 @@ +package models + +import ( + "net/http" + "net/http/httptest" + "ragflow/internal/common" + "testing" + "time" +) + +func TestProviderChatUsage(t *testing.T) { + type providerCase struct { + name string + new func(string) ModelDriver + } + + cases := []providerCase{ + { + name: "deepseek", + new: func(baseURL string) ModelDriver { + return NewDeepSeekModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + }, + { + name: "hunyuan", + new: func(baseURL string) ModelDriver { + return NewHunyuanModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + }, + { + name: "jina", + new: func(baseURL string) ModelDriver { + return NewJinaModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + }, + { + name: "gitee", + new: func(baseURL string) ModelDriver { + return NewGiteeModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + }, + { + name: "openrouter", + new: func(baseURL string) ModelDriver { + return NewOpenRouterModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + }, + { + name: "jiekouai", + new: func(baseURL string) ModelDriver { + return NewJieKouAIModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(`{"id":"chat-request","choices":[{"message":{"content":"answer"}}],"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`)) + })) + defer server.Close() + + apiKey := "test-key" + usage := &common.ModelUsage{StartAt: time.Now()} + response, err := tc.new(server.URL).ChatWithMessages( + t.Context(), + "test-model", + []Message{{Role: "user", Content: "hello"}}, + &APIConfig{ApiKey: &apiKey}, + nil, + usage, + ) + if err != nil { + t.Fatalf("ChatWithMessages: %v", err) + } + if response.Usage == nil || response.Usage.PromptTokens != 3 || response.Usage.CompletionTokens != 5 || response.Usage.TotalTokens != 8 { + t.Fatalf("response usage=%#v", response.Usage) + } + assertModelUsage(t, usage, 3, 5, 8) + if usage.Type != "chat" { + t.Fatalf("usage type=%q, want chat", usage.Type) + } + }) + } +} diff --git a/internal/entity/models/gitee.go b/internal/entity/models/gitee.go index d89df8a4b4..c5eb2883ca 100644 --- a/internal/entity/models/gitee.go +++ b/internal/entity/models/gitee.go @@ -53,6 +53,34 @@ func (g *GiteeModel) Name() string { return "GiteeAI" } +// GiteeChatResponse mirrors GiteeAI's OpenAI-compatible chat response. +type GiteeChatResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []struct { + FinishReason string `json:"finish_reason"` + Index int `json:"index"` + Logprobs any `json:"logprobs"` + Message struct { + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + Role string `json:"role"` + ToolCalls []map[string]any `json:"tool_calls"` + } `json:"message"` + } `json:"choices"` + SystemFingerprint string `json:"system_fingerprint"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + } `json:"usage"` +} + // ChatWithMessages sends multiple messages with roles and returns response func (g *GiteeModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := g.baseModel.APIConfigCheck(apiConfig); err != nil { @@ -117,69 +145,48 @@ func (g *GiteeModel) ChatWithMessages(ctx context.Context, modelName string, mes return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - // Parse response - var result map[string]interface{} - if err := json.Unmarshal(body, &result); err != nil { - return nil, fmt.Errorf("failed to parse response: %w", err) - } + return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) { + var result GiteeChatResponse + if err := json.Unmarshal(body, &result); err != nil { + return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err) + } + if len(result.Choices) == 0 { + return chatResponseParts{}, fmt.Errorf("no choices in response") + } - choices, ok := result["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil, fmt.Errorf("no choices in response") - } + choice := result.Choices[0] + content := choice.Message.Content + if content == "" && len(choice.Message.ToolCalls) == 0 { + return chatResponseParts{}, fmt.Errorf("response contains neither content nor tool calls") + } - firstChoice, ok := choices[0].(map[string]interface{}) - if !ok { - return nil, fmt.Errorf("invalid choice format") - } - - messageMap, ok := firstChoice["message"].(map[string]interface{}) - if !ok { - return nil, fmt.Errorf("invalid message format") - } - - content, _ := messageMap["content"].(string) - - // Handle thinking/reasoning if enabled - var reasonContent string - if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { - // Try to get reasoning_content directly first - if rc, ok := messageMap["reasoning_content"].(string); ok && rc != "" { - reasonContent = rc - if reasonContent[0] == '\n' { + reasonContent := "" + if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking { + reasonContent = choice.Message.ReasoningContent + if reasonContent == "" { + reasoning, answer := GetThinkingAndAnswer(chatConfig.ModelClass, &content) + if reasoning != nil { + reasonContent = *reasoning + content = *answer + } + } + if strings.HasPrefix(reasonContent, "\n") { reasonContent = reasonContent[1:] } - } else { - // Fall back to parsing tags from content - reasoning, answer := GetThinkingAndAnswer(chatModelConfig.ModelClass, &content) - if reasoning != nil { - reasonContent = *reasoning - content = *answer - } } - } - var toolCalls []map[string]any - if tcs, ok := messageMap["tool_calls"].([]any); ok { - for _, tc := range tcs { - if toolCall, ok := tc.(map[string]any); ok { - toolCalls = append(toolCalls, toolCall) - } - } - } - - chatResponse := &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - } - if pt, ct, tt := extractUsageFromMap(result); tt > 0 { - chatResponse.Usage = &TokenUsage{ - PromptTokens: pt, CompletionTokens: ct, TotalTokens: tt, - } - } - - return chatResponse, nil + return chatResponseParts{ + RequestID: result.ID, + Content: &content, + ReasonContent: &reasonContent, + ToolCalls: choice.Message.ToolCalls, + Usage: &TokenUsage{ + PromptTokens: result.Usage.PromptTokens, + CompletionTokens: result.Usage.CompletionTokens, + TotalTokens: result.Usage.TotalTokens, + }, + }, nil + }) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) @@ -191,13 +198,18 @@ func (g *GiteeModel) ChatStreamlyWithSender(ctx context.Context, modelName strin if len(messages) == 0 { return fmt.Errorf("messages is empty") } + if chatModelConfig != nil { + chatModelConfig.ToolCallsResult = nil + chatModelConfig.UsageResult = nil + } resolvedBaseURL, err := g.baseModel.GetBaseURL(apiConfig) if err != nil { return err } - url := fmt.Sprintf("%s/chat/completions", resolvedBaseURL) + url := fmt.Sprintf("%s/%s", strings.TrimSuffix(resolvedBaseURL, "/"), g.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]any{"include_usage": true} if chatModelConfig != nil { if chatModelConfig.Thinking != nil { @@ -247,23 +259,38 @@ func (g *GiteeModel) ChatStreamlyWithSender(ctx context.Context, modelName strin sawTerminal := false accumulatedToolCalls := make(map[int]map[string]any) - done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { - choices, ok := event["choices"].([]interface{}) + done, err := ParseSSEStream[map[string]any](resp.Body, func(event map[string]any) error { + common.Info(fmt.Sprintf("%v", event)) + + tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event) + if usageErr != nil { + return usageErr + } + if found { + applyStreamUsage(chatModelConfig, modelUsage, tokenUsage) + } + choices, ok := event["choices"].([]any) if !ok || len(choices) == 0 { return nil } - - firstChoice, ok := choices[0].(map[string]interface{}) + choice, ok := choices[0].(map[string]any) if !ok { return nil } - - delta, ok := firstChoice["delta"].(map[string]interface{}) + if finishReason, ok := choice["finish_reason"].(string); ok && finishReason != "" { + sawTerminal = true + } + delta, ok := choice["delta"].(map[string]any) if !ok { return nil } accumulateToolCallDeltas(delta, accumulatedToolCalls) + if reasoning, ok := delta["reasoning_content"].(string); ok && reasoning != "" { + if err := sender(nil, &reasoning); err != nil { + return err + } + } content, ok := delta["content"].(string) if ok && content != "" { @@ -297,10 +324,6 @@ func (g *GiteeModel) ChatStreamlyWithSender(ctx context.Context, modelName strin } } - finishReason, ok := firstChoice["finish_reason"].(string) - if ok && finishReason != "" { - sawTerminal = true - } return nil }) if err != nil { @@ -327,22 +350,21 @@ func (g *GiteeModel) ChatStreamlyWithSender(ctx context.Context, modelName strin return nil } -type giteeEmbeddingResponse struct { - Object string `json:"object"` - Data []giteeEmbeddingData `json:"data"` - Model string `json:"model"` - Usage giteeUsage `json:"usage"` -} - -type giteeEmbeddingData struct { - Object string `json:"object"` - Embedding []float64 `json:"embedding"` - Index int `json:"index"` -} - -type giteeUsage struct { - PromptTokens int `json:"prompt_tokens"` - TotalTokens int `json:"total_tokens"` +// GiteeEmbeddingResponse mirrors GiteeAI's embeddings response. +type GiteeEmbeddingResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Model string `json:"model"` + Data []struct { + Object string `json:"object"` + Embedding []float64 `json:"embedding"` + Index int `json:"index"` + } `json:"data"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` } // Embed embeds a list of texts into embeddings @@ -406,10 +428,15 @@ func (g *GiteeModel) Embed(ctx context.Context, modelName *string, texts []strin return nil, fmt.Errorf("Gitee embeddings API error: %s, body: %s", resp.Status, string(body)) } - var parsed giteeEmbeddingResponse + var parsed GiteeEmbeddingResponse if err = json.Unmarshal(body, &parsed); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } + recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{ + PromptTokens: parsed.Usage.PromptTokens, + CompletionTokens: parsed.Usage.CompletionTokens, + TotalTokens: parsed.Usage.TotalTokens, + }, "embedding") var embeddings []EmbeddingData for _, dataElem := range parsed.Data { @@ -430,6 +457,24 @@ type giteeRerankRequest struct { ReturnDocuments bool `json:"return_documents"` } +// GiteeRerankResponse mirrors GiteeAI's rerank response. +type GiteeRerankResponse struct { + ID string `json:"id"` + Model string `json:"model"` + Results []struct { + Index int `json:"index"` + Document struct { + Text string `json:"text"` + } `json:"document"` + RelevanceScore float64 `json:"relevance_score"` + } `json:"results"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` +} + // Rerank calculates similarity scores between query and documents func (g *GiteeModel) Rerank(ctx context.Context, modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig, modelUsage *common.ModelUsage) (*RerankResponse, error) { if err := g.baseModel.APIConfigCheck(apiConfig); err != nil { @@ -452,9 +497,9 @@ func (g *GiteeModel) Rerank(ctx context.Context, modelName *string, query string url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), g.baseModel.URLSuffix.Rerank) - var topN = rerankConfig.TopN - if rerankConfig.TopN == 0 { - topN = len(documents) + topN := len(documents) + if rerankConfig != nil && rerankConfig.TopN > 0 { + topN = rerankConfig.TopN } reqBody := giteeRerankRequest{ @@ -496,10 +541,23 @@ func (g *GiteeModel) Rerank(ctx context.Context, modelName *string, query string return nil, fmt.Errorf("gitee rerank API error: %s, body: %s", resp.Status, string(body)) } - var rerankResponse RerankResponse - if err = json.Unmarshal(body, &rerankResponse); err != nil { + var parsed GiteeRerankResponse + if err = json.Unmarshal(body, &parsed); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } + recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{ + PromptTokens: parsed.Usage.PromptTokens, + CompletionTokens: parsed.Usage.CompletionTokens, + TotalTokens: parsed.Usage.TotalTokens, + }, "rerank") + + rerankResponse := RerankResponse{Data: make([]RerankResult, 0, len(parsed.Results))} + for _, result := range parsed.Results { + rerankResponse.Data = append(rerankResponse.Data, RerankResult{ + Index: result.Index, + RelevanceScore: result.RelevanceScore, + }) + } return &rerankResponse, nil } diff --git a/internal/entity/models/gitee_test.go b/internal/entity/models/gitee_test.go index 6dfeb62620..2764a29666 100644 --- a/internal/entity/models/gitee_test.go +++ b/internal/entity/models/gitee_test.go @@ -46,6 +46,40 @@ func newGiteeForListModelsTest(baseURL string) *GiteeModel { return NewGiteeModel(map[string]string{"default": baseURL}, URLSuffix{Models: "models"}) } +func newGiteeForChatTest(baseURL string) *GiteeModel { + return NewGiteeModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) +} + +func TestGiteeStreamAcceptsTerminalWithoutDelta(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("method=%s, want POST", r.Method) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, `data: {"choices":[{"finish_reason":"stop"}]}`+"\n\n") + })) + defer srv.Close() + + apiKey := "test-key" + var sawDone bool + err := newGiteeForChatTest(srv.URL).ChatStreamlyWithSender( + t.Context(), + "gitee-model", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil, nil, + func(content *string, _ *string) error { + sawDone = content != nil && *content == "[DONE]" + return nil + }, + ) + if err != nil { + t.Fatalf("stream: %v", err) + } + if !sawDone { + t.Fatal("expected [DONE] sentinel") + } +} + func deepSeekAliasModelsForTest(t *testing.T) map[string]Model { t.Helper() diff --git a/internal/entity/models/hunyuan.go b/internal/entity/models/hunyuan.go index 4ce824e322..be0dbf0fd3 100644 --- a/internal/entity/models/hunyuan.go +++ b/internal/entity/models/hunyuan.go @@ -51,6 +51,49 @@ func (h *HunyuanModel) Name() string { return "Tencent Hunyuan" } +// HunyuanChatResponse mirrors Tencent Hunyuan's OpenAI-compatible chat response. +type HunyuanChatResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []struct { + FinishReason string `json:"finish_reason"` + Index int `json:"index"` + Logprobs any `json:"logprobs"` + Message struct { + Content string `json:"content"` + Reasoning string `json:"reasoning"` + ReasoningContent string `json:"reasoning_content"` + Role string `json:"role"` + ToolCalls []map[string]any `json:"tool_calls"` + } `json:"message"` + } `json:"choices"` + SystemFingerprint string `json:"system_fingerprint"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` +} + +// HunyuanEmbeddingResponse mirrors Tencent Hunyuan's embeddings response. +type HunyuanEmbeddingResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Model string `json:"model"` + Data []struct { + Object string `json:"object"` + Embedding []float64 `json:"embedding"` + Index int `json:"index"` + } `json:"data"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` +} + func (h *HunyuanModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := h.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err @@ -96,42 +139,35 @@ func (h *HunyuanModel) ChatWithMessages(ctx context.Context, modelName string, m return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - var result map[string]interface{} - if err = json.Unmarshal(body, &result); err != nil { - return nil, fmt.Errorf("failed to parse response: %w", err) - } + return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, _ *ChatConfig) (chatResponseParts, error) { + var result HunyuanChatResponse + if err := json.Unmarshal(body, &result); err != nil { + return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err) + } + if len(result.Choices) == 0 { + return chatResponseParts{}, fmt.Errorf("no choices in response") + } - choices, ok := result["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil, fmt.Errorf("no choices in response") - } - - firstChoice, ok := choices[0].(map[string]interface{}) - if !ok { - return nil, fmt.Errorf("invalid choice format") - } - - messageMap, ok := firstChoice["message"].(map[string]interface{}) - if !ok { - return nil, fmt.Errorf("invalid message format") - } - - content, ok := messageMap["content"].(string) - toolCalls := extractToolCalls(messageMap) - if !ok && len(toolCalls) == 0 { - return nil, fmt.Errorf("invalid content format") - } - - reasonContent := "" - if r, ok := messageMap["reasoning_content"].(string); ok { - reasonContent = r - } - - return &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - }, nil + choice := result.Choices[0] + if choice.Message.Content == "" && len(choice.Message.ToolCalls) == 0 { + return chatResponseParts{}, fmt.Errorf("response contains neither content nor tool calls") + } + reasonContent := choice.Message.ReasoningContent + if reasonContent == "" { + reasonContent = choice.Message.Reasoning + } + return chatResponseParts{ + RequestID: result.ID, + Content: &choice.Message.Content, + ReasonContent: &reasonContent, + ToolCalls: choice.Message.ToolCalls, + Usage: &TokenUsage{ + PromptTokens: result.Usage.PromptTokens, + CompletionTokens: result.Usage.CompletionTokens, + TotalTokens: result.Usage.TotalTokens, + }, + }, nil + }) } // ChatStreamlyWithSender opens the SSE chat-completions @@ -149,6 +185,10 @@ func (h *HunyuanModel) ChatStreamlyWithSender(ctx context.Context, modelName str if err := validateStreamConfig(chatModelConfig); err != nil { return err } + if chatModelConfig != nil { + chatModelConfig.ToolCallsResult = nil + chatModelConfig.UsageResult = nil + } baseURL, err := h.baseModel.GetBaseURL(apiConfig) if err != nil { @@ -157,13 +197,17 @@ func (h *HunyuanModel) ChatStreamlyWithSender(ctx context.Context, modelName str baseURL = strings.TrimSuffix(baseURL, "/") url := fmt.Sprintf("%s/%s", baseURL, h.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]any{"include_usage": true} jsonData, err := json.Marshal(reqBody) if err != nil { return fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) + ctx, cancel := context.WithTimeout(ctx, streamCallTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData)) if err != nil { return fmt.Errorf("failed to create request: %w", err) } @@ -183,36 +227,50 @@ func (h *HunyuanModel) ChatStreamlyWithSender(ctx context.Context, modelName str sawTerminal := false accumulatedToolCalls := make(map[int]map[string]any) - done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { + done, err := ParseSSEStream[map[string]any](resp.Body, func(event map[string]any) error { + common.Info(fmt.Sprintf("%v", event)) + if apiErr, ok := event["error"]; ok { return fmt.Errorf("hunyuan: upstream stream error: %v", apiErr) } - - choices, ok := event["choices"].([]interface{}) + tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event) + if usageErr != nil { + return usageErr + } + if found { + applyStreamUsage(chatModelConfig, modelUsage, tokenUsage) + } + choices, ok := event["choices"].([]any) if !ok || len(choices) == 0 { return nil } - firstChoice, ok := choices[0].(map[string]interface{}) + + choice, ok := choices[0].(map[string]any) if !ok { return nil } - if delta, ok := firstChoice["delta"].(map[string]interface{}); ok { - accumulateToolCallDeltas(delta, accumulatedToolCalls) - if r, ok := delta["reasoning_content"].(string); ok && r != "" { - rr := r - if err := sender(nil, &rr); err != nil { - return err - } - } - if c, ok := delta["content"].(string); ok && c != "" { - cc := c - if err := sender(&cc, nil); err != nil { - return err - } + if finishReason, ok := choice["finish_reason"].(string); ok && finishReason != "" { + sawTerminal = true + } + delta, ok := choice["delta"].(map[string]any) + if !ok { + return nil + } + accumulateToolCallDeltas(delta, accumulatedToolCalls) + if reasoning, ok := delta["reasoning_content"].(string); ok && reasoning != "" { + if err := sender(nil, &reasoning); err != nil { + return err } } - if finish, ok := firstChoice["finish_reason"].(string); ok && finish != "" { - sawTerminal = true + if reasoning, ok := delta["reasoning"].(string); ok && reasoning != "" { + if err := sender(nil, &reasoning); err != nil { + return err + } + } + if content, ok := delta["content"].(string); ok && content != "" { + if err := sender(&content, nil); err != nil { + return err + } } return nil }) @@ -338,12 +396,7 @@ func (h *HunyuanModel) Embed(ctx context.Context, modelName *string, texts []str return nil, fmt.Errorf("Hunyuan embedding API error: status %d, body: %s", resp.StatusCode, string(body)) } - var parsedResponse struct { - Data []struct { - Embedding []float64 `json:"embedding"` - Index int `json:"index"` - } `json:"data"` - } + var parsedResponse HunyuanEmbeddingResponse if err = json.Unmarshal(body, &parsedResponse); err != nil { return nil, fmt.Errorf("failed to decode response: %w", err) @@ -352,6 +405,11 @@ func (h *HunyuanModel) Embed(ctx context.Context, modelName *string, texts []str if len(parsedResponse.Data) == 0 { return nil, fmt.Errorf("hunyuan embedding response contains no data: %s", string(body)) } + recordResponseUsage(modelUsage, parsedResponse.ID, &TokenUsage{ + PromptTokens: parsedResponse.Usage.PromptTokens, + CompletionTokens: parsedResponse.Usage.CompletionTokens, + TotalTokens: parsedResponse.Usage.TotalTokens, + }, "embedding") embeddings := make([]EmbeddingData, 0, len(parsedResponse.Data)) for _, dataElem := range parsedResponse.Data { diff --git a/internal/entity/models/hunyuan_test.go b/internal/entity/models/hunyuan_test.go index 1c521c42c6..cdf2ec6305 100644 --- a/internal/entity/models/hunyuan_test.go +++ b/internal/entity/models/hunyuan_test.go @@ -347,6 +347,32 @@ func TestHunyuanStreamFailsWithoutTerminal(t *testing.T) { } } +func TestHunyuanStreamAcceptsTerminalWithoutDelta(t *testing.T) { + srv := newHunyuanSSEServer(t, "/chat/completions", + `data: {"choices":[{"finish_reason":"stop"}]}`+"\n\n", + ) + defer srv.Close() + + apiKey := "test-key" + var sawDone bool + err := newHunyuanForTest(srv.URL).ChatStreamlyWithSender( + t.Context(), + "hunyuan-pro", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil, nil, + func(content *string, _ *string) error { + sawDone = content != nil && *content == "[DONE]" + return nil + }, + ) + if err != nil { + t.Fatalf("stream: %v", err) + } + if !sawDone { + t.Fatal("expected [DONE] sentinel") + } +} + func TestHunyuanStreamRejectsMalformedFrame(t *testing.T) { srv := newHunyuanSSEServer(t, "/chat/completions", `data: {"choices":[{"delta":{"content":"ok"}}]}`+"\n"+ diff --git a/internal/entity/models/jiekouai.go b/internal/entity/models/jiekouai.go index 05d5501cf0..98631a5ca3 100644 --- a/internal/entity/models/jiekouai.go +++ b/internal/entity/models/jiekouai.go @@ -62,6 +62,67 @@ func validateJieKouAIModelName(modelName *string) (string, error) { return strings.TrimSpace(*modelName), nil } +// JieKouAIChatResponse mirrors Jiekou.AI's OpenAI-compatible chat response. +type JieKouAIChatResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []struct { + FinishReason string `json:"finish_reason"` + Index int `json:"index"` + Logprobs any `json:"logprobs"` + Message struct { + Content string `json:"content"` + Reasoning string `json:"reasoning"` + ReasoningContent string `json:"reasoning_content"` + Role string `json:"role"` + ToolCalls []map[string]any `json:"tool_calls"` + } `json:"message"` + } `json:"choices"` + SystemFingerprint string `json:"system_fingerprint"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` +} + +// JieKouAIEmbeddingResponse mirrors Jiekou.AI's embeddings response. +type JieKouAIEmbeddingResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Model string `json:"model"` + Data []struct { + Object string `json:"object"` + Embedding []float64 `json:"embedding"` + Index int `json:"index"` + } `json:"data"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` +} + +// JieKouAIRerankResponse mirrors Jiekou.AI's rerank response. +type JieKouAIRerankResponse struct { + ID string `json:"id"` + Model string `json:"model"` + Results []struct { + Index int `json:"index"` + Document struct { + Text string `json:"text"` + } `json:"document"` + RelevanceScore float64 `json:"relevance_score"` + } `json:"results"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` +} + func (j *JieKouAIModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := j.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err @@ -86,9 +147,7 @@ func (j *JieKouAIModel) ChatWithMessages(ctx context.Context, modelName string, if chatModelConfig != nil { // For zai-org/glm-4.5: https://docs.jiekou.ai/docs/models/reference-llm-create-chat-completion if chatModelConfig.Thinking != nil { - if *chatModelConfig.Thinking { - reqBody["enable_thinking"] = *chatModelConfig.Thinking - } + reqBody["enable_thinking"] = *chatModelConfig.Thinking } } @@ -97,7 +156,10 @@ func (j *JieKouAIModel) ChatWithMessages(ctx context.Context, modelName string, return nil, fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData)) + ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData)) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } @@ -121,52 +183,39 @@ func (j *JieKouAIModel) ChatWithMessages(ctx context.Context, modelName string, return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - // Parse response - var result map[string]interface{} - if err = json.Unmarshal(body, &result); err != nil { - return nil, fmt.Errorf("failed to parse response: %w", err) - } - - choices, ok := result["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil, fmt.Errorf("no choices in response") - } - - firstChoice, ok := choices[0].(map[string]interface{}) - if !ok { - return nil, fmt.Errorf("invalid choice format") - } - - messageMap, ok := firstChoice["message"].(map[string]interface{}) - if !ok { - return nil, fmt.Errorf("invalid message format") - } - - content, ok := messageMap["content"].(string) - toolCalls := extractToolCalls(messageMap) - if !ok && len(toolCalls) == 0 { - return nil, fmt.Errorf("invalid content format") - } - - var reasonContent string - if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { - reasonContent, ok = messageMap["reasoning_content"].(string) - if !ok { - return nil, fmt.Errorf("invalid content format") + return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) { + var result JieKouAIChatResponse + if err := json.Unmarshal(body, &result); err != nil { + return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err) } - // if first char of reasonContent is \n remove the '\n' - if reasonContent != "" && reasonContent[0] == '\n' { - reasonContent = reasonContent[1:] + if len(result.Choices) == 0 { + return chatResponseParts{}, fmt.Errorf("no choices in response") } - } - chatResponse := &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - } - - return chatResponse, nil + choice := result.Choices[0] + if choice.Message.Content == "" && len(choice.Message.ToolCalls) == 0 { + return chatResponseParts{}, fmt.Errorf("response contains neither content nor tool calls") + } + reasonContent := "" + if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking { + reasonContent = choice.Message.ReasoningContent + if reasonContent == "" { + reasonContent = choice.Message.Reasoning + } + reasonContent = strings.TrimPrefix(reasonContent, "\n") + } + return chatResponseParts{ + RequestID: result.ID, + Content: &choice.Message.Content, + ReasonContent: &reasonContent, + ToolCalls: choice.Message.ToolCalls, + Usage: &TokenUsage{ + PromptTokens: result.Usage.PromptTokens, + CompletionTokens: result.Usage.CompletionTokens, + TotalTokens: result.Usage.TotalTokens, + }, + }, nil + }) } func (j *JieKouAIModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { @@ -183,6 +232,10 @@ func (j *JieKouAIModel) ChatStreamlyWithSender(ctx context.Context, modelName st if len(messages) == 0 { return fmt.Errorf("messages is empty") } + if modelConfig != nil { + modelConfig.ToolCallsResult = nil + modelConfig.UsageResult = nil + } resolvedBaseURL, err := j.baseModel.GetBaseURL(apiConfig) if err != nil { @@ -192,12 +245,11 @@ func (j *JieKouAIModel) ChatStreamlyWithSender(ctx context.Context, modelName st // Build request body with streaming enabled reqBody := buildRequestBody(modelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]any{"include_usage": true} if modelConfig != nil { if modelConfig.Thinking != nil { - if *modelConfig.Thinking { - reqBody["enable_thinking"] = *modelConfig.Thinking - } + reqBody["enable_thinking"] = *modelConfig.Thinking } } @@ -206,7 +258,10 @@ func (j *JieKouAIModel) ChatStreamlyWithSender(ctx context.Context, modelName st return fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData)) + ctx, cancel := context.WithTimeout(ctx, streamCallTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData)) if err != nil { return fmt.Errorf("failed to create request: %w", err) } @@ -226,20 +281,31 @@ func (j *JieKouAIModel) ChatStreamlyWithSender(ctx context.Context, modelName st return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - // SSE parsing: read line by line + sawTerminal := false accumulatedToolCalls := make(map[int]map[string]any) - if _, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { - choices, ok := event["choices"].([]interface{}) + done, err := ParseSSEStream[map[string]any](resp.Body, func(event map[string]any) error { + common.Info(fmt.Sprintf("%v", event)) + + tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event) + if usageErr != nil { + return usageErr + } + if found { + applyStreamUsage(modelConfig, modelUsage, tokenUsage) + } + choices, ok := event["choices"].([]any) if !ok || len(choices) == 0 { return nil } - firstChoice, ok := choices[0].(map[string]interface{}) + choice, ok := choices[0].(map[string]any) if !ok { return nil } - - delta, ok := firstChoice["delta"].(map[string]interface{}) + if finishReason, ok := choice["finish_reason"].(string); ok && finishReason != "" { + sawTerminal = true + } + delta, ok := choice["delta"].(map[string]any) if !ok { return nil } @@ -261,9 +327,13 @@ func (j *JieKouAIModel) ChatStreamlyWithSender(ctx context.Context, modelName st } return nil - }); err != nil { + }) + if err != nil { return fmt.Errorf("failed to scan response body: %w", err) } + if !done && !sawTerminal { + return fmt.Errorf("jiekouai: stream ended before [DONE] or finish_reason") + } setSortedToolCallsResult(modelConfig, accumulatedToolCalls) // Send [DONE] marker for OpenAI compatibility @@ -305,7 +375,10 @@ func (j *JieKouAIModel) Embed(ctx context.Context, modelName *string, texts []st return nil, fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData)) + ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData)) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } @@ -328,12 +401,7 @@ func (j *JieKouAIModel) Embed(ctx context.Context, modelName *string, texts []st return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - var parsedResponse struct { - Data []struct { - Embedding []float64 `json:"embedding"` - Index int `json:"index"` - } `json:"data"` - } + var parsedResponse JieKouAIEmbeddingResponse if err = json.Unmarshal(body, &parsedResponse); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) @@ -342,6 +410,11 @@ func (j *JieKouAIModel) Embed(ctx context.Context, modelName *string, texts []st if len(parsedResponse.Data) == 0 { return nil, fmt.Errorf("failed to parse response") } + recordResponseUsage(modelUsage, parsedResponse.ID, &TokenUsage{ + PromptTokens: parsedResponse.Usage.PromptTokens, + CompletionTokens: parsedResponse.Usage.CompletionTokens, + TotalTokens: parsedResponse.Usage.TotalTokens, + }, "embedding") var embeddings []EmbeddingData for _, embedding := range parsedResponse.Data { @@ -391,7 +464,10 @@ func (j *JieKouAIModel) Rerank(ctx context.Context, modelName *string, query str return nil, fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData)) + ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData)) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } @@ -414,16 +490,16 @@ func (j *JieKouAIModel) Rerank(ctx context.Context, modelName *string, query str return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - var rerankResp struct { - Results []struct { - Index int `json:"index"` - RelevanceScore float64 `json:"relevance_score"` - } `json:"results"` - } + var rerankResp JieKouAIRerankResponse if err = json.Unmarshal(body, &rerankResp); err != nil { return nil, fmt.Errorf("failed to decode response: %w", err) } + recordResponseUsage(modelUsage, rerankResp.ID, &TokenUsage{ + PromptTokens: rerankResp.Usage.PromptTokens, + CompletionTokens: rerankResp.Usage.CompletionTokens, + TotalTokens: rerankResp.Usage.TotalTokens, + }, "rerank") var rerankResponse RerankResponse for _, result := range rerankResp.Results { diff --git a/internal/entity/models/jiekouai_test.go b/internal/entity/models/jiekouai_test.go index c87243a474..8037502827 100644 --- a/internal/entity/models/jiekouai_test.go +++ b/internal/entity/models/jiekouai_test.go @@ -108,6 +108,35 @@ func TestJieKouAIChatForcesNonStreaming(t *testing.T) { } } +func TestJieKouAIChatSendsExplicitThinkingFalse(t *testing.T) { + srv := newJieKouAIServer(t, func(t *testing.T, r *http.Request, body map[string]interface{}, w http.ResponseWriter) { + if body["enable_thinking"] != false { + t.Errorf("enable_thinking=%v, want false", body["enable_thinking"]) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "id": "jiekou-chat", + "choices": []map[string]interface{}{{ + "message": map[string]interface{}{"content": "answer"}, + }}, + }) + }) + defer srv.Close() + + apiKey := "test-key" + thinking := false + _, err := newJieKouAIForTest(srv.URL).ChatWithMessages( + t.Context(), + "gpt-5", + []Message{{Role: "user", Content: "ping"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{Thinking: &thinking}, + nil, + ) + if err != nil { + t.Fatalf("ChatWithMessages: %v", err) + } +} + func TestJieKouAIStreamForcesStreaming(t *testing.T) { srv := newJieKouAIServer(t, func(t *testing.T, r *http.Request, body map[string]interface{}, w http.ResponseWriter) { if r.URL.Path != "/openai/v1/chat/completions" { diff --git a/internal/entity/models/jina.go b/internal/entity/models/jina.go index a55b5c61a2..fe2200f6ef 100644 --- a/internal/entity/models/jina.go +++ b/internal/entity/models/jina.go @@ -55,6 +55,70 @@ func (j *JinaModel) Name() string { return "jina" } +// JinaChatResponse mirrors Jina's OpenAI-compatible chat response. +type JinaChatResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []struct { + FinishReason string `json:"finish_reason"` + Index int `json:"index"` + Logprobs any `json:"logprobs"` + Message struct { + Content string `json:"content"` + Reasoning string `json:"reasoning"` + ReasoningContent string `json:"reasoning_content"` + Role string `json:"role"` + ToolCalls []map[string]any `json:"tool_calls"` + } `json:"message"` + } `json:"choices"` + SystemFingerprint string `json:"system_fingerprint"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + } `json:"usage"` +} + +// JinaEmbeddingResponse mirrors Jina's embeddings response. Embeddings is +// populated by multivector models such as jina-embeddings-v4. +type JinaEmbeddingResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Model string `json:"model"` + Data []struct { + Object string `json:"object"` + Embedding []float64 `json:"embedding"` + Embeddings [][]float64 `json:"embeddings"` + Index int `json:"index"` + } `json:"data"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` +} + +// JinaRerankResponse mirrors Jina's rerank response. +type JinaRerankResponse struct { + ID string `json:"id"` + Model string `json:"model"` + Results []struct { + Index int `json:"index"` + Document struct { + Text string `json:"text"` + } `json:"document"` + RelevanceScore float64 `json:"relevance_score"` + } `json:"results"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` +} + func (j *JinaModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := j.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err @@ -106,42 +170,40 @@ func (j *JinaModel) ChatWithMessages(ctx context.Context, modelName string, mess return nil, fmt.Errorf("Jina chat API error: status %d, body: %s", resp.StatusCode, string(body)) } - var result map[string]interface{} - if err = json.Unmarshal(body, &result); err != nil { - return nil, fmt.Errorf("failed to parse response: %w", err) - } + return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, _ *ChatConfig) (chatResponseParts, error) { + var result JinaChatResponse + if err := json.Unmarshal(body, &result); err != nil { + return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err) + } + if len(result.Choices) == 0 { + return chatResponseParts{}, fmt.Errorf("no choices in response") + } - choices, ok := result["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil, fmt.Errorf("no choices in response") - } - - firstChoice, ok := choices[0].(map[string]interface{}) - if !ok { - return nil, fmt.Errorf("invalid choice format") - } - - messageMap, ok := firstChoice["message"].(map[string]interface{}) - if !ok { - return nil, fmt.Errorf("invalid message format") - } - - toolCalls := extractToolCalls(messageMap) - content, hasContent := messageMap["content"].(string) - if !hasContent && len(toolCalls) == 0 { - return nil, fmt.Errorf("invalid content format") - } - - reasonContent := "" - return &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - }, nil + choice := result.Choices[0] + if choice.Message.Content == "" && len(choice.Message.ToolCalls) == 0 { + return chatResponseParts{}, fmt.Errorf("response contains neither content nor tool calls") + } + reasonContent := choice.Message.ReasoningContent + if reasonContent == "" { + reasonContent = choice.Message.Reasoning + } + reasonContent = strings.TrimPrefix(reasonContent, "\n") + return chatResponseParts{ + RequestID: result.ID, + Content: &choice.Message.Content, + ReasonContent: &reasonContent, + ToolCalls: choice.Message.ToolCalls, + Usage: &TokenUsage{ + PromptTokens: result.Usage.PromptTokens, + CompletionTokens: result.Usage.CompletionTokens, + TotalTokens: result.Usage.TotalTokens, + }, + }, nil + }) } func (j *JinaModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { - //TODO implement me: https://api.jina.ai/docs#/Search%20Foundation%20Models/chat_completions_v1_chat_completions_post + // Jina's public API does not expose a streaming chat-completions endpoint. return fmt.Errorf("jina does not implement ChatStreamlyWithSender(not available for now)") } @@ -153,6 +215,9 @@ func (j *JinaModel) Embed(ctx context.Context, modelName *string, texts []string if len(texts) == 0 { return []EmbeddingData{}, nil } + if modelName == nil || strings.TrimSpace(*modelName) == "" { + return nil, fmt.Errorf("model name is required") + } resolvedBaseURL, err := j.baseModel.GetBaseURL(apiConfig) if err != nil { @@ -170,7 +235,10 @@ func (j *JinaModel) Embed(ctx context.Context, modelName *string, texts []string return nil, fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData)) + ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData)) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } @@ -193,12 +261,7 @@ func (j *JinaModel) Embed(ctx context.Context, modelName *string, texts []string return nil, fmt.Errorf("Jina embedding API error: status %d, body: %s", resp.StatusCode, string(body)) } - var parsedResponse struct { - Data []struct { - Embedding []float64 `json:"embedding"` - Index int `json:"index"` - } `json:"data"` - } + var parsedResponse JinaEmbeddingResponse if err = json.Unmarshal(body, &parsedResponse); err != nil { return nil, fmt.Errorf("failed to decode response: %w", err) @@ -207,11 +270,37 @@ func (j *JinaModel) Embed(ctx context.Context, modelName *string, texts []string if len(parsedResponse.Data) == 0 { return nil, fmt.Errorf("Jina embedding response contains no data: %s", string(body)) } + recordResponseUsage(modelUsage, parsedResponse.ID, &TokenUsage{ + PromptTokens: parsedResponse.Usage.PromptTokens, + TotalTokens: parsedResponse.Usage.TotalTokens, + }, "embedding") var embeddings []EmbeddingData for _, dataElem := range parsedResponse.Data { + embedding := dataElem.Embedding + if len(embedding) == 0 && len(dataElem.Embeddings) > 0 { + dimensions := len(dataElem.Embeddings[0]) + if dimensions == 0 { + return nil, fmt.Errorf("Jina embedding response contains an empty multivector at index %d", dataElem.Index) + } + embedding = make([]float64, dimensions) + for _, vector := range dataElem.Embeddings { + if len(vector) != dimensions { + return nil, fmt.Errorf("Jina embedding response contains inconsistent multivector dimensions at index %d", dataElem.Index) + } + for i, value := range vector { + embedding[i] += value + } + } + for i := range embedding { + embedding[i] /= float64(len(dataElem.Embeddings)) + } + } + if len(embedding) == 0 { + return nil, fmt.Errorf("Jina embedding response contains an empty vector at index %d", dataElem.Index) + } embeddings = append(embeddings, EmbeddingData{ - Embedding: dataElem.Embedding, + Embedding: embedding, Index: dataElem.Index, }) } @@ -227,6 +316,9 @@ func (j *JinaModel) Rerank(ctx context.Context, modelName *string, query string, if len(documents) == 0 { return &RerankResponse{}, nil } + if modelName == nil || strings.TrimSpace(*modelName) == "" { + return nil, fmt.Errorf("model name is required") + } resolvedBaseURL, err := j.baseModel.GetBaseURL(apiConfig) if err != nil { @@ -251,7 +343,10 @@ func (j *JinaModel) Rerank(ctx context.Context, modelName *string, query string, return nil, fmt.Errorf("failed to marshal request: %w", err) } - req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData)) + ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData)) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) } @@ -274,16 +369,15 @@ func (j *JinaModel) Rerank(ctx context.Context, modelName *string, query string, return nil, fmt.Errorf("Jina Rerank API error: status %d, body: %s", resp.StatusCode, string(body)) } - var rerankResp struct { - Results []struct { - Index int `json:"index"` - RelevanceScore float64 `json:"relevance_score"` - } `json:"results"` - } + var rerankResp JinaRerankResponse if err = json.Unmarshal(body, &rerankResp); err != nil { return nil, fmt.Errorf("failed to decode response: %w", err) } + recordResponseUsage(modelUsage, rerankResp.ID, &TokenUsage{ + PromptTokens: rerankResp.Usage.PromptTokens, + TotalTokens: rerankResp.Usage.TotalTokens, + }, "rerank") var rerankResponse RerankResponse for _, result := range rerankResp.Results { diff --git a/internal/entity/models/jina_test.go b/internal/entity/models/jina_test.go index 3dd4418b08..24084a1790 100644 --- a/internal/entity/models/jina_test.go +++ b/internal/entity/models/jina_test.go @@ -90,6 +90,38 @@ func TestJinaChatHappyPath(t *testing.T) { } } +func TestJinaChatPreservesReasoningContent(t *testing.T) { + srv := newJinaServer(t, "/chat/completions", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "id": "jina-chat", + "choices": []map[string]interface{}{{ + "message": map[string]interface{}{ + "content": "answer", + "reasoning_content": "\nthought", + }, + }}, + }) + }) + defer srv.Close() + + apiKey := "test-key" + thinking := true + response, err := newJinaForTest(srv.URL).ChatWithMessages( + t.Context(), + "jina-vlm", + []Message{{Role: "user", Content: "ping"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{Thinking: &thinking}, + nil, + ) + if err != nil { + t.Fatalf("ChatWithMessages: %v", err) + } + if response.ReasonContent == nil || *response.ReasonContent != "thought" { + t.Fatalf("ReasonContent=%v, want thought", response.ReasonContent) + } +} + func TestJinaChatSupportsToolCalls(t *testing.T) { testNonStreamingToolCall(t, "jina-vlm", "/chat/completions", func(baseURL string) ModelDriver { return newJinaForTest(baseURL) @@ -191,6 +223,51 @@ func TestJinaChatValidation(t *testing.T) { } } +func TestJinaChatStreamIsNotSupported(t *testing.T) { + apiKey := "test-key" + err := newJinaForTest("http://unused").ChatStreamlyWithSender( + t.Context(), + "jina-vlm", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, + nil, + nil, + func(*string, *string) error { return nil }, + ) + if err == nil || !strings.Contains(err.Error(), "ChatStreamlyWithSender") { + t.Fatalf("expected unsupported streaming error, got %v", err) + } +} + +func TestJinaEmbedMeanPoolsMultivectorResponse(t *testing.T) { + srv := newJinaServer(t, "/embeddings", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "data": []map[string]interface{}{{ + "embeddings": [][]float64{{1, 3}, {3, 5}}, + "index": 0, + }}, + }) + }) + defer srv.Close() + + apiKey := "test-key" + modelName := "jina-embeddings-v4" + embeddings, err := newJinaForTest(srv.URL).Embed( + t.Context(), + &modelName, + []string{"text"}, + &APIConfig{ApiKey: &apiKey}, + nil, + nil, + ) + if err != nil { + t.Fatalf("Embed: %v", err) + } + if len(embeddings) != 1 || len(embeddings[0].Embedding) != 2 || embeddings[0].Embedding[0] != 2 || embeddings[0].Embedding[1] != 4 { + t.Fatalf("embeddings=%v, want [[2 4]]", embeddings) + } +} + func TestJinaChatRejectsHTTPError(t *testing.T) { ctx := t.Context() srv := newJinaServer(t, "/chat/completions", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) { diff --git a/internal/entity/models/non_chat_usage_test.go b/internal/entity/models/non_chat_usage_test.go index 3eb88c6ff9..e260c20012 100644 --- a/internal/entity/models/non_chat_usage_test.go +++ b/internal/entity/models/non_chat_usage_test.go @@ -60,6 +60,71 @@ func TestProviderEmbeddingAndRerankUsage(t *testing.T) { embeddingBody: `{"id":"embed-volc","data":{"embedding":[0.1,0.2]},"usage":{"prompt_tokens":7,"total_tokens":7}}`, embedInput: 7, embedTotal: 7, }, + { + name: "jina", + newModel: func(baseURL string) ModelDriver { + return NewJinaModel(map[string]string{"default": baseURL}, URLSuffix{Embedding: "embedding", Rerank: "rerank"}) + }, + embeddingBody: `{"id":"embed-jina","data":[{"embedding":[0.1,0.2],"index":0}],"usage":{"prompt_tokens":7,"total_tokens":7}}`, + rerankBody: `{"id":"rerank-jina","results":[{"index":0,"relevance_score":0.9}],"usage":{"prompt_tokens":7,"total_tokens":9}}`, + embedInput: 7, + embedTotal: 7, + rerankInput: 7, + rerankOutput: 0, + rerankTotal: 9, + supportsRerank: true, + }, + { + name: "gitee", + newModel: func(baseURL string) ModelDriver { + return NewGiteeModel(map[string]string{"default": baseURL}, URLSuffix{Embedding: "embedding", Rerank: "rerank"}) + }, + embeddingBody: `{"id":"embed-gitee","data":[{"embedding":[0.1,0.2],"index":0}],"usage":{"prompt_tokens":7,"total_tokens":7}}`, + rerankBody: `{"id":"rerank-gitee","results":[{"index":0,"relevance_score":0.9}],"usage":{"prompt_tokens":7,"completion_tokens":2,"total_tokens":9}}`, + embedInput: 7, + embedTotal: 7, + rerankInput: 7, + rerankOutput: 2, + rerankTotal: 9, + supportsRerank: true, + }, + { + name: "openrouter", + newModel: func(baseURL string) ModelDriver { + return NewOpenRouterModel(map[string]string{"default": baseURL}, URLSuffix{Embedding: "embedding", Rerank: "rerank"}) + }, + embeddingBody: `{"id":"embed-openrouter","data":[{"embedding":[0.1,0.2],"index":0}],"usage":{"prompt_tokens":7,"total_tokens":7}}`, + rerankBody: `{"id":"rerank-openrouter","results":[{"index":0,"relevance_score":0.9}],"usage":{"prompt_tokens":7,"completion_tokens":2,"total_tokens":9}}`, + embedInput: 7, + embedTotal: 7, + rerankInput: 7, + rerankOutput: 2, + rerankTotal: 9, + supportsRerank: true, + }, + { + name: "jiekouai", + newModel: func(baseURL string) ModelDriver { + return NewJieKouAIModel(map[string]string{"default": baseURL}, URLSuffix{Embedding: "embedding", Rerank: "rerank"}) + }, + embeddingBody: `{"id":"embed-jiekouai","data":[{"embedding":[0.1,0.2],"index":0}],"usage":{"prompt_tokens":7,"total_tokens":7}}`, + rerankBody: `{"id":"rerank-jiekouai","results":[{"index":0,"relevance_score":0.9}],"usage":{"prompt_tokens":7,"completion_tokens":2,"total_tokens":9}}`, + embedInput: 7, + embedTotal: 7, + rerankInput: 7, + rerankOutput: 2, + rerankTotal: 9, + supportsRerank: true, + }, + { + name: "hunyuan", + newModel: func(baseURL string) ModelDriver { + return NewHunyuanModel(map[string]string{"default": baseURL}, URLSuffix{Embedding: "embedding"}) + }, + embeddingBody: `{"id":"embed-hunyuan","data":[{"embedding":[0.1,0.2],"index":0}],"usage":{"prompt_tokens":7,"total_tokens":7}}`, + embedInput: 7, + embedTotal: 7, + }, } for _, tc := range cases { diff --git a/internal/entity/models/openai.go b/internal/entity/models/openai.go index 289d9c84a7..b8166d7341 100644 --- a/internal/entity/models/openai.go +++ b/internal/entity/models/openai.go @@ -341,10 +341,10 @@ func (o *OpenAIModel) ChatStreamlyWithSender(ctx context.Context, modelName stri } type openaiEmbeddingResponse struct { - Data []openrouterEmbeddingData `json:"data"` - Model string `json:"model"` - Object string `json:"object"` - Usage openrouterUsage `json:"usage"` + Data []openaiEmbeddingData `json:"data"` + Model string `json:"model"` + Object string `json:"object"` + Usage openaiUsage `json:"usage"` } type openaiEmbeddingData struct { diff --git a/internal/entity/models/openai_compatible_usage_test.go b/internal/entity/models/openai_compatible_usage_test.go index 9b75555984..79b7476514 100644 --- a/internal/entity/models/openai_compatible_usage_test.go +++ b/internal/entity/models/openai_compatible_usage_test.go @@ -58,6 +58,34 @@ func TestOpenAICompatibleProvidersExtractStreamingUsage(t *testing.T) { }, path: "/chat/completions", }, + { + name: "gitee", + new: func(baseURL string) ModelDriver { + return NewGiteeModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + path: "/chat/completions", + }, + { + name: "openrouter", + new: func(baseURL string) ModelDriver { + return NewOpenRouterModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + path: "/chat/completions", + }, + { + name: "jiekouai", + new: func(baseURL string) ModelDriver { + return NewJieKouAIModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "openai/v1/chat/completions"}) + }, + path: "/openai/v1/chat/completions", + }, + { + name: "hunyuan", + new: func(baseURL string) ModelDriver { + return NewHunyuanModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + path: "/chat/completions", + }, } for _, provider := range providers { diff --git a/internal/entity/models/openrouter.go b/internal/entity/models/openrouter.go index 81ee38b08c..da26203c64 100644 --- a/internal/entity/models/openrouter.go +++ b/internal/entity/models/openrouter.go @@ -54,6 +54,37 @@ func (o *OpenRouterModel) Name() string { return "openrouter" } +// OpenRouterChatResponse mirrors OpenRouter's chat-completions response. +type OpenRouterChatResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []struct { + FinishReason string `json:"finish_reason"` + Index int `json:"index"` + Logprobs any `json:"logprobs"` + Message struct { + Content string `json:"content"` + Reasoning string `json:"reasoning"` + Role string `json:"role"` + ToolCalls []map[string]any `json:"tool_calls"` + } `json:"message"` + } `json:"choices"` + SystemFingerprint string `json:"system_fingerprint"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + CompletionTokensDetails struct { + ReasoningTokens int `json:"reasoning_tokens"` + } `json:"completion_tokens_details"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + } `json:"usage"` +} + func (o *OpenRouterModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := o.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err @@ -108,51 +139,36 @@ func (o *OpenRouterModel) ChatWithMessages(ctx context.Context, modelName string return nil, fmt.Errorf("failed to send request: %d %s", resp.StatusCode, string(body)) } - // Parse response - var result map[string]interface{} - if err := json.Unmarshal(body, &result); err != nil { - return nil, fmt.Errorf("failed to unmarshal response: %w", err) - } - - choices, ok := result["choices"].([]interface{}) - if !ok { - return nil, fmt.Errorf("no choices in response") - } - - firstChoice, ok := choices[0].(map[string]interface{}) - if !ok { - return nil, fmt.Errorf("no choices in response") - } - - messageMap, ok := firstChoice["message"].(map[string]interface{}) - if !ok { - return nil, fmt.Errorf("no message in response") - } - - content, ok := messageMap["content"].(string) - toolCalls := extractToolCalls(messageMap) - if !ok && len(toolCalls) == 0 { - return nil, fmt.Errorf("no message in response") - } - - var reasonContent string - if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { - reasonContent, ok = messageMap["reasoning"].(string) - if !ok { - return nil, fmt.Errorf("invalid content format") + return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) { + var result OpenRouterChatResponse + if err := json.Unmarshal(body, &result); err != nil { + return chatResponseParts{}, fmt.Errorf("failed to unmarshal response: %w", err) } - if reasonContent != "" && reasonContent[0] == '\n' { - reasonContent = reasonContent[1:] + if len(result.Choices) == 0 { + return chatResponseParts{}, fmt.Errorf("no choices in response") } - } - chatResponse := &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - } + choice := result.Choices[0] + if choice.Message.Content == "" && len(choice.Message.ToolCalls) == 0 { + return chatResponseParts{}, fmt.Errorf("response contains neither content nor tool calls") + } + reasonContent := "" + if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking { + reasonContent = strings.TrimPrefix(choice.Message.Reasoning, "\n") + } - return chatResponse, nil + return chatResponseParts{ + RequestID: result.ID, + Content: &choice.Message.Content, + ReasonContent: &reasonContent, + ToolCalls: choice.Message.ToolCalls, + Usage: &TokenUsage{ + PromptTokens: result.Usage.PromptTokens, + CompletionTokens: result.Usage.CompletionTokens, + TotalTokens: result.Usage.TotalTokens, + }, + }, nil + }) } func (o *OpenRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { @@ -163,6 +179,10 @@ func (o *OpenRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName if len(messages) == 0 { return fmt.Errorf("messages is empty") } + if modelConfig != nil { + modelConfig.ToolCallsResult = nil + modelConfig.UsageResult = nil + } resolvedBaseURL, err := o.baseModel.GetBaseURL(apiConfig) if err != nil { @@ -175,6 +195,7 @@ func (o *OpenRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName url := fmt.Sprintf("%s/%s", resolvedBaseURL, o.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(modelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]any{"include_usage": true} if modelConfig != nil { if modelConfig.Thinking != nil { @@ -212,21 +233,31 @@ func (o *OpenRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName return fmt.Errorf("invalid status code: %d, body: %s", resp.StatusCode, string(body)) } + sawTerminal := false accumulatedToolCalls := make(map[int]map[string]any) - if _, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { + done, err := ParseSSEStream[map[string]any](resp.Body, func(event map[string]any) error { common.Info(fmt.Sprintf("%v", event)) - choices, ok := event["choices"].([]interface{}) + tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event) + if usageErr != nil { + return usageErr + } + if found { + applyStreamUsage(modelConfig, modelUsage, tokenUsage) + } + choices, ok := event["choices"].([]any) if !ok || len(choices) == 0 { return nil } - firstChoice, ok := choices[0].(map[string]interface{}) + choice, ok := choices[0].(map[string]any) if !ok { return nil } - - delta, ok := firstChoice["delta"].(map[string]interface{}) + if finishReason, ok := choice["finish_reason"].(string); ok && finishReason != "" { + sawTerminal = true + } + delta, ok := choice["delta"].(map[string]any) if !ok { return nil } @@ -248,9 +279,13 @@ func (o *OpenRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName } return nil - }); err != nil { + }) + if err != nil { return fmt.Errorf("failed to scan response body: %w", err) } + if !done && !sawTerminal { + return fmt.Errorf("openrouter: stream ended before [DONE] or finish_reason") + } setSortedToolCallsResult(modelConfig, accumulatedToolCalls) // Send [DONE] marker for OpenAI compatibility @@ -262,22 +297,21 @@ func (o *OpenRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName return nil } -type openrouterEmbeddingResponse struct { - Data []openrouterEmbeddingData `json:"data"` - Model string `json:"model"` - Object string `json:"object"` - Usage openrouterUsage `json:"usage"` -} - -type openrouterEmbeddingData struct { - Embedding []float64 `json:"embedding"` - Object string `json:"object"` - Index int `json:"index"` -} - -type openrouterUsage struct { - PromptTokens int `json:"prompt_tokens"` - TotalTokens int `json:"total_tokens"` +// OpenRouterEmbeddingResponse mirrors OpenRouter's embeddings response. +type OpenRouterEmbeddingResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Model string `json:"model"` + Data []struct { + Object string `json:"object"` + Embedding []float64 `json:"embedding"` + Index int `json:"index"` + } `json:"data"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` } func (o *OpenRouterModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) { @@ -338,10 +372,15 @@ func (o *OpenRouterModel) Embed(ctx context.Context, modelName *string, texts [] return nil, fmt.Errorf("OpenRouter embedding API error: status %d, body: %s", resp.StatusCode, string(body)) } - var parsed openrouterEmbeddingResponse + var parsed OpenRouterEmbeddingResponse if err = json.Unmarshal(body, &parsed); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } + recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{ + PromptTokens: parsed.Usage.PromptTokens, + CompletionTokens: parsed.Usage.CompletionTokens, + TotalTokens: parsed.Usage.TotalTokens, + }, "embedding") var embeddings []EmbeddingData for _, dataElem := range parsed.Data { @@ -373,6 +412,11 @@ type OpenRouterRerankResponse struct { Text string `json:"text"` } `json:"document,omitempty"` } `json:"results"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` } func (o *OpenRouterModel) Rerank(ctx context.Context, modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig, modelUsage *common.ModelUsage) (*RerankResponse, error) { @@ -383,10 +427,13 @@ func (o *OpenRouterModel) Rerank(ctx context.Context, modelName *string, query s if len(documents) == 0 { return &RerankResponse{}, nil } + if modelName == nil || strings.TrimSpace(*modelName) == "" { + return nil, fmt.Errorf("model name is required") + } - var topN = rerankConfig.TopN - if rerankConfig.TopN == 0 { - topN = len(documents) + topN := len(documents) + if rerankConfig != nil && rerankConfig.TopN > 0 { + topN = rerankConfig.TopN } reqBody := OpenRouterRerankRequest{ @@ -437,6 +484,11 @@ func (o *OpenRouterModel) Rerank(ctx context.Context, modelName *string, query s if err = json.Unmarshal(body, &rerankResp); err != nil { return nil, fmt.Errorf("failed to decode response: %w", err) } + recordResponseUsage(modelUsage, rerankResp.ID, &TokenUsage{ + PromptTokens: rerankResp.Usage.PromptTokens, + CompletionTokens: rerankResp.Usage.CompletionTokens, + TotalTokens: rerankResp.Usage.TotalTokens, + }, "rerank") var rerankResponse RerankResponse for _, result := range rerankResp.Results { diff --git a/internal/entity/models/provider_response_usage_test.go b/internal/entity/models/provider_response_usage_test.go new file mode 100644 index 0000000000..cca2e6b1b6 --- /dev/null +++ b/internal/entity/models/provider_response_usage_test.go @@ -0,0 +1,108 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package models + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" +) + +func TestProviderLocalChatResponsesExposeUsage(t *testing.T) { + type providerCase struct { + name string + path string + new func(string) ModelDriver + } + + cases := []providerCase{ + { + name: "jina", + path: "/chat/completions", + new: func(baseURL string) ModelDriver { + return NewJinaModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + }, + { + name: "gitee", + path: "/chat/completions", + new: func(baseURL string) ModelDriver { + return NewGiteeModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + }, + { + name: "openrouter", + path: "/chat/completions", + new: func(baseURL string) ModelDriver { + return NewOpenRouterModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + }, + { + name: "jiekouai", + path: "/openai/v1/chat/completions", + new: func(baseURL string) ModelDriver { + return NewJieKouAIModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "openai/v1/chat/completions"}) + }, + }, + { + name: "hunyuan", + path: "/chat/completions", + new: func(baseURL string) ModelDriver { + return NewHunyuanModel(map[string]string{"default": baseURL}, URLSuffix{Chat: "chat/completions"}) + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != tc.path { + t.Errorf("path=%q, want %q", r.URL.Path, tc.path) + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "id": "chat-request", + "choices": []map[string]any{{ + "message": map[string]any{"content": "answer"}, + }}, + "usage": map[string]any{ + "prompt_tokens": 3, + "completion_tokens": 5, + "total_tokens": 8, + }, + }) + })) + defer server.Close() + + apiKey := "test-key" + response, err := tc.new(server.URL).ChatWithMessages( + t.Context(), + "model", + []Message{{Role: "user", Content: "hello"}}, + &APIConfig{ApiKey: &apiKey}, + nil, + nil, + ) + if err != nil { + t.Fatalf("ChatWithMessages: %v", err) + } + if response.Usage == nil || response.Usage.PromptTokens != 3 || response.Usage.CompletionTokens != 5 || response.Usage.TotalTokens != 8 { + t.Fatalf("Usage=%#v, want prompt=3 completion=5 total=8", response.Usage) + } + }) + } +}