diff --git a/internal/entity/models/azure_openai.go b/internal/entity/models/azure_openai.go index 27224edad6..bebef259b3 100644 --- a/internal/entity/models/azure_openai.go +++ b/internal/entity/models/azure_openai.go @@ -17,12 +17,10 @@ package models import ( - "bytes" "context" "encoding/json" "fmt" "io" - "net/http" "ragflow/internal/common" "strings" ) @@ -42,6 +40,11 @@ func NewAzureOpenAIModel(baseURL map[string]string, urlSuffix URLSuffix) *AzureO BaseURL: baseURL, URLSuffix: urlSuffix, httpClient: NewDriverHTTPClient(false), + // Azure OpenAI authenticates with the non-standard "api-key" + // header instead of "Authorization: Bearer". + authHeader: func(cfg *APIConfig) (string, string) { + return "api-key", *cfg.ApiKey + }, }, } } @@ -81,16 +84,11 @@ func (a *AzureOpenAIModel) ChatWithMessages(ctx context.Context, modelName strin baseURL = strings.TrimSuffix(baseURL, "/") url := a.deploymentURL(baseURL, modelName, a.baseModel.URLSuffix.Chat) - apiMessages := make([]map[string]interface{}, len(messages)) - for i, msg := range messages { - apiMessages[i] = map[string]interface{}{ - "role": msg.Role, - "content": msg.Content, - } - } - + // Azure preserves its own body shape (no "model" field; deployment is in + // the URL; temperature defaults to 1) rather than using buildRequestBody + // which would inject an unknown "model" field Azure rejects. reqBody := map[string]interface{}{ - "messages": apiMessages, + "messages": buildChatMessages(messages), "stream": false, "temperature": 1, } @@ -110,35 +108,9 @@ func (a *AzureOpenAIModel) ChatWithMessages(ctx context.Context, modelName strin } } - jsonData, err := json.Marshal(reqBody) + body, err := a.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) 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("api-key", *apiConfig.ApiKey) - - resp, err := a.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 nil, err } return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) @@ -170,16 +142,8 @@ func (a *AzureOpenAIModel) ChatStreamlyWithSender(ctx context.Context, modelName baseURL = strings.TrimSuffix(baseURL, "/") url := a.deploymentURL(baseURL, modelName, a.baseModel.URLSuffix.Chat) - apiMessages := make([]map[string]interface{}, len(messages)) - for i, msg := range messages { - apiMessages[i] = map[string]interface{}{ - "role": msg.Role, - "content": msg.Content, - } - } - reqBody := map[string]interface{}{ - "messages": apiMessages, + "messages": buildChatMessages(messages), "stream": true, } @@ -206,34 +170,9 @@ func (a *AzureOpenAIModel) ChatStreamlyWithSender(ctx context.Context, modelName // is set. 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, "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 := a.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)) - } - - return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + return a.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + }) } type azureEmbeddingResponse struct { @@ -272,35 +211,9 @@ func (a *AzureOpenAIModel) Embed(ctx context.Context, modelName *string, texts [ reqBody["dimensions"] = embeddingConfig.Dimension } - jsonData, err := json.Marshal(reqBody) + body, err := a.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) 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("api-key", *apiConfig.ApiKey) - - resp, err := a.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("azure OpenAI embeddings API error: %s, body: %s", resp.Status, string(body)) + return nil, err } var parsed azureEmbeddingResponse @@ -335,29 +248,9 @@ func (a *AzureOpenAIModel) ListModels(ctx context.Context, apiConfig *APIConfig) url := fmt.Sprintf("%s/%s?api-version=%s", strings.TrimRight(baseURL, "/"), a.baseModel.URLSuffix.Models, azureAPIVersion) - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "GET", url, nil) + body, err := a.baseModel.doGetRequest(ctx, url, apiConfig, nonStreamCallTimeout) if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("api-key", *apiConfig.ApiKey) - - resp, err := a.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 nil, err } // Parse response diff --git a/internal/entity/models/azure_openai_test.go b/internal/entity/models/azure_openai_test.go new file mode 100644 index 0000000000..80d8ee28ab --- /dev/null +++ b/internal/entity/models/azure_openai_test.go @@ -0,0 +1,355 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package models + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func newAzureForTest(baseURL string) *AzureOpenAIModel { + return NewAzureOpenAIModel( + map[string]string{"default": baseURL}, + URLSuffix{Chat: "chat/completions", Models: "deployments", Embedding: "embeddings"}, + ) +} + +func newAzureServer(t *testing.T, deployment, op string, handler func(t *testing.T, body map[string]interface{}, w http.ResponseWriter)) *httptest.Server { + t.Helper() + expectedPath := "/deployments/" + deployment + "/" + op + "?api-version=2024-10-21" + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.String() != expectedPath { + t.Errorf("expected path=%s, got %s", expectedPath, r.URL.String()) + return + } + if got := r.Header.Get("api-key"); got != "test-key" { + t.Errorf("expected api-key=test-key, got %q", got) + return + } + if got := r.Header.Get("Authorization"); got != "" { + t.Errorf("expected no Authorization header, got %q", got) + return + } + if got := r.Header.Get("Content-Type"); !strings.HasPrefix(got, "application/json") { + t.Errorf("expected Content-Type to start with application/json, got %q", got) + return + } + if r.Method == http.MethodGet { + handler(t, nil, w) + return + } + raw, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read body: %v", err) + return + } + var body map[string]interface{} + if err := json.Unmarshal(raw, &body); err != nil { + t.Errorf("unmarshal: %v\nraw=%s", err, string(raw)) + return + } + handler(t, body, w) + })) +} + +func newAzureSSEServer(t *testing.T, deployment, op, ssePayload string) *httptest.Server { + t.Helper() + expectedPath := "/deployments/" + deployment + "/" + op + "?api-version=2024-10-21" + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("expected POST, got %s", r.Method) + return + } + if r.URL.String() != expectedPath { + t.Errorf("expected path=%s, got %s", expectedPath, r.URL.String()) + return + } + if got := r.Header.Get("api-key"); got != "test-key" { + t.Errorf("expected api-key=test-key, got %q", got) + return + } + if got := r.Header.Get("Authorization"); got != "" { + t.Errorf("expected no Authorization header, got %q", got) + return + } + if got := r.Header.Get("Content-Type"); !strings.HasPrefix(got, "application/json") { + t.Errorf("expected Content-Type to start with application/json, got %q", got) + return + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, ssePayload) + })) +} + +func TestAzureName(t *testing.T) { + if got := newAzureForTest("http://unused").Name(); got != "azure-openai" { + t.Errorf("Name()=%q, want azure-openai", got) + } +} + +func TestAzureChatHappyPath(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + srv := newAzureServer(t, "gpt-4o", "chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + messages, ok := body["messages"].([]interface{}) + if !ok || len(messages) != 1 { + t.Errorf("messages=%v", body["messages"]) + } + if body["stream"] != false { + t.Errorf("stream=%v want false", body["stream"]) + } + if body["temperature"] != float64(1) { + t.Errorf("temperature=%v want 1", body["temperature"]) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "choices": []map[string]interface{}{{ + "message": map[string]interface{}{"content": "pong"}, + }}, + "usage": map[string]interface{}{ + "prompt_tokens": 3, + "completion_tokens": 5, + "total_tokens": 8, + }, + }) + }) + defer srv.Close() + + apiKey := "test-key" + resp, err := newAzureForTest(srv.URL).ChatWithMessages( + ctx, + "gpt-4o", + []Message{{Role: "user", Content: "ping"}}, + &APIConfig{ApiKey: &apiKey}, nil, nil, + ) + if err != nil { + t.Fatalf("Chat: %v", err) + } + if *resp.Answer != "pong" { + t.Errorf("Answer=%q, want pong", *resp.Answer) + } +} + +func TestAzureChatRequiresDeployment(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + apiKey := "test-key" + _, err := newAzureForTest("http://unused").ChatWithMessages( + ctx, + "", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil, nil, + ) + if err == nil || !strings.Contains(err.Error(), "deployment name is required") { + t.Fatalf("expected deployment name error, got %v", err) + } +} + +func TestAzureChatRequiresMessages(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + apiKey := "test-key" + _, err := newAzureForTest("http://unused").ChatWithMessages( + ctx, + "gpt-4o", + nil, + &APIConfig{ApiKey: &apiKey}, nil, nil, + ) + if err == nil || !strings.Contains(err.Error(), "messages is empty") { + t.Fatalf("expected messages error, got %v", err) + } +} + +func TestAzureChatRejectsHTTPError(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + srv := newAzureServer(t, "gpt-4o", "chat/completions", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"unauthorized"}`)) + }) + defer srv.Close() + + apiKey := "test-key" + _, err := newAzureForTest(srv.URL).ChatWithMessages( + ctx, + "gpt-4o", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil, nil, + ) + if err == nil { + t.Fatal("expected error, got nil") + } +} + +func TestAzureChatStreamHappyPath(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + srv := newAzureSSEServer(t, "gpt-4o", "chat/completions", + `data: {"choices":[{"index":0,"delta":{"role":"assistant","content":"hello"}}]}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{"content":" world"}}]}`+"\n"+ + `data: {"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`+"\n"+ + `data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`+"\n"+ + `data: [DONE]`+"\n", + ) + defer srv.Close() + + apiKey := "test-key" + var content []string + var sawDone bool + err := newAzureForTest(srv.URL).ChatStreamlyWithSender( + ctx, + "gpt-4o", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil, nil, + func(c *string, _ *string) error { + if c != nil && *c == "[DONE]" { + sawDone = true + } + if c != nil && *c != "[DONE]" { + content = append(content, *c) + } + return nil + }, + ) + if err != nil { + t.Fatalf("stream: %v", err) + } + if got := strings.Join(content, ""); got != "hello world" { + t.Errorf("content=%q, want 'hello world'", got) + } + if !sawDone { + t.Error("expected [DONE] sentinel") + } +} + +func TestAzureChatRequiresAPIKey(t *testing.T) { + ctx := t.Context() + _, err := newAzureForTest("http://unused").ChatWithMessages( + ctx, + "gpt-4o", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{}, nil, nil, + ) + if err == nil { + t.Fatal("expected error, got nil") + } +} + +func TestAzureChatStreamRejectsExplicitFalse(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + apiKey := "test-key" + stream := false + err := newAzureForTest("http://unused").ChatStreamlyWithSender( + ctx, + "gpt-4o", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, &ChatConfig{Stream: &stream}, nil, + func(c *string, _ *string) error { return nil }, + ) + if err == nil || !strings.Contains(err.Error(), "stream must be true") { + t.Fatalf("expected stream error, got %v", err) + } +} + +func TestAzureStreamRequiresSender(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + apiKey := "test-key" + err := newAzureForTest("http://unused").ChatStreamlyWithSender( + ctx, + "gpt-4o", + []Message{{Role: "user", Content: "x"}}, + &APIConfig{ApiKey: &apiKey}, nil, nil, nil, + ) + if err == nil || !strings.Contains(err.Error(), "sender is required") { + t.Fatalf("expected sender error, got %v", err) + } +} + +func TestAzureListModelsHappyPath(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + expectedPath := "/deployments?api-version=2024-10-21" + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.String() != expectedPath { + t.Errorf("expected path=%s, got %s", expectedPath, r.URL.String()) + return + } + if got := r.Header.Get("api-key"); got != "test-key" { + t.Errorf("expected api-key=test-key, got %q", got) + return + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "data": []map[string]interface{}{ + {"id": "gpt-4o"}, + {"id": "gpt-4-turbo"}, + }, + }) + })) + defer srv.Close() + + apiKey := "test-key" + models, err := newAzureForTest(srv.URL).ListModels(ctx, &APIConfig{ApiKey: &apiKey}) + if err != nil { + t.Fatalf("ListModels: %v", err) + } + if len(models) != 2 { + t.Fatalf("got %d models, want 2", len(models)) + } + if models[0].Name != "gpt-4o" || models[1].Name != "gpt-4-turbo" { + t.Errorf("models=%v", models) + } +} + +func TestAzureEmbedHappyPath(t *testing.T) { + withSSRFBypass(t) + ctx := t.Context() + srv := newAzureServer(t, "text-embedding-3-small", "embeddings", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + input, ok := body["input"].([]interface{}) + if !ok || len(input) != 1 || input[0] != "hello" { + t.Errorf("input=%v", body["input"]) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "data": []map[string]interface{}{ + {"index": 0, "embedding": []float64{0.1, 0.2, 0.3}}, + }, + }) + }) + defer srv.Close() + + modelName := "text-embedding-3-small" + apiKey := "test-key" + embeddings, err := newAzureForTest(srv.URL).Embed( + ctx, + &modelName, + []string{"hello"}, + &APIConfig{ApiKey: &apiKey}, nil, nil, + ) + if err != nil { + t.Fatalf("Embed: %v", err) + } + if len(embeddings) != 1 { + t.Fatalf("got %d embeddings, want 1", len(embeddings)) + } + if len(embeddings[0].Embedding) != 3 || embeddings[0].Index != 0 { + t.Errorf("embedding=%#v", embeddings[0]) + } +} diff --git a/internal/entity/models/deepseek.go b/internal/entity/models/deepseek.go index 4c77eae453..f34f277845 100644 --- a/internal/entity/models/deepseek.go +++ b/internal/entity/models/deepseek.go @@ -155,7 +155,7 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(ctx context.Context, modelName st if err != nil { return err } - url := fmt.Sprintf("%s/chat/completions", resolvedBaseURL) + url := fmt.Sprintf("%s/%s", resolvedBaseURL, d.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) if chatModelConfig != nil && chatModelConfig.Thinking != nil && *chatModelConfig.Thinking { @@ -196,34 +196,9 @@ func (d *DeepSeekModel) ChatStreamlyWithSender(ctx context.Context, modelName st // usage when stream_options.include_usage is set. 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, "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 := d.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)) - } - - return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + return d.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + }) } // Embed embeds a list of texts into embeddings diff --git a/internal/entity/models/gpustack.go b/internal/entity/models/gpustack.go index cb4bbf598c..bb740b6b19 100644 --- a/internal/entity/models/gpustack.go +++ b/internal/entity/models/gpustack.go @@ -107,36 +107,9 @@ func (g *GPUStackModel) ChatStreamlyWithSender(ctx context.Context, modelName st 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") - if auth := BearerAuth(apiConfig); auth != "" { - req.Header.Set("Authorization", auth) - } - req.Header.Set("Accept", "text/event-stream") - - 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)) - } - - return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + return g.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + }) } type gpustackModelInfo struct { diff --git a/internal/entity/models/groq.go b/internal/entity/models/groq.go index fa84d9b4a1..e6451ce8d2 100644 --- a/internal/entity/models/groq.go +++ b/internal/entity/models/groq.go @@ -157,34 +157,9 @@ func (g *GroqModel) ChatStreamlyWithSender(ctx context.Context, modelName string reqBody := buildRequestBody(chatModelConfig, modelName, messages, true) applyGroqReasoningRequestParams(reqBody, modelName, chatModelConfig) 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)) - req.Header.Set("Accept", "text/event-stream") - - 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)) - } - - return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + return g.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + }) } type groqModelInfo struct { diff --git a/internal/entity/models/longcat.go b/internal/entity/models/longcat.go index 846a9a2d52..c594eda989 100644 --- a/internal/entity/models/longcat.go +++ b/internal/entity/models/longcat.go @@ -17,7 +17,6 @@ package models import ( - "bytes" "context" "encoding/json" "fmt" @@ -138,34 +137,9 @@ func (l *LongCatModel) ChatStreamlyWithSender(ctx context.Context, modelName str chatModelConfig.UsageResult = nil } - 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 := 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)) - } - - return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + return l.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + }) } type longCatModelInfo struct { diff --git a/internal/entity/models/minimax.go b/internal/entity/models/minimax.go index 20a08d8d48..a5fbb04bb6 100644 --- a/internal/entity/models/minimax.go +++ b/internal/entity/models/minimax.go @@ -166,7 +166,6 @@ func (m *MinimaxModel) ChatStreamlyWithSender(ctx context.Context, modelName str if err := m.baseModel.APIConfigCheck(apiConfig); err != nil { return err } - apiKey := strings.TrimSpace(*apiConfig.ApiKey) modelName, err := validateMinimaxModelName(modelName) if err != nil { return err @@ -205,87 +204,57 @@ func (m *MinimaxModel) ChatStreamlyWithSender(ctx context.Context, modelName str 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) - } + return m.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + // Pipe the response through a base_resp checker. MiniMax can send + // error events (e.g. rate limits) without a choices array, and the + // shared handler skips those silently. We surface them so the retry + // predicates can match and the caller sees the real reason. + pr, pw := io.Pipe() + defer pr.Close() + streamErr := make(chan error, 1) + go func() { + defer pw.Close() + // resp.Body is owned by doStreamRequest — do NOT close it here. - ctx, cancel := context.WithTimeout(ctx, streamCallTimeout) - defer cancel() + var scanErr error + // Ensure streamErr always receives a result, on every exit + // path, so the final receive below can never block. + defer func() { + select { + case streamErr <- scanErr: + default: + } + }() - 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", apiKey)) - req.Header.Set("Accept", "text/event-stream") - - resp, err := m.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("minimax API error: status %d: %s", resp.StatusCode, extractMinimaxErrorBody(body)) - } - - // Pipe the response through a base_resp checker. MiniMax can send - // error events (e.g. rate limits) without a choices array, and the - // shared handler skips those silently. We surface them so the retry - // predicates can match and the caller sees the real reason. - pr, pw := io.Pipe() - // Close pr when this function returns so an early exit from - // HandleStreamingResponse unblocks the producer goroutine below - // (its pw.Write fails) and releases resp.Body instead of leaving - // the reader blocked on a live pipe. - defer pr.Close() - streamErr := make(chan error, 1) - go func() { - defer pw.Close() - defer resp.Body.Close() - - var scanErr error - // Ensure streamErr always receives a result, on every exit - // path, so the final receive below can never block. - defer func() { - select { - case streamErr <- scanErr: - default: - } - }() - - scanner := bufio.NewScanner(resp.Body) - scanner.Buffer(make([]byte, 64*1024), 1024*1024) - for scanner.Scan() { - line := scanner.Text() - if strings.HasPrefix(line, "data:") { - data := strings.TrimSpace(line[5:]) - if data != "" && data != "[DONE]" { - var event map[string]any - if json.Unmarshal([]byte(data), &event) == nil { - if errMsg := extractMinimaxAPIError(event); errMsg != "" { - pw.CloseWithError(fmt.Errorf("minimax API error: %s", errMsg)) - return + scanner := bufio.NewScanner(body) + scanner.Buffer(make([]byte, 64*1024), 1024*1024) + for scanner.Scan() { + line := scanner.Text() + if strings.HasPrefix(line, "data:") { + data := strings.TrimSpace(line[5:]) + if data != "" && data != "[DONE]" { + var event map[string]any + if json.Unmarshal([]byte(data), &event) == nil { + if errMsg := extractMinimaxAPIError(event); errMsg != "" { + pw.CloseWithError(fmt.Errorf("minimax API error: %s", errMsg)) + return + } } } } + if _, err := pw.Write([]byte(line + "\n")); err != nil { + scanErr = err + return + } } - if _, err := pw.Write([]byte(line + "\n")); err != nil { - scanErr = err - return - } - } - scanErr = scanner.Err() - }() + scanErr = scanner.Err() + }() - if err := HandleStreamingResponse(pr, modelUsage, modelConfig, OpenAIParserConfig, sender); err != nil { - return err - } - return <-streamErr + if err := HandleStreamingResponse(pr, modelUsage, modelConfig, OpenAIParserConfig, sender); err != nil { + return err + } + return <-streamErr + }) } // Embed embeds a list of texts into embeddings diff --git a/internal/entity/models/moonshot.go b/internal/entity/models/moonshot.go index 8f921fb193..40ae62e327 100644 --- a/internal/entity/models/moonshot.go +++ b/internal/entity/models/moonshot.go @@ -17,7 +17,6 @@ package models import ( - "bytes" "context" "encoding/json" "fmt" @@ -128,7 +127,6 @@ func (m *MoonshotModel) ChatStreamlyWithSender(ctx context.Context, modelName st if err := m.baseModel.APIConfigCheck(apiConfig); err != nil { return err } - apiKey := strings.TrimSpace(*apiConfig.ApiKey) modelName, err := validateMoonshotModelName(modelName) if err != nil { return err @@ -161,35 +159,9 @@ func (m *MoonshotModel) 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, "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", apiKey)) - req.Header.Set("Accept", "text/event-stream") - - resp, err := m.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)) - } - - return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + return m.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + }) } // Embed embeds a list of texts into embeddings diff --git a/internal/entity/models/nvidia.go b/internal/entity/models/nvidia.go index 1d08d704a2..e08cd94987 100644 --- a/internal/entity/models/nvidia.go +++ b/internal/entity/models/nvidia.go @@ -56,7 +56,7 @@ func NewNvidiaModel(baseURL map[string]string, urlSuffix URLSuffix) *NvidiaModel } } -func (n NvidiaModel) NewInstance(baseURL map[string]string) ModelDriver { +func (n *NvidiaModel) NewInstance(baseURL map[string]string) ModelDriver { return NewNvidiaModel(baseURL, n.baseModel.URLSuffix) } @@ -77,11 +77,7 @@ func (n *NvidiaModel) ChatWithMessages(ctx context.Context, modelName string, me if err != nil { return nil, err } - baseURL := resolvedBaseURL - if baseURL == "" { - baseURL = resolvedBaseURL - } - url := fmt.Sprintf("%s/%s", baseURL, n.baseModel.URLSuffix.Chat) + url := fmt.Sprintf("%s/%s", resolvedBaseURL, n.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) if chatModelConfig != nil { @@ -116,11 +112,7 @@ func (n *NvidiaModel) ChatStreamlyWithSender(ctx context.Context, modelName stri if err != nil { return err } - baseURL := resolvedBaseURL - if baseURL == "" { - baseURL = resolvedBaseURL - } - url := fmt.Sprintf("%s/%s", baseURL, n.baseModel.URLSuffix.Chat) + url := fmt.Sprintf("%s/%s", resolvedBaseURL, n.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(modelConfig, modelName, messages, true) if modelConfig != nil { @@ -134,34 +126,10 @@ func (n *NvidiaModel) ChatStreamlyWithSender(ctx context.Context, modelName stri } 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, "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 := n.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)) - } - - return HandleStreamingResponse(resp.Body, modelUsage, modelConfig, OpenAIParserConfig, sender) + return n.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender) + }) } type nvidiaEmbeddingResponse struct { @@ -188,12 +156,8 @@ func (n NvidiaModel) Embed(ctx context.Context, modelName *string, texts []strin if err != nil { return nil, err } - baseURL := resolvedBaseURL - if baseURL == "" { - baseURL = resolvedBaseURL - } - url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), n.baseModel.URLSuffix.Embedding) + url := fmt.Sprintf("%s/%s", strings.TrimSuffix(resolvedBaseURL, "/"), n.baseModel.URLSuffix.Embedding) reqBody := map[string]interface{}{ "model": *modelName, @@ -304,12 +268,8 @@ func (n NvidiaModel) Rerank(ctx context.Context, modelName *string, query string if err != nil { return nil, err } - baseURL := resolvedBaseURL - if baseURL == "" { - baseURL = resolvedBaseURL - } - url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), n.baseModel.URLSuffix.Rerank) + url := fmt.Sprintf("%s/%s", strings.TrimSuffix(resolvedBaseURL, "/"), n.baseModel.URLSuffix.Rerank) topN := len(documents) if rerankConfig != nil && rerankConfig.TopN > 0 && rerankConfig.TopN < topN { diff --git a/internal/entity/models/openai.go b/internal/entity/models/openai.go index 9538264e46..3da7e98438 100644 --- a/internal/entity/models/openai.go +++ b/internal/entity/models/openai.go @@ -131,31 +131,9 @@ func (o *OpenAIModel) ChatStreamlyWithSender(ctx context.Context, modelName stri } } - jsonData, err := json.Marshal(reqBody) - if err != nil { - return fmt.Errorf("failed to marshal request: %w", err) - } - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - 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 := o.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)) - } - - return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + return o.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + }) } type openaiEmbeddingResponse struct { diff --git a/internal/entity/models/siliconflow.go b/internal/entity/models/siliconflow.go index efb0ae71fb..5dac0642a6 100644 --- a/internal/entity/models/siliconflow.go +++ b/internal/entity/models/siliconflow.go @@ -119,34 +119,9 @@ func (s *SiliconflowModel) ChatStreamlyWithSender(ctx context.Context, modelName reqBody["enable_thinking"] = false } - 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 := s.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)) - } - - return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + return s.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error { + return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender) + }) } type siliconflowEmbeddingResponse struct {