diff --git a/internal/entity/models/gitee.go b/internal/entity/models/gitee.go index 8ba04c792c..976cdf4427 100644 --- a/internal/entity/models/gitee.go +++ b/internal/entity/models/gitee.go @@ -53,34 +53,6 @@ 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 { @@ -112,81 +84,12 @@ func (g *GiteeModel) ChatWithMessages(ctx context.Context, modelName string, mes } } - jsonData, err := json.Marshal(reqBody) + body, err := g.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) + return nil, err } - common.Info(fmt.Sprintf("GiteeAPI request body: %s", string(jsonData))) - - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) - - resp, err := g.baseModel.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read response: %w", err) - } - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) - } - - 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") - } - - 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") - } - - 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:] - } - } - - 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 - }) + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) @@ -226,128 +129,9 @@ func (g *GiteeModel) ChatStreamlyWithSender(ctx context.Context, modelName strin } - jsonData, err := json.Marshal(reqBody) - if err != nil { - return fmt.Errorf("failed to marshal request: %w", err) - } - - ctx, cancel := context.WithTimeout(ctx, streamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) - - resp, err := g.baseModel.httpClient.Do(req) - if err != nil { - return fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) - } - - reserveText := "" - thinkingPhase := false - answerPhase := false - sawTerminal := false - - accumulatedToolCalls := make(map[int]map[string]any) - 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 - } - choice, ok := choices[0].(map[string]any) - if !ok { - return nil - } - 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 != "" { - common.Info(content) - if content == "" { - thinkingPhase = true - return nil - - } else if content == "" { - thinkingPhase = false - answerPhase = true - return nil - } - - if thinkingPhase { - if err = sender(nil, &content); err != nil { - return err - } - reserveText = "" - } else if answerPhase { - if err = sender(&content, nil); err != nil { - return err - } - reserveText = "" - } else { - content = strings.Trim(content, "\n") - content = strings.Trim(content, " ") - if content != "" { - reserveText += content - } - } - } - - return nil + return g.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) }) - if err != nil { - return fmt.Errorf("failed to scan response body: %w", err) - } - if !done && !sawTerminal { - return fmt.Errorf("gitee: stream ended before [DONE] or finish_reason") - } - - if reserveText != "" { - if err = sender(&reserveText, nil); err != nil { - return err - } - } - - setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls) - - // Send [DONE] marker for OpenAI compatibility - endOfStream := "[DONE]" - if err = sender(&endOfStream, nil); err != nil { - return err - } - - return nil } // GiteeEmbeddingResponse mirrors GiteeAI's embeddings response. diff --git a/internal/entity/models/gitee_test.go b/internal/entity/models/gitee_test.go index 280fefcae4..957e9fd1ff 100644 --- a/internal/entity/models/gitee_test.go +++ b/internal/entity/models/gitee_test.go @@ -57,7 +57,8 @@ func TestGiteeStreamAcceptsTerminalWithoutDelta(t *testing.T) { 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") + _, _ = io.WriteString(w, `data: {"choices":[{"finish_reason":"stop"}]}`+"\n\n"+ + `data: [DONE]`+"\n\n") })) defer srv.Close() diff --git a/internal/entity/models/gpustack.go b/internal/entity/models/gpustack.go index fb87f4afde..cb4bbf598c 100644 --- a/internal/entity/models/gpustack.go +++ b/internal/entity/models/gpustack.go @@ -71,70 +71,12 @@ func (g *GPUStackModel) ChatWithMessages(ctx context.Context, modelName string, url := fmt.Sprintf("%s/%s", baseURL, g.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) - jsonData, err := json.Marshal(reqBody) + body, err := g.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) + return nil, err } - 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) - } - req.Header.Set("Content-Type", "application/json") - if auth := BearerAuth(apiConfig); auth != "" { - req.Header.Set("Authorization", auth) - } - - resp, err := g.baseModel.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read response: %w", err) - } - if resp.StatusCode != http.StatusOK { - 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) - } - - 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 + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender streams the response via the sender. @@ -163,12 +105,16 @@ func (g *GPUStackModel) ChatStreamlyWithSender(ctx context.Context, modelName st baseURL = strings.TrimSuffix(baseURL, "/") url := fmt.Sprintf("%s/%s", baseURL, g.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) } + 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) @@ -190,60 +136,7 @@ func (g *GPUStackModel) ChatStreamlyWithSender(ctx context.Context, modelName st return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - sawTerminal := false - accumulatedToolCalls := make(map[int]map[string]any) - done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { - if apiErr, ok := event["error"]; ok { - return fmt.Errorf("gpustack: upstream stream error: %v", apiErr) - } - - choices, ok := event["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil - } - firstChoice, ok := choices[0].(map[string]interface{}) - if !ok { - return nil - } - delta, ok := firstChoice["delta"].(map[string]interface{}) - if !ok { - return nil - } - - accumulateToolCallDeltas(delta, accumulatedToolCalls) - - if r, ok := delta["reasoning_content"].(string); ok && r != "" { - if err := sender(nil, &r); err != nil { - return err - } - } - - content, ok := delta["content"].(string) - if ok && content != "" { - if err := sender(&content, nil); err != nil { - return err - } - } - - finishReason, ok := firstChoice["finish_reason"].(string) - if ok && finishReason != "" { - sawTerminal = true - } - return nil - }) - if err != nil { - return fmt.Errorf("failed to scan response body: %w", err) - } - setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls) - if !done && !sawTerminal { - return fmt.Errorf("gpustack: stream ended before [DONE] or finish_reason") - } - - endOfStream := "[DONE]" - if err := sender(&endOfStream, nil); err != nil { - return err - } - return nil + return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) } type gpustackModelInfo struct { diff --git a/internal/entity/models/greenpt.go b/internal/entity/models/greenpt.go index baa57244b2..3711d293c5 100644 --- a/internal/entity/models/greenpt.go +++ b/internal/entity/models/greenpt.go @@ -67,6 +67,88 @@ func (m *GreenPTModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([] return models, nil } +// ChatWithMessages sends multiple messages with roles and returns response +func (m *GreenPTModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { + if err := m.baseModel.APIConfigCheck(apiConfig); err != nil { + return nil, err + } + + if len(messages) == 0 { + return nil, fmt.Errorf("messages is empty") + } + + resolvedBaseURL, err := m.baseModel.GetBaseURL(apiConfig) + if err != nil { + return nil, err + } + url := fmt.Sprintf("%s/%s", resolvedBaseURL, m.baseModel.URLSuffix.Chat) + + // Build request body + reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) + + if chatModelConfig != nil && chatModelConfig.Thinking != nil { + if *chatModelConfig.Thinking { + reqBody["thinking"] = map[string]interface{}{ + "type": "enabled", + } + } else { + reqBody["thinking"] = map[string]interface{}{ + "type": "disabled", + } + } + } + + body, err := m.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) + if err != nil { + return nil, err + } + + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) +} + +// ChatStreamlyWithSender sends messages and streams response via sender function +func (m *GreenPTModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { + if err := m.baseModel.APIConfigCheck(apiConfig); err != nil { + return err + } + + if len(messages) == 0 { + return fmt.Errorf("messages is empty") + } + if chatModelConfig != nil { + chatModelConfig.ToolCallsResult = nil + chatModelConfig.UsageResult = nil + } + + resolvedBaseURL, err := m.baseModel.GetBaseURL(apiConfig) + if err != nil { + return err + } + url := fmt.Sprintf("%s/%s", resolvedBaseURL, m.baseModel.URLSuffix.Chat) + + // Build request body with streaming enabled + reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]interface{}{ + "include_usage": true, + } + + if chatModelConfig != nil && chatModelConfig.Thinking != nil { + if *chatModelConfig.Thinking { + reqBody["thinking"] = map[string]interface{}{ + "type": "enabled", + } + } else { + reqBody["thinking"] = map[string]interface{}{ + "type": "disabled", + } + } + } + + return m.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + }) +} + // TranscribeAudio calls GreenPT's Deepgram-compatible /v1/listen endpoint. func (m *GreenPTModel) TranscribeAudio(ctx context.Context, modelName *string, file *string, apiConfig *APIConfig, asrConfig *ASRConfig, modelUsage *common.ModelUsage) (*ASRResponse, error) { if err := m.baseModel.APIConfigCheck(apiConfig); err != nil { diff --git a/internal/entity/models/huaweicloud.go b/internal/entity/models/huaweicloud.go index a032463c45..4368348095 100644 --- a/internal/entity/models/huaweicloud.go +++ b/internal/entity/models/huaweicloud.go @@ -50,37 +50,6 @@ func (h *HuaweiCloudModel) Name() string { return "huaweicloud" } -// HuaweiCloudChatResponse captures the OpenAI-compatible fields consumed by RAGFlow. -type HuaweiCloudChatResponse 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"` - 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 huaweiCloudRegion(api *APIConfig) string { region := "default" if api != nil && api.Region != nil && *api.Region != "" { @@ -222,73 +191,12 @@ func (h *HuaweiCloudModel) ChatWithMessages(ctx context.Context, modelName strin huaweiCloudApplyChatConfig(reqb, modelName, chatModelConfig) applyChatToolConfig(reqb, chatModelConfig) - jsonData, err := json.Marshal(reqb) + body, err := h.baseModel.doRequest(ctx, url, apiConfig, reqb, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) + return nil, err } - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", huaweiCloudAuthorization(*apiConfig.ApiKey)) - - rep, err := h.baseModel.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer rep.Body.Close() - - body, err := io.ReadAll(rep.Body) - if err != nil { - return nil, fmt.Errorf("failed to read response: %w", err) - } - if rep.StatusCode != http.StatusOK { - return nil, fmt.Errorf("Huawei Cloud chat API error: status %d, body: %s", rep.StatusCode, string(body)) - } - - return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) { - var result HuaweiCloudChatResponse - 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") - } - - choice := &result.Choices[0] - content := choice.Message.Content - if content == "" && len(choice.Message.ToolCalls) == 0 { - return chatResponseParts{}, fmt.Errorf("invalid content format") - } - reasonContent := choice.Message.ReasoningContent - if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking && reasonContent == "" { - reasoning, answer := GetThinkingAndAnswer(chatConfig.ModelClass, &content) - if reasoning != nil { - reasonContent = *reasoning - content = *answer - } - } - if reasonContent != "" && reasonContent[0] == '\n' { - reasonContent = reasonContent[1:] - } - - 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 - }) + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (h *HuaweiCloudModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { @@ -325,93 +233,13 @@ func (h *HuaweiCloudModel) ChatStreamlyWithSender(ctx context.Context, modelName huaweiCloudApplyChatConfig(reqBody, modelName, chatModelConfig) if chatModelConfig != nil { chatModelConfig.ToolCallsResult = nil + chatModelConfig.UsageResult = nil } + reqBody["stream_options"] = map[string]interface{}{"include_usage": true} - jsonData, err := json.Marshal(reqBody) - if err != nil { - return fmt.Errorf("failed to marshal request: %w", err) - } - - 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) - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", huaweiCloudAuthorization(*apiConfig.ApiKey)) - - resp, err := h.baseModel.httpClient.Do(req) - if err != nil { - return fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - return fmt.Errorf("Huawei Cloud stream API error: status %d, body: %s", resp.StatusCode, string(body)) - } - - sawTerminal := false - accumulatedToolCalls := make(map[int]map[string]interface{}) - done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { - if apiErr, ok := event["error"]; ok { - return fmt.Errorf("huaweicloud: upstream stream error: %v", apiErr) - } - - tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event) - if usageErr != nil { - return usageErr - } - if found { - applyStreamUsage(chatModelConfig, modelUsage, tokenUsage) - } - - choices, ok := event["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil - } - firstChoice, ok := choices[0].(map[string]interface{}) - if !ok { - return nil - } - if finishReason, ok := firstChoice["finish_reason"].(string); ok && finishReason != "" { - sawTerminal = true - } - delta, ok := firstChoice["delta"].(map[string]interface{}) - if !ok { - return nil - } - if accumulateToolCallDeltas(delta, accumulatedToolCalls) { - return nil - } - - if r, ok := delta["reasoning_content"].(string); ok && r != "" { - if err := sender(nil, &r); err != nil { - return err - } - } - if content, ok := delta["content"].(string); ok && content != "" { - if err := sender(&content, nil); err != nil { - return err - } - } - return nil + return h.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) }) - if err != nil { - return fmt.Errorf("failed to scan response body: %w", err) - } - if !done && !sawTerminal { - return fmt.Errorf("huaweicloud: stream ended before [DONE] or finish_reason") - } - setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls) - - endOfStream := "[DONE]" - if err := sender(&endOfStream, nil); err != nil { - return err - } - return nil } type huaweiCloudEmbeddingResponse struct { diff --git a/internal/entity/models/huggingface.go b/internal/entity/models/huggingface.go index bc3d5145be..974fb2b903 100644 --- a/internal/entity/models/huggingface.go +++ b/internal/entity/models/huggingface.go @@ -66,7 +66,6 @@ func (h *HuggingFaceModel) ChatWithMessages(ctx context.Context, modelName strin reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) if chatModelConfig != nil { - if chatModelConfig.Thinking != nil { if *chatModelConfig.Thinking { reqBody["thinking"] = map[string]interface{}{ @@ -80,83 +79,12 @@ func (h *HuggingFaceModel) ChatWithMessages(ctx context.Context, modelName strin } } - jsonData, err := json.Marshal(reqBody) + body, err := h.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) + return nil, err } - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) - - resp, err := h.baseModel.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read response: %w", err) - } - - if resp.StatusCode != http.StatusOK { - 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") - } - // if first char of reasonContent is \n remove the \n - if reasonContent != "" && reasonContent[0] == '\n' { - reasonContent = reasonContent[1:] - } - } - - chatResponse := &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - } - - return chatResponse, nil + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (h *HuggingFaceModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { @@ -187,77 +115,11 @@ func (h *HuggingFaceModel) ChatStreamlyWithSender(ctx context.Context, modelName } } - jsonData, err := json.Marshal(reqBody) - if err != nil { - return fmt.Errorf("failed to marshal request: %w", err) - } + reqBody["stream_options"] = map[string]interface{}{"include_usage": true} - ctx, cancel := context.WithTimeout(ctx, streamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) - - resp, err := h.baseModel.httpClient.Do(req) - if err != nil { - return fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) - } - - accumulatedToolCalls := make(map[int]map[string]any) - if _, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { - common.Info(fmt.Sprintf("%v", event)) - - choices, ok := event["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil - } - - firstChoice, ok := choices[0].(map[string]interface{}) - if !ok { - return nil - } - - delta, ok := firstChoice["delta"].(map[string]interface{}) - if !ok { - return nil - } - - accumulateToolCallDeltas(delta, accumulatedToolCalls) - - reasoningContent, ok := delta["reasoning_content"].(string) - if ok && reasoningContent != "" { - if err := sender(nil, &reasoningContent); err != nil { - return err - } - } - - content, ok := delta["content"].(string) - if ok && content != "" { - if err := sender(&content, nil); err != nil { - return err - } - } - - return nil - }); err != nil { - return fmt.Errorf("failed to scan response body: %w", err) - } - setSortedToolCallsResult(modelConfig, accumulatedToolCalls) - - // Send [DONE] marker for OpenAI compatibility - endOfStream := "[DONE]" - return sender(&endOfStream, nil) + return h.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender) + }) } func (h *HuggingFaceModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) { diff --git a/internal/entity/models/hunyuan.go b/internal/entity/models/hunyuan.go index 33d83c2da2..085ac55568 100644 --- a/internal/entity/models/hunyuan.go +++ b/internal/entity/models/hunyuan.go @@ -51,32 +51,6 @@ 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"` @@ -110,64 +84,12 @@ func (h *HunyuanModel) ChatWithMessages(ctx context.Context, modelName string, m url := fmt.Sprintf("%s/%s", baseURL, h.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) - jsonData, err := json.Marshal(reqBody) + body, err := h.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) + return nil, err } - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) - - resp, err := h.baseModel.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read response: %w", err) - } - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) - } - - 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") - } - - 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 - }) + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender opens the SSE chat-completions @@ -199,94 +121,9 @@ func (h *HunyuanModel) ChatStreamlyWithSender(ctx context.Context, modelName str 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) - } - - 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) - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) - - resp, err := h.baseModel.httpClient.Do(req) - if err != nil { - return fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) - } - - sawTerminal := false - accumulatedToolCalls := make(map[int]map[string]any) - 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) - } - 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 - } - - choice, ok := choices[0].(map[string]any) - if !ok { - return nil - } - 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 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 + return h.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) }) - if err != nil { - return fmt.Errorf("failed to scan response body: %w", err) - } - setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls) - if !done && !sawTerminal { - return fmt.Errorf("hunyuan: stream ended before [DONE] or finish_reason") - } - - endOfStream := "[DONE]" - if err := sender(&endOfStream, nil); err != nil { - return err - } - return nil } func (h *HunyuanModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([]ListModelResponse, error) { diff --git a/internal/entity/models/hunyuan_test.go b/internal/entity/models/hunyuan_test.go index dc3ba54e86..cdd2de77df 100644 --- a/internal/entity/models/hunyuan_test.go +++ b/internal/entity/models/hunyuan_test.go @@ -360,7 +360,8 @@ func TestHunyuanStreamFailsWithoutTerminal(t *testing.T) { func TestHunyuanStreamAcceptsTerminalWithoutDelta(t *testing.T) { withSSRFBypass(t) srv := newHunyuanSSEServer(t, "/chat/completions", - `data: {"choices":[{"finish_reason":"stop"}]}`+"\n\n", + `data: {"choices":[{"finish_reason":"stop"}]}`+"\n\n"+ + `data: [DONE]`+"\n\n", ) defer srv.Close() diff --git a/internal/entity/models/jiekouai.go b/internal/entity/models/jiekouai.go index 156f69c99e..8e6b035a67 100644 --- a/internal/entity/models/jiekouai.go +++ b/internal/entity/models/jiekouai.go @@ -62,32 +62,6 @@ 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"` @@ -127,7 +101,6 @@ func (j *JieKouAIModel) ChatWithMessages(ctx context.Context, modelName string, if err := j.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err } - apiKey := strings.TrimSpace(*apiConfig.ApiKey) if modelName = strings.TrimSpace(modelName); modelName == "" { return nil, fmt.Errorf("model name is required") } @@ -151,78 +124,18 @@ func (j *JieKouAIModel) ChatWithMessages(ctx context.Context, modelName string, } } - jsonData, err := json.Marshal(reqBody) + body, err := j.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) + return nil, err } - 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) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) - req.Header.Set("Accept", "application/json") - - resp, err := j.baseModel.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read response: %w", err) - } - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) - } - - 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 len(result.Choices) == 0 { - return chatResponseParts{}, fmt.Errorf("no choices in response") - } - - 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 - }) + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (j *JieKouAIModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { if err := j.baseModel.APIConfigCheck(apiConfig); err != nil { return err } - apiKey := strings.TrimSpace(*apiConfig.ApiKey) if modelName = strings.TrimSpace(modelName); modelName == "" { return fmt.Errorf("model name is required") } @@ -253,96 +166,9 @@ func (j *JieKouAIModel) ChatStreamlyWithSender(ctx context.Context, modelName st } } - jsonData, err := json.Marshal(reqBody) - if err != nil { - return fmt.Errorf("failed to marshal request: %w", err) - } - - 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) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) - req.Header.Set("Accept", "text/event-stream") - - resp, err := j.baseModel.httpClient.Do(req) - if err != nil { - return fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) - } - - sawTerminal := false - accumulatedToolCalls := make(map[int]map[string]any) - 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 - } - - choice, ok := choices[0].(map[string]any) - if !ok { - return nil - } - 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) - - reasoningContent, ok := delta["reasoning_content"].(string) - if ok && reasoningContent != "" { - if err := sender(nil, &reasoningContent); err != nil { - return err - } - } - - content, ok := delta["content"].(string) - if ok && content != "" { - if err := sender(&content, nil); err != nil { - return err - } - } - - return nil + return j.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender) }) - 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 - endOfStream := "[DONE]" - if err = sender(&endOfStream, nil); err != nil { - return err - } - - return nil } func (j *JieKouAIModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) { diff --git a/internal/entity/models/jina.go b/internal/entity/models/jina.go index 241e27e909..f418ffdbb5 100644 --- a/internal/entity/models/jina.go +++ b/internal/entity/models/jina.go @@ -55,35 +55,6 @@ 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 { @@ -139,72 +110,55 @@ func (j *JinaModel) ChatWithMessages(ctx context.Context, modelName string, mess reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) - jsonData, err := json.Marshal(reqBody) - if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) - } - - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) - - resp, err := j.baseModel.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read response: %w", err) - } - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("Jina chat API error: status %d, body: %s", resp.StatusCode, string(body)) - } - - 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") + if chatModelConfig != nil { + if chatModelConfig.Thinking != nil { + reqBody["enable_thinking"] = *chatModelConfig.Thinking } + } - 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 - }) + body, err := j.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) + if err != nil { + return nil, err + } + + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } -func (j *JinaModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { - // Jina's public API does not expose a streaming chat-completions endpoint. - return fmt.Errorf("jina does not implement ChatStreamlyWithSender(not available for now)") +func (j *JinaModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { + if err := j.baseModel.APIConfigCheck(apiConfig); err != nil { + return err + } + if modelName == "" { + return fmt.Errorf("model name is required") + } + if len(messages) == 0 { + return fmt.Errorf("messages is empty") + } + if err := validateStreamConfig(chatModelConfig); err != nil { + return err + } + + baseURL, err := j.baseModel.GetBaseURL(apiConfig) + if err != nil { + return err + } + baseURL = strings.TrimSuffix(baseURL, "/") + url := fmt.Sprintf("%s/%s", baseURL, j.baseModel.URLSuffix.Chat) + + reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]interface{}{"include_usage": true} + + if chatModelConfig != nil { + chatModelConfig.ToolCallsResult = nil + chatModelConfig.UsageResult = nil + if chatModelConfig.Thinking != nil { + reqBody["enable_thinking"] = *chatModelConfig.Thinking + } + } + + return j.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + }) } func (j *JinaModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) { diff --git a/internal/entity/models/jina_test.go b/internal/entity/models/jina_test.go index f2648c50d7..9528ddc8cf 100644 --- a/internal/entity/models/jina_test.go +++ b/internal/entity/models/jina_test.go @@ -227,23 +227,6 @@ func TestJinaChatValidation(t *testing.T) { } } -func TestJinaChatStreamIsNotSupported(t *testing.T) { - withSSRFBypass(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) { withSSRFBypass(t) srv := newJinaServer(t, "/embeddings", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) { diff --git a/internal/entity/models/lmstudio.go b/internal/entity/models/lmstudio.go index 29b0d5ffaa..af815a0678 100644 --- a/internal/entity/models/lmstudio.go +++ b/internal/entity/models/lmstudio.go @@ -91,84 +91,12 @@ func (l *LmStudioModel) ChatWithMessages(ctx context.Context, modelName string, } } - jsonData, err := json.Marshal(reqBody) + body, err := l.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) + return nil, err } - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - if auth := BearerAuth(apiConfig); auth != "" { - req.Header.Set("Authorization", auth) - } - - resp, err := l.baseModel.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read response: %w", err) - } - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("API request failed with status %d: %s :%s", resp.StatusCode, string(body), messages[0].Content) - } - - // 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") - } - if reasonContent != "" && reasonContent[0] == '\n' { - reasonContent = reasonContent[1:] - } - } - - chatResponse := &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - } - - return chatResponse, nil + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) @@ -205,84 +133,11 @@ func (l *LmStudioModel) ChatStreamlyWithSender(ctx context.Context, modelName st } } - jsonData, err := json.Marshal(reqBody) - if err != nil { - return fmt.Errorf("failed to marshal request: %w", err) - } + reqBody["stream_options"] = map[string]any{"include_usage": true} - ctx, cancel := context.WithTimeout(ctx, streamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - if auth := BearerAuth(apiConfig); auth != "" { - req.Header.Set("Authorization", auth) - } - - resp, err := l.baseModel.httpClient.Do(req) - if err != nil { - return fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) - } - - // SSE parsing: read line by line - accumulatedToolCalls := make(map[int]map[string]any) - if _, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { - common.Info(fmt.Sprintf("%v", event)) - - choices, ok := event["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil - } - - firstChoice, ok := choices[0].(map[string]interface{}) - if !ok { - return nil - } - - delta, ok := firstChoice["delta"].(map[string]interface{}) - if !ok { - return nil - } - - accumulateToolCallDeltas(delta, accumulatedToolCalls) - - reasoningContent, ok := delta["reasoning_content"].(string) - if ok && reasoningContent != "" { - if err := sender(nil, &reasoningContent); err != nil { - return err - } - } - - content, ok := delta["content"].(string) - if ok && content != "" { - if err := sender(&content, nil); err != nil { - return err - } - } - - return nil - }); err != nil { - return fmt.Errorf("failed to scan response body: %w", err) - } - setSortedToolCallsResult(modelConfig, accumulatedToolCalls) - - // Send [DONE] marker for OpenAI compatibility - endOfStream := "[DONE]" - if err = sender(&endOfStream, nil); err != nil { - return err - } - - return nil + return l.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender) + }) } func (l *LmStudioModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) { diff --git a/internal/entity/models/localai.go b/internal/entity/models/localai.go index a2e16e4800..442a67d31d 100644 --- a/internal/entity/models/localai.go +++ b/internal/entity/models/localai.go @@ -25,13 +25,8 @@ import ( "net/http" "ragflow/internal/common" "strings" - "sync" - "time" ) -// localAIStreamIdleTimeout bounds how long ChatStreamlyWithSender -var localAIStreamIdleTimeout = 60 * time.Second - // LocalAIModel implements ModelDriver for LocalAI type LocalAIModel struct { baseModel BaseModel @@ -57,17 +52,6 @@ func (l *LocalAIModel) Name() string { return "LocalAI" } -var localAIReasoningFields = []string{"reasoning_content", "reasoning", "thinking"} - -func extractLocalAIReasoning(m map[string]interface{}) string { - for _, k := range localAIReasoningFields { - if v, ok := m[k].(string); ok && v != "" { - return v - } - } - return "" -} - func addLocalAIReasoningRequestParams(reqBody map[string]interface{}, cfg *ChatConfig) { if cfg == nil { return @@ -100,74 +84,12 @@ func (l *LocalAIModel) ChatWithMessages(ctx context.Context, modelName string, m addLocalAIReasoningRequestParams(reqBody, chatModelConfig) } - jsonData, err := json.Marshal(reqBody) + body, err := l.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) + return nil, err } - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - if auth := BearerAuth(apiConfig); auth != "" { - req.Header.Set("Authorization", auth) - } - - resp, err := l.baseModel.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read response: %w", err) - } - - if resp.StatusCode != http.StatusOK { - 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) - } - - 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") - } - - // Pull the chain-of-thought from whichever field the upstream model - // used. See localAIReasoningFields for the priority order. - reasonContent := extractLocalAIReasoning(messageMap) - - return &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - }, nil + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams the response via the @@ -197,118 +119,13 @@ func (l *LocalAIModel) ChatStreamlyWithSender(ctx context.Context, modelName str addLocalAIReasoningRequestParams(reqBody, chatModelConfig) } - jsonData, err := json.Marshal(reqBody) - if err != nil { - return fmt.Errorf("failed to marshal request: %w", err) + reqBody["stream_options"] = map[string]interface{}{ + "include_usage": true, } - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - if auth := BearerAuth(apiConfig); auth != "" { - req.Header.Set("Authorization", auth) - } - - resp, err := l.baseModel.httpClient.Do(req) - if err != nil { - return fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - body, _ := io.ReadAll(resp.Body) - return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) - } - - lastActive := time.Now() - var lastActiveMu sync.Mutex - done := make(chan struct{}) - defer close(done) - go func() { - ticker := time.NewTicker(localAIStreamIdleTimeout / 4) - defer ticker.Stop() - for { - select { - case <-done: - return - case now := <-ticker.C: - lastActiveMu.Lock() - idle := now.Sub(lastActive) - lastActiveMu.Unlock() - if idle >= localAIStreamIdleTimeout { - cancel() - return - } - } - } - }() - - sawTerminal := false - accumulatedToolCalls := make(map[int]map[string]any) - streamDone, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { - lastActiveMu.Lock() - lastActive = time.Now() - lastActiveMu.Unlock() - - choices, ok := event["choices"].([]interface{}) - if !ok || len(choices) == 0 { - return nil - } - - firstChoice, ok := choices[0].(map[string]interface{}) - if !ok { - return nil - } - - delta, ok := firstChoice["delta"].(map[string]interface{}) - if !ok { - return nil - } - - accumulateToolCallDeltas(delta, accumulatedToolCalls) - - if reasoning := extractLocalAIReasoning(delta); reasoning != "" { - if err := sender(nil, &reasoning); err != nil { - return err - } - } - - content, ok := delta["content"].(string) - if ok && content != "" { - if err := sender(&content, nil); err != nil { - return err - } - } - - finishReason, ok := firstChoice["finish_reason"].(string) - if ok && finishReason != "" { - sawTerminal = true - } - return nil + return l.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) }) - if err != nil { - if ctx.Err() != nil { - return fmt.Errorf("localai: stream idle for more than %s, aborted", localAIStreamIdleTimeout) - } - return fmt.Errorf("failed to scan response body: %w", err) - } - setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls) - if !streamDone && !sawTerminal { - return fmt.Errorf("localai: stream ended before [DONE] or finish_reason") - } - - endOfStream := "[DONE]" - if err := sender(&endOfStream, nil); err != nil { - return err - } - - return nil } type localAIEmbeddingData struct { diff --git a/internal/entity/models/localai_test.go b/internal/entity/models/localai_test.go index b4d03d6544..48ea8af3e5 100644 --- a/internal/entity/models/localai_test.go +++ b/internal/entity/models/localai_test.go @@ -6,9 +6,7 @@ import ( "net/http" "net/http/httptest" "strings" - "sync" "testing" - "time" ) func newLocalAIForTest(baseURL string) *LocalAIModel { @@ -23,18 +21,6 @@ func newLocalAIForTest(baseURL string) *LocalAIModel { ) } -// withLocalAIIdleTimeout swaps the package-level idle timeout for the -// duration of the test. Tests that exercise the stall watchdog use a -// sub-second value so they finish quickly. -func withLocalAIIdleTimeout(t *testing.T, d time.Duration) { - t.Helper() - original := localAIStreamIdleTimeout - localAIStreamIdleTimeout = d - t.Cleanup(func() { - localAIStreamIdleTimeout = original - }) -} - func TestLocalAIName(t *testing.T) { l := newLocalAIForTest("http://unused") if got := l.Name(); got != "LocalAI" { @@ -42,114 +28,6 @@ func TestLocalAIName(t *testing.T) { } } -func TestLocalAIStreamCancelsOnIdle(t *testing.T) { - withSSRFBypass(t) - ctx := t.Context() - // The server emits one valid chunk and then stalls. Without the - // watchdog, scanner.Scan() would hang forever. With the watchdog - // at 200ms, it must return a clear "stream idle" error in well - // under a second. - withLocalAIIdleTimeout(t, 200*time.Millisecond) - - // hold blocks the handler until the test closes it. Register - // close(hold) FIRST so it runs LAST (defers are LIFO) — wait, - // that's the opposite. We want close(hold) to run BEFORE - // srv.Close() so the handler can return. Use t.Cleanup, which - // runs in reverse-registration order: register srv.Close first - // so it runs last, then close(hold) so it runs first. - hold := make(chan struct{}) - - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/event-stream") - w.WriteHeader(http.StatusOK) - if f, ok := w.(http.Flusher); ok { - _, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"hi"}}]}`+"\n") - f.Flush() - } - // Hang until either the client disconnects (watchdog cancels - // the request, which causes r.Context() to fire) or the test - // teardown signals via `hold`. - select { - case <-hold: - case <-r.Context().Done(): - } - })) - t.Cleanup(srv.Close) - t.Cleanup(func() { close(hold) }) - - l := newLocalAIForTest(srv.URL) - var got []string - var mu sync.Mutex - err := l.ChatStreamlyWithSender(ctx, "gpt-4", - []Message{{Role: "user", Content: "x"}}, - &APIConfig{}, nil, nil, - func(content *string, _ *string) error { - if content == nil || *content == "" { - return nil - } - mu.Lock() - got = append(got, *content) - mu.Unlock() - return nil - }, - ) - - if err == nil { - t.Fatal("expected an idle-timeout error, got nil") - } - if !strings.Contains(err.Error(), "idle for more than") { - t.Errorf("expected idle-timeout error, got %v", err) - } - mu.Lock() - defer mu.Unlock() - if len(got) == 0 || got[0] != "hi" { - t.Errorf("expected first chunk before stall, got %v", got) - } -} - -func TestLocalAIStreamCompletesWithoutTriggeringWatchdog(t *testing.T) { - withSSRFBypass(t) - ctx := t.Context() - // Sanity check: a fast, complete stream should not trip the - // watchdog even with a moderately tight idle window. - withLocalAIIdleTimeout(t, 500*time.Millisecond) - - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Type", "text/event-stream") - w.WriteHeader(http.StatusOK) - f, _ := w.(http.Flusher) - _, _ = io.WriteString(w, - `data: {"choices":[{"delta":{"content":"a"}}]}`+"\n"+ - `data: {"choices":[{"delta":{"content":"b"}}]}`+"\n"+ - `data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`+"\n"+ - `data: [DONE]`+"\n", - ) - if f != nil { - f.Flush() - } - })) - defer srv.Close() - - l := newLocalAIForTest(srv.URL) - var chunks []string - err := l.ChatStreamlyWithSender(ctx, "gpt-4", - []Message{{Role: "user", Content: "x"}}, - &APIConfig{}, nil, nil, - func(content *string, _ *string) error { - if content != nil && *content != "" && *content != "[DONE]" { - chunks = append(chunks, *content) - } - return nil - }, - ) - if err != nil { - t.Fatalf("stream: %v", err) - } - if strings.Join(chunks, "") != "ab" { - t.Errorf("chunks=%v want [a b]", chunks) - } -} - func TestLocalAIStreamRequiresSender(t *testing.T) { withSSRFBypass(t) ctx := t.Context() @@ -335,54 +213,7 @@ func TestLocalAIEmbedEmptyInputShortCircuits(t *testing.T) { } } -// ---------- reasoning extraction (multi-field) ---------- - -// Table-driven unit coverage of the helper. Mirrors the priority order -// reasoning_content > reasoning > thinking declared in -// localAIReasoningFields. New upstream field names can be added by -// extending that slice without touching call sites. -func TestExtractLocalAIReasoning(t *testing.T) { - cases := []struct { - name string - in map[string]interface{} - want string - }{ - {"empty map", map[string]interface{}{}, ""}, - {"reasoning_content wins", map[string]interface{}{ - "reasoning_content": "rc", - "reasoning": "r", - "thinking": "t", - }, "rc"}, - {"reasoning when no reasoning_content", map[string]interface{}{ - "reasoning": "r", - "thinking": "t", - }, "r"}, - {"thinking when only that is set", map[string]interface{}{ - "thinking": "qwen3-thought", - }, "qwen3-thought"}, - {"empty string treated as absent", map[string]interface{}{ - "reasoning_content": "", - "reasoning": "fallback", - }, "fallback"}, - {"non-string ignored", map[string]interface{}{ - "reasoning_content": 42, - "reasoning": "fallback", - }, "fallback"}, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - got := extractLocalAIReasoning(tc.in) - if got != tc.want { - t.Errorf("got=%q want=%q", got, tc.want) - } - }) - } -} - -// Non-streaming chat against an upstream that puts the trace in -// message.reasoning_content (kimi-k2.6, OpenAI o-series, DeepSeek-R1 -// when proxied through OpenAI-shim). The driver must surface it on -// ChatResponse.ReasonContent. +// ---------- reasoning content (message.reasoning_content) ---------- func TestLocalAIChatExtractsReasoningContent(t *testing.T) { withSSRFBypass(t) ctx := t.Context() @@ -411,33 +242,6 @@ func TestLocalAIChatExtractsReasoningContent(t *testing.T) { } } -// Non-streaming chat that uses message.thinking (Qwen3 via Ollama-shim -// inside LocalAI). The driver must surface it on ReasonContent too. -func TestLocalAIChatExtractsThinking(t *testing.T) { - withSSRFBypass(t) - ctx := t.Context() - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - _, _ = io.WriteString(w, `{"choices":[{"message":{ - "role":"assistant", - "content":"12", - "thinking":"Compute 15/100 * 80" - }}]}`) - })) - defer srv.Close() - - l := newLocalAIForTest(srv.URL) - resp, err := l.ChatWithMessages(ctx, "qwen3-32b", - []Message{{Role: "user", Content: "15% of 80?"}}, - &APIConfig{}, nil, nil, - ) - if err != nil { - t.Fatalf("Chat: %v", err) - } - if *resp.ReasonContent != "Compute 15/100 * 80" { - t.Errorf("ReasonContent=%q want %q", *resp.ReasonContent, "Compute 15/100 * 80") - } -} - // Regression net: a response with no reasoning field at all (any // non-reasoning model) must produce empty ReasonContent without // crashing or erroring. @@ -515,46 +319,6 @@ func TestLocalAIStreamExtractsReasoningContentDelta(t *testing.T) { } } -// Streaming chat where the upstream uses delta.thinking (Qwen3 shape). -// The same handler must work. -func TestLocalAIStreamExtractsThinkingDelta(t *testing.T) { - withSSRFBypass(t) - ctx := t.Context() - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Type", "text/event-stream") - w.WriteHeader(http.StatusOK) - _, _ = io.WriteString(w, - `data: {"choices":[{"index":0,"delta":{"thinking":"qwen-trace"}}]}`+"\n"+ - `data: {"choices":[{"index":0,"delta":{"content":"final"},"finish_reason":"stop"}]}`+"\n"+ - `data: [DONE]`+"\n", - ) - })) - defer srv.Close() - - l := newLocalAIForTest(srv.URL) - var got []string - err := l.ChatStreamlyWithSender(ctx, "qwen3-32b", - []Message{{Role: "user", Content: "x"}}, - &APIConfig{}, nil, nil, - func(c *string, r *string) error { - if r != nil && *r != "" { - got = append(got, "R:"+*r) - } - if c != nil && *c != "" && *c != "[DONE]" { - got = append(got, "C:"+*c) - } - return nil - }, - ) - if err != nil { - t.Fatalf("stream: %v", err) - } - want := []string{"R:qwen-trace", "C:final"} - if len(got) != 2 || got[0] != want[0] || got[1] != want[1] { - t.Errorf("seq=%v want %v", got, want) - } -} - // Request-side: ChatConfig.Effort must flow into request body as // reasoning_effort. func TestLocalAIChatPropagatesReasoningEffort(t *testing.T) {