diff --git a/internal/entity/models/base_model.go b/internal/entity/models/base_model.go index 465d3778a4..51a6f67eb1 100644 --- a/internal/entity/models/base_model.go +++ b/internal/entity/models/base_model.go @@ -30,8 +30,6 @@ import ( "sort" "strings" "time" - - "github.com/mitchellh/mapstructure" ) type BaseModel struct { @@ -39,6 +37,11 @@ type BaseModel struct { URLSuffix URLSuffix httpClient *http.Client AllowEmptyAPIKey bool + // authHeader, when non-nil, supplies the (name, value) pair used for + // authentication instead of the default "Authorization: Bearer ". + // Drivers with non-standard auth (e.g. Xiaomi's api-key header, Xunfei's + // spark_api_password bundle) set it in their constructor. + authHeader func(*APIConfig) (string, string) } // chatResponseParts is the provider-normalized result of a non-streaming chat @@ -105,20 +108,6 @@ func collectModelUsage(modelUsage *common.ModelUsage, usage *TokenUsage) error { return clickhouse.GetDriver().CollectModelUsage(modelUsage) } -// decodeOpenAICompatibleStreamUsage extracts aggregate token usage from one -// OpenAI-compatible streaming event. A missing usage field is not an error. -func decodeOpenAICompatibleStreamUsage(event map[string]any) (*TokenUsage, bool, error) { - rawUsage, ok := event["usage"].(map[string]any) - if !ok { - return nil, false, nil - } - usage := &TokenUsage{} - if err := mapstructure.Decode(rawUsage, usage); err != nil { - return nil, false, err - } - return usage, true, nil -} - // applyStreamUsage exposes streamed token usage to the caller and records it // for model-usage analytics when a usage event is received. Analytics failures // are logged but do not interrupt the stream. @@ -149,7 +138,21 @@ func (b *BaseModel) APIConfigCheck(apiConfig *APIConfig) error { return nil } -func newJSONPostRequest(ctx context.Context, url string, apiConfig *APIConfig, reqBody map[string]any) (*http.Request, error) { +// applyAuth sets the authentication header on req. Drivers with a custom +// authHeader hook (e.g. Xiaomi's api-key header, Xunfei's spark_api_password +// bundle) use it; the default is "Authorization: Bearer ". +func (b *BaseModel) applyAuth(req *http.Request, apiConfig *APIConfig) { + if b.authHeader != nil { + name, value := b.authHeader(apiConfig) + req.Header.Set(name, value) + return + } + if auth := BearerAuth(apiConfig); auth != "" { + req.Header.Set("Authorization", auth) + } +} + +func (b *BaseModel) newJSONPostRequest(ctx context.Context, url string, apiConfig *APIConfig, reqBody map[string]any) (*http.Request, error) { jsonData, err := json.Marshal(reqBody) if err != nil { return nil, fmt.Errorf("failed to marshal request: %w", err) @@ -161,9 +164,7 @@ func newJSONPostRequest(ctx context.Context, url string, apiConfig *APIConfig, r } req.Header.Set("Content-Type", "application/json") - if auth := BearerAuth(apiConfig); auth != "" { - req.Header.Set("Authorization", auth) - } + b.applyAuth(req, apiConfig) return req, nil } @@ -173,7 +174,7 @@ func (b *BaseModel) doRequest(ctx context.Context, url string, apiConfig *APICon ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() - req, err := newJSONPostRequest(ctx, url, apiConfig, reqBody) + req, err := b.newJSONPostRequest(ctx, url, apiConfig, reqBody) if err != nil { return nil, err } @@ -196,6 +197,37 @@ func (b *BaseModel) doRequest(ctx context.Context, url string, apiConfig *APICon return body, nil } +// doGetRequest sends a GET request and returns the response body. +func (b *BaseModel) doGetRequest(ctx context.Context, url string, apiConfig *APIConfig, timeout time.Duration) ([]byte, error) { + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, "GET", url, nil) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + b.applyAuth(req, apiConfig) + + resp, err := b.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 body, nil +} + // mustMarshal marshals v to JSON, panicking on error. func mustMarshal(v any) []byte { b, err := json.Marshal(v) @@ -210,7 +242,7 @@ func (b *BaseModel) doStreamRequest(ctx context.Context, url string, apiConfig * ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() - req, err := newJSONPostRequest(ctx, url, apiConfig, reqBody) + req, err := b.newJSONPostRequest(ctx, url, apiConfig, reqBody) if err != nil { return err } diff --git a/internal/entity/models/common.go b/internal/entity/models/common.go index bfc99bc280..c0dbd2607b 100644 --- a/internal/entity/models/common.go +++ b/internal/entity/models/common.go @@ -18,44 +18,6 @@ package models import "strings" -func GetThinkingAndAnswer(modelType *string, content *string) (*string, *string) { - if content == nil { - return nil, nil - } - - switch NormalizeModelFamily(modelType) { - case "qwen3": - return extractThinkContent(content) - } - return nil, content -} - -// NormalizeModelFamily normalizes provider-prefixed model class/name strings for shared response parsing. -func NormalizeModelFamily(modelType *string) string { - if modelType == nil { - return "" - } - - family := strings.ToLower(strings.TrimSpace(*modelType)) - if family == "" { - return "" - } - - if slash := strings.LastIndex(family, "/"); slash >= 0 && slash < len(family)-1 { - family = family[slash+1:] - } - - if family == "qwen3" || strings.HasPrefix(family, "qwen3-") || strings.HasPrefix(family, "qwen3.") { - return "qwen3" - } - - if dash := strings.Index(family, "-"); dash >= 0 { - family = family[:dash] - } - - return family -} - func extractThinkContent(content *string) (*string, *string) { if content == nil { return nil, nil diff --git a/internal/entity/models/common_test.go b/internal/entity/models/common_test.go index 22fbf5b74e..61aafa5324 100644 --- a/internal/entity/models/common_test.go +++ b/internal/entity/models/common_test.go @@ -10,101 +10,6 @@ import ( "time" ) -func TestNormalizeModelFamily(t *testing.T) { - tests := []struct { - name string - input *string - want string - }{ - {name: "nil", input: nil, want: ""}, - {name: "empty", input: modelFamilyTestString(""), want: ""}, - {name: "qwen3", input: modelFamilyTestString("qwen3"), want: "qwen3"}, - {name: "qwen3 hyphen variant", input: modelFamilyTestString("qwen3-8b"), want: "qwen3"}, - {name: "qwen3 dot variant", input: modelFamilyTestString("qwen3.5-4b"), want: "qwen3"}, - {name: "provider-prefixed qwen3", input: modelFamilyTestString("qwen/qwen3-8b"), want: "qwen3"}, - {name: "case-varied qwen3", input: modelFamilyTestString("Qwen/Qwen3.5-4B"), want: "qwen3"}, - {name: "provider-prefixed non-qwen", input: modelFamilyTestString("deepseek/deepseek-r1"), want: "deepseek"}, - {name: "qwen plus not qwen3", input: modelFamilyTestString("qwen-plus"), want: "qwen"}, - {name: "provider-prefixed qwen2.5 not qwen3", input: modelFamilyTestString("qwen/qwen2.5-32b-instruct"), want: "qwen2.5"}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if got := NormalizeModelFamily(tt.input); got != tt.want { - t.Fatalf("NormalizeModelFamily()=%q, want %q", got, tt.want) - } - }) - } -} - -func TestGetThinkingAndAnswerExtractsQwenThinking(t *testing.T) { - content := "\nplan\nanswer" - - for _, modelType := range []string{ - "qwen3", - "qwen3-8b", - "qwen/qwen3", - "qwen/qwen3-8b", - "Qwen/Qwen3.5-4B", - } { - t.Run(modelType, func(t *testing.T) { - thinking, answer := GetThinkingAndAnswer(&modelType, &content) - if thinking == nil || *thinking != "plan" { - t.Fatalf("thinking=%v, want plan", thinking) - } - if answer == nil || *answer != "answer" { - t.Fatalf("answer=%v, want answer", answer) - } - }) - } -} - -func TestGetThinkingAndAnswerLeavesUnknownModelFamiliesUnchanged(t *testing.T) { - content := "\nplan\nanswer" - - for _, modelType := range []string{ - "deepseek", - "deepseek/deepseek-r1", - "qwen-plus", - "qwen/qwen2.5-32b-instruct", - } { - t.Run(modelType, func(t *testing.T) { - thinking, answer := GetThinkingAndAnswer(&modelType, &content) - if thinking != nil { - t.Fatalf("thinking=%v, want nil", thinking) - } - if answer != &content { - t.Fatalf("answer pointer changed") - } - }) - } -} - -func TestGetThinkingAndAnswerHandlesNilInputs(t *testing.T) { - thinking, answer := GetThinkingAndAnswer(nil, nil) - if thinking != nil || answer != nil { - t.Fatalf("GetThinkingAndAnswer(nil, nil)=(%v, %v), want nils", thinking, answer) - } - - content := "\nplan\nanswer" - thinking, answer = GetThinkingAndAnswer(nil, &content) - if thinking != nil { - t.Fatalf("thinking=%v, want nil", thinking) - } - if answer != &content { - t.Fatalf("answer pointer changed") - } - - modelType := "qwen3" - thinking, answer = GetThinkingAndAnswer(&modelType, nil) - if thinking != nil { - t.Fatalf("thinking=%v, want nil", thinking) - } - if answer != nil { - t.Fatalf("answer=%v, want nil", answer) - } -} - func TestBaseModelDoRequestAuthorizationHeader(t *testing.T) { tests := []struct { name string diff --git a/internal/entity/models/response_handler.go b/internal/entity/models/response_handler.go index 4832d565cc..ee6023645b 100644 --- a/internal/entity/models/response_handler.go +++ b/internal/entity/models/response_handler.go @@ -27,12 +27,7 @@ import ( // HandleNonStreamingResponse processes a complete non-streaming chat // response using the ParserConfig's ResponseParser to extract usage. -func HandleNonStreamingResponse( - body []byte, - modelUsage *common.ModelUsage, - chatConfig *ChatConfig, - cfg *ParserConfig, -) (*ChatResponse, error) { +func HandleNonStreamingResponse(body []byte, modelUsage *common.ModelUsage, chatConfig *ChatConfig, cfg *ParserConfig) (*ChatResponse, error) { var result map[string]any if err := json.Unmarshal(body, &result); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) @@ -75,13 +70,7 @@ func HandleNonStreamingResponse( // HandleStreamingResponse processes a streaming chat response using the // ParserConfig's StreamParser to extract usage from each event. -func HandleStreamingResponse( - body io.Reader, - modelUsage *common.ModelUsage, - chatConfig *ChatConfig, - cfg *ParserConfig, - sender func(*string, *string) error, -) error { +func HandleStreamingResponse(body io.Reader, modelUsage *common.ModelUsage, chatConfig *ChatConfig, cfg *ParserConfig, sender func(*string, *string) error) error { if sender == nil { return fmt.Errorf("sender is required") } @@ -103,6 +92,14 @@ func HandleStreamingResponse( return fmt.Errorf("upstream stream error: %v", apiErr) } + // Some providers emit a terminal event that carries only a root-level + // finish_reason with no choices. Check it before the choices guard so + // the stream terminates cleanly instead of being rejected as "ended + // before [DONE] or finish_reason". + if finishReason, ok := event["finish_reason"].(string); ok && finishReason != "" { + sawTerminal = true + } + choices, ok := event["choices"].([]any) if !ok || len(choices) == 0 { return nil @@ -142,9 +139,6 @@ func HandleStreamingResponse( if finishReason, ok := firstChoice["finish_reason"].(string); ok && finishReason != "" { sawTerminal = true } - if finishReason, ok := event["finish_reason"].(string); ok && finishReason != "" { - sawTerminal = true - } return nil }) diff --git a/internal/entity/models/response_handler_test.go b/internal/entity/models/response_handler_test.go new file mode 100644 index 0000000000..c83dda80ab --- /dev/null +++ b/internal/entity/models/response_handler_test.go @@ -0,0 +1,81 @@ +// +// 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 ( + "bytes" + "io" + "strings" + "testing" +) + +// TestHandleStreamingResponseBareRootFinishReason verifies the shared handler +// terminates the stream when an event carries only a root-level finish_reason +// with no choices and no trailing [DONE]. Before the fix this event was +// skipped by the choices guard and the handler returned "stream ended before +// [DONE] or finish_reason". +func TestHandleStreamingResponseBareRootFinishReason(t *testing.T) { + sse := "data: {\"finish_reason\":\"stop\"}\n" + + var contentChunks, reasonChunks []string + err := HandleStreamingResponse( + io.NopCloser(bytes.NewBufferString(sse)), + nil, + nil, + OpenAIParserConfig, + func(content *string, reason *string) error { + if content != nil && *content != "" { + contentChunks = append(contentChunks, *content) + } + if reason != nil && *reason != "" { + reasonChunks = append(reasonChunks, *reason) + } + return nil + }, + ) + if err != nil { + t.Fatalf("HandleStreamingResponse: %v", err) + } + + // The terminal [DONE] marker is the only content the handler emits for a + // finish_reason-only event: no tokens arrived and no reasoning was sent. + if got := strings.Join(contentChunks, ""); got != "[DONE]" { + t.Errorf("content=%q, want only the [DONE] marker", got) + } + if len(reasonChunks) != 0 { + t.Errorf("reasoning=%q, want none for a finish_reason-only event", strings.Join(reasonChunks, "")) + } +} + +// TestHandleStreamingResponseTruncatedStream verifies a stream that ends +// without [DONE], a root-level finish_reason, or a per-choice finish_reason +// is rejected as truncated. This is the counterpart of the bare +// finish_reason case: only a genuine truncation must fail. +func TestHandleStreamingResponseTruncatedStream(t *testing.T) { + sse := "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n" + + err := HandleStreamingResponse( + io.NopCloser(bytes.NewBufferString(sse)), + nil, + nil, + OpenAIParserConfig, + func(*string, *string) error { return nil }, + ) + if err == nil || !strings.Contains(err.Error(), "stream ended before [DONE] or finish_reason") { + t.Fatalf("expected truncation error, got %v", err) + } +} diff --git a/internal/entity/models/vllm.go b/internal/entity/models/vllm.go index f7bbef29c9..2187554e4c 100644 --- a/internal/entity/models/vllm.go +++ b/internal/entity/models/vllm.go @@ -91,84 +91,12 @@ func (v *VllmModel) ChatWithMessages(ctx context.Context, modelName string, mess } } - jsonData, err := json.Marshal(reqBody) + body, err := v.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 := v.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 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) @@ -206,84 +134,11 @@ func (v *VllmModel) ChatStreamlyWithSender(ctx context.Context, modelName string } } - 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") - if auth := BearerAuth(apiConfig); auth != "" { - req.Header.Set("Authorization", auth) - } - - resp, err := v.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 v.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender) + }) } // Encode encodes a list of texts into embeddings diff --git a/internal/entity/models/volcengine.go b/internal/entity/models/volcengine.go index 12e64cdd14..5304c7fe83 100644 --- a/internal/entity/models/volcengine.go +++ b/internal/entity/models/volcengine.go @@ -24,7 +24,6 @@ import ( "io" "net/http" "ragflow/internal/common" - "sort" "strings" ) @@ -52,37 +51,6 @@ func (v *VolcEngine) Name() string { return "volcengine" } -// VolcEngineChatResponse mirrors Ark's OpenAI-compatible chat response. -type VolcEngineChatResponse 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"` -} - // getAPIKey extracts the actual API key from VolcEngine's stored format. // VolcEngine stores the api_key as a JSON string like: // @@ -103,6 +71,20 @@ func (v *VolcEngine) getAPIKey(apiConfig *APIConfig) string { return key } +// effectiveAPIConfig returns an APIConfig with the extracted API key +// suitable for use with doRequest/doStreamRequest. +func (v *VolcEngine) effectiveAPIConfig(apiConfig *APIConfig) *APIConfig { + if apiConfig == nil { + return nil + } + key := v.getAPIKey(apiConfig) + return &APIConfig{ + ApiKey: &key, + Region: apiConfig.Region, + BaseURL: apiConfig.BaseURL, + } +} + // ChatWithMessages sends multiple messages with roles and returns response func (v *VolcEngine) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := v.baseModel.APIConfigCheck(apiConfig); err != nil { @@ -165,75 +147,13 @@ func (v *VolcEngine) ChatWithMessages(ctx context.Context, modelName string, mes } - jsonData, err := json.Marshal(reqBody) + effectiveConfig := v.effectiveAPIConfig(apiConfig) + body, err := v.baseModel.doRequest(ctx, url, effectiveConfig, 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", v.getAPIKey(apiConfig))) - - resp, err := v.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 VolcEngineChatResponse - if err := json.Unmarshal(body, &result); err != nil { - return chatResponseParts{}, fmt.Errorf("failed to unmarshal 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) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) @@ -250,7 +170,7 @@ func (v *VolcEngine) ChatStreamlyWithSender(ctx context.Context, modelName strin if err != nil { return err } - url := fmt.Sprintf("%s/chat/completions", resolvedBaseURL) + url := fmt.Sprintf("%s/%s", resolvedBaseURL, v.baseModel.URLSuffix.Chat) // Build request body with streaming enabled reqBody := buildRequestBody(modelConfig, modelName, messages, true) @@ -309,157 +229,10 @@ func (v *VolcEngine) 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", v.getAPIKey(apiConfig))) - - resp, err := v.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) - sawTerminal := false - done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) 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"].([]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 - } - - content, ok := delta["content"].(string) - if ok && content != "" { - if err := sender(&content, nil); err != nil { - return err - } - } - - if tcs, ok := delta["tool_calls"].([]interface{}); ok { - for _, tc := range tcs { - tcMap, ok := tc.(map[string]interface{}) - if !ok { - continue - } - idxF, ok := tcMap["index"].(float64) - if !ok { - continue - } - idx := int(idxF) - existing, hasExisting := accumulatedToolCalls[idx] - if !hasExisting { - accumulatedToolCalls[idx] = cloneMap(tcMap) - continue - } - if id, ok := tcMap["id"].(string); ok && id != "" { - if eid, ok := existing["id"].(string); ok { - existing["id"] = eid + id - } else { - existing["id"] = id - } - } - if typ, ok := tcMap["type"].(string); ok && typ != "" { - existing["type"] = typ - } - if fn, ok := tcMap["function"].(map[string]interface{}); ok { - ef, ok := existing["function"].(map[string]interface{}) - if !ok { - ef = make(map[string]interface{}) - existing["function"] = ef - } - if name, ok := fn["name"].(string); ok && name != "" { - if en, ok := ef["name"].(string); ok { - ef["name"] = en + name - } else { - ef["name"] = name - } - } - if args, ok := fn["arguments"].(string); ok && args != "" { - if ea, ok := ef["arguments"].(string); ok { - ef["arguments"] = ea + args - } else { - ef["arguments"] = args - } - } - } - } - } - - reasoningContent, ok := delta["reasoning_content"].(string) - if ok && reasoningContent != "" { - if err := sender(nil, &reasoningContent); err != nil { - return err - } - } - - return nil + effectiveConfig := v.effectiveAPIConfig(apiConfig) + return v.baseModel.doStreamRequest(ctx, url, effectiveConfig, 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("volcengine: stream ended before [DONE] or finish_reason") - } - - if len(accumulatedToolCalls) > 0 && modelConfig != nil { - indices := make([]int, 0, len(accumulatedToolCalls)) - for idx := range accumulatedToolCalls { - indices = append(indices, idx) - } - sort.Ints(indices) - tcs := make([]map[string]interface{}, 0, len(accumulatedToolCalls)) - for _, idx := range indices { - tcs = append(tcs, accumulatedToolCalls[idx]) - } - modelConfig.ToolCallsResult = &tcs - } - - // Send [DONE] marker for OpenAI compatibility - endOfStream := "[DONE]" - if err = sender(&endOfStream, nil); err != nil { - return err - } - - return nil } type volcengineEmbeddingResponse struct { diff --git a/internal/entity/models/volcengine_test.go b/internal/entity/models/volcengine_test.go index 32d45f75dd..d733aa36a4 100644 --- a/internal/entity/models/volcengine_test.go +++ b/internal/entity/models/volcengine_test.go @@ -120,6 +120,11 @@ func TestVolcEngineChatStreamSupportsMaxEffortAndUsage(t *testing.T) { withSSRFBypass(t) ctx := t.Context() srv := newVolcEngineServer(t, func(t *testing.T, r *http.Request, w http.ResponseWriter) { + // The streaming endpoint must honor URLSuffix.Chat, not a hardcoded + // "chat/completions" path. + if r.URL.Path != "/v1/chat/completions" { + t.Errorf("path=%s, want /v1/chat/completions", r.URL.Path) + } var body map[string]interface{} if err := json.NewDecoder(r.Body).Decode(&body); err != nil { t.Fatalf("decode request: %v", err) @@ -142,7 +147,11 @@ func TestVolcEngineChatStreamSupportsMaxEffortAndUsage(t *testing.T) { thinking := true effort := "max" config := &ChatConfig{Thinking: &thinking, Effort: &effort} - if err := newVolcEngineForTest(srv.URL).ChatStreamlyWithSender( + driver := NewVolcEngine( + map[string]string{"default": srv.URL}, + URLSuffix{Chat: "v1/chat/completions", Models: "models"}, + ) + if err := driver.ChatStreamlyWithSender( ctx, "doubao-seed-2-0-pro-260215", []Message{{Role: "user", Content: "hello"}}, diff --git a/internal/entity/models/xai.go b/internal/entity/models/xai.go index b509ce3750..87d9379d74 100644 --- a/internal/entity/models/xai.go +++ b/internal/entity/models/xai.go @@ -80,81 +80,12 @@ func (x *XAIModel) ChatWithMessages(ctx context.Context, modelName string, messa url := fmt.Sprintf("%s/%s", baseURL, x.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) - jsonData, err := json.Marshal(reqBody) + body, err := x.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 := x.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") - } - - // xAI reasoning models (grok-3-mini and similar) return reasoning text in - // the reasoning_content field. Pass it through when present. - var reasonContent string - if rc, ok := messageMap["reasoning_content"].(string); ok { - reasonContent = rc - 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 the response @@ -174,89 +105,13 @@ func (x *XAIModel) ChatStreamlyWithSender(ctx context.Context, modelName string, baseURL = strings.TrimSuffix(baseURL, "/") url := fmt.Sprintf("%s/%s", baseURL, x.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) - - 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 := x.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]interface{}](resp.Body, func(event map[string]interface{}) error { - 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 - } - } - - finishReason, ok := firstChoice["finish_reason"].(string) - if ok && finishReason != "" { - sawTerminal = true - } - return nil + return x.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("xai: stream ended before [DONE] or finish_reason") - } - - // Send the [DONE] marker for OpenAI compatibility - endOfStream := "[DONE]" - if err := sender(&endOfStream, nil); err != nil { - return err - } - - return nil } // Embed embeds a list of texts into embeddings. xAI does not expose a @@ -282,33 +137,11 @@ func (x *XAIModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([]List } url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), modelsSuffix) - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "GET", url, nil) + body, err := x.baseModel.doGetRequest(ctx, url, apiConfig, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) + return nil, err } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) - - resp, err := x.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 // Parse response var modelList ModelList if err = json.Unmarshal(body, &modelList); err != nil { diff --git a/internal/entity/models/xai_test.go b/internal/entity/models/xai_test.go index 6a3e46bff7..7311d04bff 100644 --- a/internal/entity/models/xai_test.go +++ b/internal/entity/models/xai_test.go @@ -2,6 +2,7 @@ package models import ( "encoding/json" + "io" "net/http" "net/http/httptest" "os" @@ -137,3 +138,109 @@ func TestXAIListModelsRequiresModelsSuffix(t *testing.T) { t.Fatalf("expected missing models suffix error, got %v", err) } } + +func TestXAIChatHappyPath(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + 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) + } + if got := r.Header.Get("Authorization"); got != "Bearer test-key" { + t.Errorf("Authorization=%q, want Bearer test-key", got) + } + var body map[string]interface{} + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Errorf("decode request: %v", err) + return + } + if body["stream"] != false { + t.Errorf("stream=%v, want false", body["stream"]) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "choices": []map[string]interface{}{{ + "message": map[string]interface{}{ + "content": "pong", + "reasoning_content": "\nthought", + }, + }}, + }) + })) + defer srv.Close() + + apiKey := "test-key" + resp, err := newXAIForTest(srv.URL).ChatWithMessages( + ctx, + "grok-3-mini", + []Message{{Role: "user", Content: "ping"}}, + &APIConfig{ApiKey: &apiKey}, + nil, + nil, + ) + if err != nil { + t.Fatalf("ChatWithMessages: %v", err) + } + if resp.Answer == nil || *resp.Answer != "pong" { + t.Errorf("Answer=%v, want pong", resp.Answer) + } + if resp.ReasonContent == nil || *resp.ReasonContent != "thought" { + t.Errorf("ReasonContent=%v, want thought", resp.ReasonContent) + } +} + +func TestXAIStreamHappyPath(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Accept"); got != "text/event-stream" { + t.Errorf("Accept=%q, want text/event-stream", got) + } + var body map[string]interface{} + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Errorf("decode request: %v", err) + return + } + streamOptions, ok := body["stream_options"].(map[string]interface{}) + if !ok || streamOptions["include_usage"] != true { + t.Errorf("stream_options=%#v, want include_usage=true", body["stream_options"]) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, strings.Join([]string{ + `data: {"choices":[{"delta":{"reasoning_content":"step "}}]}`, + `data: {"choices":[{"delta":{"content":"Hello"}}]}`, + `data: {"choices":[{"delta":{"content":" world"},"finish_reason":"stop"}]}`, + `data: [DONE]`, + ``, + }, "\n")) + })) + defer srv.Close() + + apiKey := "test-key" + var content, reasoning []string + err := newXAIForTest(srv.URL).ChatStreamlyWithSender( + ctx, + "grok-3-mini", + []Message{{Role: "user", Content: "hi"}}, + &APIConfig{ApiKey: &apiKey}, + nil, + nil, + func(answer, reason *string) error { + if answer != nil { + content = append(content, *answer) + } + if reason != nil { + reasoning = append(reasoning, *reason) + } + return nil + }, + ) + if err != nil { + t.Fatalf("ChatStreamlyWithSender: %v", err) + } + if strings.Join(reasoning, "") != "step " { + t.Errorf("reasoning=%q", strings.Join(reasoning, "")) + } + if got := strings.Join(content, ""); got != "Hello world[DONE]" { + t.Errorf("content=%q, want Hello world[DONE]", got) + } +} diff --git a/internal/entity/models/xiaomi.go b/internal/entity/models/xiaomi.go index 23440b2051..70b3a18fea 100644 --- a/internal/entity/models/xiaomi.go +++ b/internal/entity/models/xiaomi.go @@ -41,6 +41,11 @@ func NewXiaomiModel(baseURL map[string]string, urlSuffix URLSuffix) *XiaomiModel BaseURL: baseURL, URLSuffix: urlSuffix, httpClient: NewDriverHTTPClient(false), + // Xiaomi authenticates with the non-standard "api-key" header + // instead of "Authorization: Bearer". + authHeader: func(cfg *APIConfig) (string, string) { + return "api-key", *cfg.ApiKey + }, }, } } @@ -75,7 +80,6 @@ func (x *XiaomiModel) ChatWithMessages(ctx context.Context, modelName string, me delete(reqBody, "max_tokens") if chatModelConfig != nil { - if chatModelConfig.MaxTokens != nil { reqBody["max_completion_tokens"] = *chatModelConfig.MaxTokens } @@ -93,94 +97,12 @@ func (x *XiaomiModel) ChatWithMessages(ctx context.Context, modelName string, me } } - jsonData, err := json.Marshal(reqBody) + body, err := x.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("api-key", *apiConfig.ApiKey) - - resp, err := x.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) - if !ok { - return nil, fmt.Errorf("invalid content format") - } - - var reasonContent string - reasonContent, _ = messageMap["reasoning_content"].(string) - if reasonContent == "" && chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { - // If reasoning_content not in response, try parsing from content tags - reasoning, answer := GetThinkingAndAnswer(chatModelConfig.ModelClass, &content) - if reasoning != nil { - reasonContent = *reasoning - content = *answer - } - } - // if first char of reasonContent is \n remove the '\n' - if reasonContent != "" && reasonContent[0] == '\n' { - reasonContent = reasonContent[1:] - } - - var toolCalls []map[string]interface{} - if tcs, ok := messageMap["tool_calls"].([]interface{}); ok { - for _, tc := range tcs { - if tcMap, ok := tc.(map[string]interface{}); ok { - toolCalls = append(toolCalls, tcMap) - } - } - } - - chatResponse := &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - ToolCalls: toolCalls, - } - - return chatResponse, nil + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (x *XiaomiModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { @@ -204,6 +126,9 @@ func (x *XiaomiModel) ChatStreamlyWithSender(ctx context.Context, modelName stri // Build request body with streaming enabled reqBody := buildRequestBody(modelConfig, modelName, messages, true) delete(reqBody, "max_tokens") + reqBody["stream_options"] = map[string]interface{}{ + "include_usage": true, + } if modelConfig != nil { if modelConfig.MaxTokens != nil { @@ -224,84 +149,9 @@ func (x *XiaomiModel) ChatStreamlyWithSender(ctx context.Context, modelName stri } - 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("api-key", *apiConfig.ApiKey) - - resp, err := x.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 - if modelConfig != nil { - modelConfig.ToolCallsResult = nil - } - accumulatedToolCalls := make(map[int]map[string]interface{}) - if _, err = ParseSSEStreamTolerant[map[string]interface{}](resp.Body, func(event map[string]interface{}) error { - 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 x.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender) + }) } func (x *XiaomiModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) { @@ -485,7 +335,7 @@ func (x *XiaomiModel) newXiaomiASRRequest(ctx context.Context, modelName *string return nil, fmt.Errorf("failed to create request: %w", err) } req.Header.Set("Content-Type", "application/json") - req.Header.Set("api-key", *apiConfig.ApiKey) + x.baseModel.applyAuth(req, apiConfig) return req, nil } @@ -679,7 +529,7 @@ func (x *XiaomiModel) newXiaomiTTSRequest(ctx context.Context, modelName *string return nil, fmt.Errorf("failed to create request: %w", err) } req.Header.Set("Content-Type", "application/json") - req.Header.Set("api-key", *apiConfig.ApiKey) + x.baseModel.applyAuth(req, apiConfig) return req, nil } diff --git a/internal/entity/models/xiaomi_test.go b/internal/entity/models/xiaomi_test.go index b46f8b3972..1fe90402b3 100644 --- a/internal/entity/models/xiaomi_test.go +++ b/internal/entity/models/xiaomi_test.go @@ -251,6 +251,10 @@ func TestXiaomiStreamHappyPath(t *testing.T) { if body["stream"] != true { t.Errorf("stream=%v want true", body["stream"]) } + streamOptions, ok := body["stream_options"].(map[string]interface{}) + if !ok || streamOptions["include_usage"] != true { + t.Errorf("stream_options=%#v, want include_usage=true", body["stream_options"]) + } w.Header().Set("Content-Type", "text/event-stream") _, _ = io.WriteString(w, `data: {"choices":[{"delta":{"reasoning_content":"step "}}]}`+"\n\n"+ @@ -344,10 +348,11 @@ func TestXiaomiStreamRejectsMalformedFrame(t *testing.T) { defer srv.Close() apiKey := "test-key" - // Malformed SSE frames are silently skipped; the stream completes and sends [DONE]. + // Malformed SSE frames abort the stream: Xiaomi now uses the strict + // OpenAIParserConfig shared by every OpenAI-compatible driver. err := newXiaomiForTest(srv.URL).ChatStreamlyWithSender(ctx, "mimo-v2.5-pro", []Message{{Role: "user", Content: "x"}}, &APIConfig{ApiKey: &apiKey}, nil, nil, func(*string, *string) error { return nil }) - if err != nil { - t.Errorf("expected no error on malformed frame, got %v", err) + if err == nil { + t.Error("expected error on malformed frame, got nil") } } diff --git a/internal/entity/models/xinference.go b/internal/entity/models/xinference.go index 5647575a8c..ebd42eeb18 100644 --- a/internal/entity/models/xinference.go +++ b/internal/entity/models/xinference.go @@ -28,30 +28,13 @@ import ( "ragflow/internal/common" "strconv" "strings" - "sync" - "time" ) -var xinferenceStreamIdleTimeout = 60 * time.Second - // XinferenceModel implements ModelDriver for Xinference chat models. type XinferenceModel struct { baseModel BaseModel } -type xinferenceChatChoice struct { - Message struct { - Content string `json:"content"` - ReasoningContent string `json:"reasoning_content"` - Reasoning string `json:"reasoning"` - Thinking string `json:"thinking"` - } `json:"message"` -} - -type xinferenceChatResponse struct { - Choices []xinferenceChatChoice `json:"choices"` -} - type xinferenceModelListResponse struct { Data []ModelListItem `json:"data"` } @@ -87,28 +70,6 @@ func normalizeXinferenceBaseURL(base string) string { return trimmed } -func xinferenceReasoningFromStrings(reasoningContent string, reasoning string, thinking string) string { - switch { - case reasoningContent != "": - return reasoningContent - case reasoning != "": - return reasoning - case thinking != "": - return thinking - default: - return "" - } -} - -func xinferenceReasoningFromMap(value map[string]interface{}) string { - for _, field := range []string{"reasoning_content", "reasoning", "thinking"} { - if text, ok := value[field].(string); ok && text != "" { - return text - } - } - return "" -} - // ChatWithMessages sends multiple messages with roles and returns the response. func (x *XinferenceModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := x.baseModel.APIConfigCheck(apiConfig); err != nil { @@ -127,56 +88,13 @@ func (x *XinferenceModel) ChatWithMessages(ctx context.Context, modelName string url := fmt.Sprintf("%s/%s", baseURL, x.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) - jsonData, err := json.Marshal(reqBody) + + body, err := x.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 := x.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 xinferenceChatResponse - if err = json.Unmarshal(body, &result); err != nil { - return nil, fmt.Errorf("failed to parse response: %w", err) - } - if len(result.Choices) == 0 { - return nil, fmt.Errorf("no choices in response") - } - - content := result.Choices[0].Message.Content - reasonContent := xinferenceReasoningFromStrings( - result.Choices[0].Message.ReasoningContent, - result.Choices[0].Message.Reasoning, - result.Choices[0].Message.Thinking, - ) - - return &ChatResponse{ - Answer: &content, - ReasonContent: &reasonContent, - }, nil + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender. @@ -203,101 +121,13 @@ func (x *XinferenceModel) ChatStreamlyWithSender(ctx context.Context, modelName url := fmt.Sprintf("%s/%s", baseURL, x.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) - 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 := x.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(xinferenceStreamIdleTimeout / 4) - defer ticker.Stop() - for { - select { - case <-done: - return - case now := <-ticker.C: - lastActiveMu.Lock() - idle := now.Sub(lastActive) - lastActiveMu.Unlock() - if idle >= xinferenceStreamIdleTimeout { - cancel() - return - } - } - } - }() - - sawTerminal := false - sseDone, parseErr := 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 - } - - if delta, ok := firstChoice["delta"].(map[string]interface{}); ok { - if reasoning := xinferenceReasoningFromMap(delta); 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 - } - } - } - - if finishReason, ok := firstChoice["finish_reason"].(string); ok && finishReason != "" { - sawTerminal = true - } - return nil + return x.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) }) - if parseErr != nil { - if ctx.Err() != nil { - return fmt.Errorf("xinference: stream idle for more than %s, aborted", xinferenceStreamIdleTimeout) - } - return fmt.Errorf("failed to scan response body: %w", parseErr) - } - if !sseDone && !sawTerminal { - return fmt.Errorf("xinference: stream ended before [DONE] or finish_reason") - } - - endOfStream := "[DONE]" - return sender(&endOfStream, nil) } // Index is *int so a missing JSON field is distinguishable from index 0. diff --git a/internal/entity/models/xinference_test.go b/internal/entity/models/xinference_test.go index 3601c5b0ca..c75ded1c41 100644 --- a/internal/entity/models/xinference_test.go +++ b/internal/entity/models/xinference_test.go @@ -7,7 +7,6 @@ import ( "net/http/httptest" "strings" "testing" - "time" ) func newXinferenceForTest(baseURL string) *XinferenceModel { @@ -22,15 +21,6 @@ func newXinferenceForTest(baseURL string) *XinferenceModel { ) } -func withXinferenceIdleTimeout(t *testing.T, d time.Duration) { - t.Helper() - original := xinferenceStreamIdleTimeout - xinferenceStreamIdleTimeout = d - t.Cleanup(func() { - xinferenceStreamIdleTimeout = original - }) -} - func TestXinferenceName(t *testing.T) { x := newXinferenceForTest("http://unused") if got := x.Name(); got != "Xinference" { @@ -229,38 +219,6 @@ func TestXinferenceStreamRejectsFalseStreamConfig(t *testing.T) { } } -func TestXinferenceStreamCancelsOnIdle(t *testing.T) { - withSSRFBypass(t) - ctx := t.Context() - withXinferenceIdleTimeout(t, 200*time.Millisecond) - - 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() - } - select { - case <-hold: - case <-r.Context().Done(): - } - })) - t.Cleanup(srv.Close) - t.Cleanup(func() { close(hold) }) - - x := newXinferenceForTest(srv.URL) - err := x.ChatStreamlyWithSender(ctx, "qwen2.5-instruct", - []Message{{Role: "user", Content: "x"}}, - &APIConfig{}, nil, - nil, - func(*string, *string) error { return nil }) - if err == nil || !strings.Contains(err.Error(), "stream idle") { - t.Errorf("expected stream-idle error, got %v", err) - } -} - func TestXinferenceListModelsAndCheckConnection(t *testing.T) { withSSRFBypass(t) ctx := t.Context() diff --git a/internal/entity/models/xunfei.go b/internal/entity/models/xunfei.go index 52ee0cd4b5..0f1825ba07 100644 --- a/internal/entity/models/xunfei.go +++ b/internal/entity/models/xunfei.go @@ -17,12 +17,10 @@ package models import ( - "bytes" "context" "encoding/json" "fmt" "io" - "net/http" "ragflow/internal/common" "strings" ) @@ -76,6 +74,11 @@ func NewXunFeiModel(baseURL map[string]string, urlSuffix URLSuffix) *XunFeiModel BaseURL: baseURL, URLSuffix: urlSuffix, httpClient: NewDriverHTTPClient(false), + // The Spark HTTP API authenticates with the credential bundle's + // spark_api_password, not the raw stored key. + authHeader: func(cfg *APIConfig) (string, string) { + return "Authorization", "Bearer " + resolveBearerToken(cfg) + }, }, } } @@ -118,82 +121,12 @@ func (x *XunFeiModel) ChatWithMessages(ctx context.Context, modelName string, me } } - jsonData, err := json.Marshal(reqBody) + body, err := x.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to marshal request body: %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", resolveBearerToken(apiConfig))) - - resp, err := x.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 body: %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 unmarshal response body: %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_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) } func (x *XunFeiModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { @@ -227,84 +160,13 @@ func (x *XunFeiModel) ChatStreamlyWithSender(ctx context.Context, modelName stri } } - 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", resolveBearerToken(apiConfig))) - - resp, err := x.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("invalid status code: %d, body: %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 { - if data, marshalErr := json.Marshal(event); marshalErr == nil { - common.Info(string(data)) - } - - 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 + // XunFei's OpenAI-compatible endpoint does not document support for + // stream_options.include_usage, but its streaming responses carry a + // usage object in the final chunk by default, so usage reporting works + // without it. + return x.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender) + }) } func (x *XunFeiModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) { @@ -350,41 +212,11 @@ func (x *XunFeiModel) ListModels(ctx context.Context, apiConfig *APIConfig) ([]L } url := fmt.Sprintf("%s/%s", resolvedBaseURL, x.baseModel.URLSuffix.Models) - // Build request body - reqBody := map[string]interface{}{} - - jsonData, err := json.Marshal(reqBody) + body, err := x.baseModel.doGetRequest(ctx, url, apiConfig, 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, "GET", 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", resolveBearerToken(apiConfig))) - - resp, err := x.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 // Parse response var modelList ModelList if err = json.Unmarshal(body, &modelList); err != nil { diff --git a/internal/entity/models/xunfei_test.go b/internal/entity/models/xunfei_test.go index f3969a6da3..a7f0b51a21 100644 --- a/internal/entity/models/xunfei_test.go +++ b/internal/entity/models/xunfei_test.go @@ -1,6 +1,10 @@ package models import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" "strings" "testing" ) @@ -119,3 +123,119 @@ func TestXunFeiUnsupportedMethodsReturnNoSuchMethod(t *testing.T) { }) } } + +func newXunFeiForTest(baseURL string) *XunFeiModel { + return NewXunFeiModel( + map[string]string{"default": baseURL}, + URLSuffix{Chat: "v1/chat/completions", Models: "v1/models"}, + ) +} + +func TestXunFeiChatUsesResolvedBearerToken(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + 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) + } + // The Spark credential bundle is stored as JSON; the request must + // authenticate with the extracted spark_api_password. + if got := r.Header.Get("Authorization"); got != "Bearer pwd" { + t.Errorf("Authorization=%q, want Bearer pwd", got) + } + var body map[string]interface{} + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Errorf("decode request: %v", err) + return + } + if body["model"] != "lite" { + t.Errorf("model=%v, want lite (resolved Spark-Lite)", body["model"]) + } + if body["stream"] != false { + t.Errorf("stream=%v, want false", body["stream"]) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "choices": []map[string]interface{}{{ + "message": map[string]interface{}{ + "content": "pong", + "reasoning_content": "\nthought", + }, + }}, + }) + })) + defer srv.Close() + + bundle := `{"spark_api_password":"pwd","spark_app_id":"app","spark_api_secret":"secret","spark_api_key":"key"}` + resp, err := newXunFeiForTest(srv.URL).ChatWithMessages( + ctx, + "Spark-Lite", + []Message{{Role: "user", Content: "ping"}}, + &APIConfig{ApiKey: &bundle}, + nil, + nil, + ) + if err != nil { + t.Fatalf("ChatWithMessages: %v", err) + } + if resp.Answer == nil || *resp.Answer != "pong" { + t.Errorf("Answer=%v, want pong", resp.Answer) + } + if resp.ReasonContent == nil || *resp.ReasonContent != "thought" { + t.Errorf("ReasonContent=%v, want thought", resp.ReasonContent) + } +} + +func TestXunFeiStreamHappyPath(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "Bearer sk-plain" { + t.Errorf("Authorization=%q, want Bearer sk-plain", got) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, strings.Join([]string{ + `data: {"choices":[{"delta":{"reasoning_content":"step "}}]}`, + `data: {"choices":[{"delta":{"content":"Hello"}}]}`, + `data: {"choices":[{"delta":{"content":" world"},"finish_reason":"stop"}]}`, + // XunFei carries usage in the final chunk without requiring + // stream_options.include_usage. + `data: {"choices":[],"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`, + `data: [DONE]`, + ``, + }, "\n")) + })) + defer srv.Close() + + apiKey := "sk-plain" + var content, reasoning []string + config := &ChatConfig{} + err := newXunFeiForTest(srv.URL).ChatStreamlyWithSender( + ctx, + "Spark-Lite", + []Message{{Role: "user", Content: "hi"}}, + &APIConfig{ApiKey: &apiKey}, + config, + nil, + func(answer, reason *string) error { + if answer != nil { + content = append(content, *answer) + } + if reason != nil { + reasoning = append(reasoning, *reason) + } + return nil + }, + ) + if err != nil { + t.Fatalf("ChatStreamlyWithSender: %v", err) + } + if strings.Join(reasoning, "") != "step " { + t.Errorf("reasoning=%q", strings.Join(reasoning, "")) + } + if got := strings.Join(content, ""); got != "Hello world[DONE]" { + t.Errorf("content=%q, want Hello world[DONE]", got) + } + if config.UsageResult == nil || config.UsageResult.TotalTokens != 8 { + t.Errorf("UsageResult=%#v, want total tokens 8", config.UsageResult) + } +} diff --git a/internal/entity/models/zhipu-ai.go b/internal/entity/models/zhipu-ai.go index 6b0d7b72be..98920be0fe 100644 --- a/internal/entity/models/zhipu-ai.go +++ b/internal/entity/models/zhipu-ai.go @@ -55,32 +55,6 @@ func (z *ZhipuAIModel) Name() string { return "zhipu" } -type ZhipuChatResponse struct { - Choices []struct { - FinishReason string `json:"finish_reason"` - Index int `json:"index"` - 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"` - Created int `json:"created"` - Id string `json:"id"` - Model string `json:"model"` - Object string `json:"object"` - RequestId string `json:"request_id"` - Usage struct { - CompletionTokens int `json:"completion_tokens"` - PromptTokens int `json:"prompt_tokens"` - PromptTokensDetails struct { - CachedTokens int `json:"cached_tokens"` - } `json:"prompt_tokens_details"` - TotalTokens int `json:"total_tokens"` - } `json:"usage"` -} - // ChatWithMessages sends multiple messages with roles and returns response func (z *ZhipuAIModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := z.baseModel.APIConfigCheck(apiConfig); err != nil { @@ -111,68 +85,14 @@ func (z *ZhipuAIModel) ChatWithMessages(ctx context.Context, modelName string, m } } } - } - jsonData, err := json.Marshal(reqBody) + body, err := z.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 := z.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 ZhipuChatResponse - 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("empty response") - } - - choice := &result.Choices[0] - var reasonContent *string - if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking { - reasonContent = &choice.Message.ReasoningContent - } - - return chatResponseParts{ - RequestID: result.RequestId, - 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 sends messages and streams response via sender function (best performance, no channel) @@ -192,6 +112,9 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(ctx context.Context, modelName str url := fmt.Sprintf("%s/%s", baseURL, z.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) + reqBody["stream_options"] = map[string]interface{}{ + "include_usage": true, + } if chatModelConfig != nil { if chatModelConfig.Thinking != nil { @@ -207,86 +130,9 @@ func (z *ZhipuAIModel) ChatStreamlyWithSender(ctx context.Context, modelName str } } - 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 := z.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)) - - 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 - } - - 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(chatModelConfig, accumulatedToolCalls) - - // Send [DONE] marker for OpenAI compatibility - endOfStream := "[DONE]" - return sender(&endOfStream, nil) + return z.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + }) } type zhipuEmbeddingResponse struct {