From d63bd81d0dc79e822217e50f4436b5decacae46f Mon Sep 17 00:00:00 2001 From: oktofeesh <287075021+oktofeesh1@users.noreply.github.com> Date: Mon, 8 Jun 2026 04:27:19 -0700 Subject: [PATCH] fix(go-models): fix Moonshot model and balance requests (#15528) ## Summary - keep Moonshot chat calls in non-streaming mode and streaming calls in SSE mode - make Moonshot model listing and balance checks use bodyless GET requests - add focused Moonshot request/response regression tests --- internal/entity/models/moonshot.go | 110 ++++--- internal/entity/models/moonshot_test.go | 409 ++++++++++++++++++++++++ 2 files changed, 468 insertions(+), 51 deletions(-) create mode 100644 internal/entity/models/moonshot_test.go diff --git a/internal/entity/models/moonshot.go b/internal/entity/models/moonshot.go index ecf1a32f9e..9aeaf36656 100644 --- a/internal/entity/models/moonshot.go +++ b/internal/entity/models/moonshot.go @@ -24,7 +24,6 @@ import ( "fmt" "io" "net/http" - "ragflow/internal/common" "strings" "time" ) @@ -60,11 +59,22 @@ func (m *MoonshotModel) Name() string { return "moonshot" } +func validateMoonshotModelName(modelName string) (string, error) { + if strings.TrimSpace(modelName) == "" { + return "", fmt.Errorf("model name is required") + } + return strings.TrimSpace(modelName), nil +} + func (m *MoonshotModel) ChatWithMessages(modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig) (*ChatResponse, error) { if err := m.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err } - + apiKey := strings.TrimSpace(*apiConfig.ApiKey) + modelName, err := validateMoonshotModelName(modelName) + if err != nil { + return nil, err + } if len(messages) == 0 { return nil, fmt.Errorf("messages is empty") } @@ -93,10 +103,6 @@ func (m *MoonshotModel) ChatWithMessages(modelName string, messages []Message, a } if chatModelConfig != nil { - if chatModelConfig.Stream != nil { - reqBody["stream"] = *chatModelConfig.Stream - } - if chatModelConfig.MaxTokens != nil { reqBody["max_tokens"] = *chatModelConfig.MaxTokens } @@ -140,7 +146,7 @@ func (m *MoonshotModel) ChatWithMessages(modelName string, messages []Message, a } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) resp, err := m.baseModel.httpClient.Do(req) if err != nil { @@ -205,10 +211,17 @@ func (m *MoonshotModel) ChatStreamlyWithSender(modelName string, messages []Mess if err := m.baseModel.APIConfigCheck(apiConfig); err != nil { return err } - + apiKey := strings.TrimSpace(*apiConfig.ApiKey) + modelName, err := validateMoonshotModelName(modelName) + if err != nil { + return err + } if len(messages) == 0 { return fmt.Errorf("messages is empty") } + if sender == nil { + return fmt.Errorf("sender is required") + } resolvedBaseURL, err := m.baseModel.GetBaseURL(apiConfig) if err != nil { @@ -233,10 +246,6 @@ func (m *MoonshotModel) ChatStreamlyWithSender(modelName string, messages []Mess } if chatModelConfig != nil { - if chatModelConfig.Stream != nil { - reqBody["stream"] = *chatModelConfig.Stream - } - if chatModelConfig.MaxTokens != nil { reqBody["max_tokens"] = *chatModelConfig.MaxTokens } @@ -284,7 +293,8 @@ func (m *MoonshotModel) ChatStreamlyWithSender(modelName string, messages []Mess } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + 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 { @@ -302,7 +312,6 @@ func (m *MoonshotModel) ChatStreamlyWithSender(modelName string, messages []Mess scanner.Buffer(make([]byte, 64*1024), 1024*1024) for scanner.Scan() { line := scanner.Text() - common.Info(line) // SSE data line starts with "data:" if !strings.HasPrefix(line, "data:") { @@ -358,13 +367,17 @@ func (m *MoonshotModel) ChatStreamlyWithSender(modelName string, messages []Mess } } + if err = scanner.Err(); err != nil { + return err + } + // Send [DONE] marker for OpenAI compatibility endOfStream := "[DONE]" if err = sender(&endOfStream, nil); err != nil { return err } - return scanner.Err() + return nil } // Embed embeds a list of texts into embeddings @@ -376,6 +389,7 @@ func (m *MoonshotModel) ListModels(apiConfig *APIConfig) ([]string, error) { if err := m.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err } + apiKey := strings.TrimSpace(*apiConfig.ApiKey) resolvedBaseURL, err := m.baseModel.GetBaseURL(apiConfig) if err != nil { @@ -383,24 +397,16 @@ func (m *MoonshotModel) ListModels(apiConfig *APIConfig) ([]string, error) { } url := fmt.Sprintf("%s/%s", resolvedBaseURL, m.baseModel.URLSuffix.Models) - // Build request body - reqBody := map[string]interface{}{} - - jsonData, err := json.Marshal(reqBody) - if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) - } - ctx, cancel := context.WithTimeout(context.Background(), nonStreamCallTimeout) defer cancel() - req, err := http.NewRequestWithContext(ctx, "GET", url, bytes.NewBuffer(jsonData)) + 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") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) + req.Header.Set("Accept", "application/json") resp, err := m.baseModel.httpClient.Do(req) if err != nil { @@ -417,18 +423,24 @@ func (m *MoonshotModel) ListModels(apiConfig *APIConfig) ([]string, error) { return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - // Parse response - var result map[string]interface{} + var result struct { + Data []struct { + ID string `json:"id"` + } `json:"data"` + } if err = json.Unmarshal(body, &result); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } + if result.Data == nil { + return nil, fmt.Errorf("models response missing data") + } - // convert result["data"] to []map[string]interface{} - models := make([]string, 0) - for _, model := range result["data"].([]interface{}) { - modelMap := model.(map[string]interface{}) - modelName := modelMap["id"].(string) - models = append(models, modelName) + models := make([]string, 0, len(result.Data)) + for _, model := range result.Data { + if strings.TrimSpace(model.ID) == "" { + return nil, fmt.Errorf("models response contains empty id") + } + models = append(models, strings.TrimSpace(model.ID)) } return models, nil @@ -438,6 +450,7 @@ func (m *MoonshotModel) Balance(apiConfig *APIConfig) (map[string]interface{}, e if err := m.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err } + apiKey := strings.TrimSpace(*apiConfig.ApiKey) baseURL, err := m.baseModel.GetBaseURL(apiConfig) if err != nil { @@ -445,24 +458,16 @@ func (m *MoonshotModel) Balance(apiConfig *APIConfig) (map[string]interface{}, e } url := fmt.Sprintf("%s/%s", baseURL, m.baseModel.URLSuffix.Balance) - // Build request body - reqBody := map[string]interface{}{} - - jsonData, err := json.Marshal(reqBody) - if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) - } - ctx, cancel := context.WithTimeout(context.Background(), nonStreamCallTimeout) defer cancel() - req, err := http.NewRequestWithContext(ctx, "GET", url, bytes.NewBuffer(jsonData)) + 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") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey)) + req.Header.Set("Accept", "application/json") resp, err := m.baseModel.httpClient.Do(req) if err != nil { @@ -479,17 +484,20 @@ func (m *MoonshotModel) Balance(apiConfig *APIConfig) (map[string]interface{}, e return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) } - // Parse response - var result map[string]interface{} + var result struct { + Data *struct { + AvailableBalance *float64 `json:"available_balance"` + } `json:"data"` + } if err = json.Unmarshal(body, &result); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) } - - data := result["data"].(map[string]interface{}) - balance := data["available_balance"].(float64) + if result.Data == nil || result.Data.AvailableBalance == nil { + return nil, fmt.Errorf("balance response missing available_balance") + } var response = map[string]interface{}{ - "balance": balance, + "balance": *result.Data.AvailableBalance, "currency": "CNY", } diff --git a/internal/entity/models/moonshot_test.go b/internal/entity/models/moonshot_test.go new file mode 100644 index 0000000000..f80013086b --- /dev/null +++ b/internal/entity/models/moonshot_test.go @@ -0,0 +1,409 @@ +package models + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func newMoonshotServer(t *testing.T, handler func(t *testing.T, r *http.Request, body map[string]interface{}, w http.ResponseWriter)) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "Bearer test-key" { + t.Errorf("expected Authorization=Bearer test-key, got %q", got) + return + } + + var body map[string]interface{} + if r.Method == http.MethodPost { + 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 + } + raw, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read body: %v", err) + return + } + if err := json.Unmarshal(raw, &body); err != nil { + t.Errorf("unmarshal: %v\nraw=%s", err, string(raw)) + return + } + } else { + if r.ContentLength > 0 { + t.Errorf("expected %s request without body, ContentLength=%d", r.Method, r.ContentLength) + return + } + raw, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read body: %v", err) + return + } + if len(raw) != 0 { + t.Errorf("expected %s request without body, got %q", r.Method, string(raw)) + return + } + } + + handler(t, r, body, w) + })) +} + +func newMoonshotForTest(baseURL string) *MoonshotModel { + return NewMoonshotModel( + map[string]string{"default": baseURL}, + URLSuffix{ + Chat: "chat/completions", + Models: "models", + Balance: "users/me/balance", + }, + ) +} + +func TestMoonshotNewInstancePreservesConfig(t *testing.T) { + model := NewMoonshotModel( + map[string]string{"default": "http://old.example"}, + URLSuffix{Chat: "chat", Models: "models", Balance: "balance"}, + ) + + instance, ok := model.NewInstance(map[string]string{"default": "http://new.example"}).(*MoonshotModel) + if !ok { + t.Fatalf("NewInstance type=%T, want *MoonshotModel", instance) + } + if instance.baseModel.BaseURL["default"] != "http://new.example" { + t.Errorf("BaseURL=%q", instance.baseModel.BaseURL["default"]) + } + if instance.baseModel.URLSuffix.Chat != "chat" || instance.baseModel.URLSuffix.Models != "models" || instance.baseModel.URLSuffix.Balance != "balance" { + t.Errorf("URLSuffix=%+v", instance.baseModel.URLSuffix) + } + if instance.baseModel.httpClient == nil { + t.Error("httpClient is nil") + } +} + +func TestMoonshotChatForcesNonStreaming(t *testing.T) { + srv := newMoonshotServer(t, func(t *testing.T, r *http.Request, body map[string]interface{}, w http.ResponseWriter) { + if r.Method != http.MethodPost { + t.Errorf("method=%s, want POST", r.Method) + } + if r.URL.Path != "/chat/completions" { + t.Errorf("path=%s, want /chat/completions", r.URL.Path) + } + if body["model"] != "kimi-k2.6" { + t.Errorf("model=%v, want kimi-k2.6", 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() + + apiKey := " test-key " + stream := true + thinking := true + resp, err := newMoonshotForTest(srv.URL).ChatWithMessages( + " kimi-k2.6 ", + []Message{{Role: "user", Content: "ping"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{Stream: &stream, Thinking: &thinking}, + ) + 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 TestMoonshotStreamForcesStreaming(t *testing.T) { + srv := newMoonshotServer(t, func(t *testing.T, r *http.Request, body map[string]interface{}, w http.ResponseWriter) { + if r.Method != http.MethodPost { + t.Errorf("method=%s, want POST", r.Method) + } + if r.URL.Path != "/chat/completions" { + t.Errorf("path=%s, want /chat/completions", r.URL.Path) + } + if body["stream"] != true { + t.Errorf("stream=%v, want true", body["stream"]) + } + if got := r.Header.Get("Accept"); got != "text/event-stream" { + t.Errorf("Accept=%q, want text/event-stream", got) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, strings.Join([]string{ + `data: {"choices":[{"delta":{"reasoning_content":"thinking"}}]}`, + `data: {"choices":[{"delta":{"content":"hello"}}]}`, + `data: [DONE]`, + ``, + }, "\n")) + }) + defer srv.Close() + + apiKey := "test-key" + stream := false + var content, reasoning []string + var sawDone bool + err := newMoonshotForTest(srv.URL).ChatStreamlyWithSender( + "kimi-k2.6", + []Message{{Role: "user", Content: "ping"}}, + &APIConfig{ApiKey: &apiKey}, + &ChatConfig{Stream: &stream}, + func(answer, reason *string) error { + if answer != nil { + if *answer == "[DONE]" { + sawDone = true + return nil + } + content = append(content, *answer) + } + if reason != nil { + reasoning = append(reasoning, *reason) + } + return nil + }, + ) + if err != nil { + t.Fatalf("ChatStreamlyWithSender: %v", err) + } + if got := strings.Join(content, ""); got != "hello" { + t.Errorf("content=%q, want hello", got) + } + if got := strings.Join(reasoning, ""); got != "thinking" { + t.Errorf("reasoning=%q, want thinking", got) + } + if !sawDone { + t.Error("expected [DONE] sentinel") + } +} + +func TestMoonshotStreamDoesNotSendDoneAfterScannerError(t *testing.T) { + srv := newMoonshotServer(t, func(t *testing.T, _ *http.Request, body map[string]interface{}, w http.ResponseWriter) { + if body["stream"] != true { + t.Errorf("stream=%v, want true", body["stream"]) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: "+strings.Repeat("x", 1024*1024+1)+"\n") + }) + defer srv.Close() + + apiKey := "test-key" + var sawDone bool + err := newMoonshotForTest(srv.URL).ChatStreamlyWithSender( + "kimi-k2.6", + []Message{{Role: "user", Content: "ping"}}, + &APIConfig{ApiKey: &apiKey}, + nil, + func(answer, _ *string) error { + if answer != nil && *answer == "[DONE]" { + sawDone = true + } + return nil + }, + ) + if err == nil { + t.Fatal("expected scanner error") + } + if sawDone { + t.Fatal("sender received [DONE] after scanner error") + } +} + +func TestMoonshotListModelsUsesBodylessGet(t *testing.T) { + srv := newMoonshotServer(t, func(t *testing.T, r *http.Request, _ map[string]interface{}, w http.ResponseWriter) { + if r.Method != http.MethodGet { + t.Errorf("method=%s, want GET", r.Method) + } + if r.URL.Path != "/models" { + t.Errorf("path=%s, want /models", r.URL.Path) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "data": []map[string]string{ + {"id": "kimi-k2.6"}, + {"id": " moonshot-v1-8k "}, + }, + }) + }) + defer srv.Close() + + apiKey := "test-key" + models, err := newMoonshotForTest(srv.URL).ListModels(&APIConfig{ApiKey: &apiKey}) + if err != nil { + t.Fatalf("ListModels: %v", err) + } + if got := strings.Join(models, ","); got != "kimi-k2.6,moonshot-v1-8k" { + t.Errorf("models=%q", got) + } +} + +func TestMoonshotBalanceUsesBodylessGet(t *testing.T) { + srv := newMoonshotServer(t, func(t *testing.T, r *http.Request, _ map[string]interface{}, w http.ResponseWriter) { + if r.Method != http.MethodGet { + t.Errorf("method=%s, want GET", r.Method) + } + if r.URL.Path != "/users/me/balance" { + t.Errorf("path=%s, want /users/me/balance", r.URL.Path) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "code": 0, + "status": true, + "data": map[string]float64{ + "available_balance": 49.5, + "voucher_balance": 46.5, + "cash_balance": 3, + }, + }) + }) + defer srv.Close() + + apiKey := "test-key" + balance, err := newMoonshotForTest(srv.URL).Balance(&APIConfig{ApiKey: &apiKey}) + if err != nil { + t.Fatalf("Balance: %v", err) + } + if balance["balance"] != 49.5 { + t.Errorf("balance=%v, want 49.5", balance["balance"]) + } + if balance["currency"] != "CNY" { + t.Errorf("currency=%v, want CNY", balance["currency"]) + } +} + +func TestMoonshotRejectsMalformedResponses(t *testing.T) { + apiKey := "test-key" + tests := []struct { + name string + response map[string]interface{} + run func(*MoonshotModel) error + }{ + { + name: "models missing data", + response: map[string]interface{}{"object": "list"}, + run: func(m *MoonshotModel) error { + _, err := m.ListModels(&APIConfig{ApiKey: &apiKey}) + return err + }, + }, + { + name: "models empty id", + response: map[string]interface{}{ + "data": []map[string]string{{"id": ""}}, + }, + run: func(m *MoonshotModel) error { + _, err := m.ListModels(&APIConfig{ApiKey: &apiKey}) + return err + }, + }, + { + name: "balance missing available balance", + response: map[string]interface{}{ + "data": map[string]float64{"cash_balance": 3}, + }, + run: func(m *MoonshotModel) error { + _, err := m.Balance(&APIConfig{ApiKey: &apiKey}) + return err + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + srv := newMoonshotServer(t, func(t *testing.T, _ *http.Request, _ map[string]interface{}, w http.ResponseWriter) { + _ = json.NewEncoder(w).Encode(tt.response) + }) + defer srv.Close() + + if err := tt.run(newMoonshotForTest(srv.URL)); err == nil { + t.Fatal("expected malformed response error") + } + }) + } +} + +func TestMoonshotValidatesInputs(t *testing.T) { + apiKey := "test-key" + emptyKey := " " + send := func(*string, *string) error { return nil } + + tests := []struct { + name string + run func() error + want string + }{ + { + name: "chat api key", + run: func() error { + _, err := newMoonshotForTest("http://unused").ChatWithMessages("kimi-k2.6", []Message{{Role: "user", Content: "x"}}, &APIConfig{ApiKey: &emptyKey}, nil) + return err + }, + want: "api key is required", + }, + { + name: "chat model", + run: func() error { + _, err := newMoonshotForTest("http://unused").ChatWithMessages(" ", []Message{{Role: "user", Content: "x"}}, &APIConfig{ApiKey: &apiKey}, nil) + return err + }, + want: "model name is required", + }, + { + name: "stream api key", + run: func() error { + return newMoonshotForTest("http://unused").ChatStreamlyWithSender("kimi-k2.6", []Message{{Role: "user", Content: "x"}}, nil, nil, send) + }, + want: "api key is required", + }, + { + name: "stream model", + run: func() error { + return newMoonshotForTest("http://unused").ChatStreamlyWithSender(" ", []Message{{Role: "user", Content: "x"}}, &APIConfig{ApiKey: &apiKey}, nil, send) + }, + want: "model name is required", + }, + { + name: "stream sender", + run: func() error { + return newMoonshotForTest("http://unused").ChatStreamlyWithSender("kimi-k2.6", []Message{{Role: "user", Content: "x"}}, &APIConfig{ApiKey: &apiKey}, nil, nil) + }, + want: "sender is required", + }, + { + name: "models api key", + run: func() error { + _, err := newMoonshotForTest("http://unused").ListModels(&APIConfig{}) + return err + }, + want: "api key is required", + }, + { + name: "balance api key", + run: func() error { + _, err := newMoonshotForTest("http://unused").Balance(nil) + return err + }, + want: "api key is required", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.run() + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("expected %q error, got %v", tt.want, err) + } + }) + } +}