Files
ragflow/internal/entity/models/mistral_test.go

817 lines
27 KiB
Go
Raw Normal View History

Go: implement Embed (embeddings) in Mistral driver (#14807) ### What problem does this PR solve? The Mistral Go driver landed in #14805 with chat, list models, and check connection. `Embed` was left as a stub that returns `"not implemented"`. This PR fills the gap. `conf/models/mistral.json` did not list any embedding model out of the box, so a tenant who wanted to use Mistral end to end (chat + embeddings) could not run an embedding call. This PR adds `mistral-embed` to the config and a real `/v1/embeddings` implementation. ### What this PR includes - `conf/models/mistral.json`: add `"embedding": "embeddings"` under `url_suffix` so the driver can build the URL from config (matches the `URLSuffix.Embedding` field already used by openai, siliconflow, zhipu-ai), and add a `mistral-embed` entry under `models` (1024-dimensional vectors, 8192 max input tokens). - `internal/entity/models/mistral.go`: replace the `Embed` stub with a real implementation that POSTs to `/v1/embeddings`. Adds local response types `mistralEmbeddingData` and `mistralEmbeddingResponse`. No factory change. No interface change. ### How the implementation works - Validate `apiConfig`, the API key, and the model name. Use the existing `baseURLForRegion` helper so an unknown region fails fast with a clear error. - Wrap the request with `context.WithTimeout(nonStreamCallTimeout)` so the call has a clear deadline. Same pattern as `ChatWithMessages` and `ListModels` already use in this file. - Send all input texts in one request. The Mistral API accepts the `input` field as an array. - Parse `data[*].embedding` and copy each slice into a `[]EmbeddingData` indexed by `data[*].index` so the output order matches the input order even if the API returns items in a different order. - An empty input slice returns `[]EmbeddingData{}` with no HTTP call. - Non-200 responses propagate the upstream status line and body. - A final pass checks that every input slot got a vector. If any slot is still empty, return a clear error so the caller does not silently use a zero vector. ### Note on stacking This PR builds on #14805 (the Mistral driver). Until #14805 merges, this PR's diff on GitHub will include both that PR's commits and this one. After #14805 lands on `main`, GitHub will auto-reduce this PR to only the `Embed` changes (one commit, ~111 line diff in `mistral.go` plus 8 lines in `mistral.json`). ### Type of change - [x] New Feature (non-breaking change which adds functionality) ### How was this tested? - `go build ./internal/entity/models/...` returns exit 0 on go 1.25 (the `go.mod` minimum). - The full method set on `MistralModel` still matches the `ModelDriver` interface. - Pattern parity with the existing OpenAI Embed implementation (`internal/entity/models/openai.go`). Closes #14806 Depends on #14805 Tracking: #14736 --------- Co-authored-by: Jin Hai <haijin.chn@gmail.com>
2026-05-11 23:45:48 -10:00
package models
import (
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
)
// newMistralServer stands up an httptest server that asserts the
// request shape and lets the caller decide what to return.
func newMistralServer(t *testing.T, expectedPath string, handler func(t *testing.T, body map[string]interface{}, w http.ResponseWriter)) *httptest.Server {
t.Helper()
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != expectedPath {
t.Errorf("expected path=%s, got %s", expectedPath, r.URL.Path)
return
}
if got := r.Header.Get("Authorization"); got != "Bearer test-key" {
t.Errorf("expected Authorization=Bearer test-key, got %q", got)
return
}
if r.Method == http.MethodPost {
if got := r.Header.Get("Content-Type"); got != "application/json" {
t.Errorf("expected Content-Type=application/json, got %q", got)
return
}
raw, err := io.ReadAll(r.Body)
if err != nil {
t.Errorf("failed to read body: %v", err)
return
}
var body map[string]interface{}
if err := json.Unmarshal(raw, &body); err != nil {
t.Errorf("invalid JSON body: %v\n%s", err, string(raw))
return
}
handler(t, body, w)
return
}
// GET path: no body
handler(t, nil, w)
}))
}
func newMistralForTest(baseURL string) *MistralModel {
return NewMistralModel(
map[string]string{"default": baseURL},
URLSuffix{
Chat: "chat/completions",
Models: "models",
Embedding: "embeddings",
},
)
}
func TestMistralName(t *testing.T) {
m := newMistralForTest("http://unused")
if got := m.Name(); got != "mistral" {
t.Errorf("Name()=%q, want %q", got, "mistral")
}
}
func TestMistralChatHappyPath(t *testing.T) {
srv := newMistralServer(t, "/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) {
if body["model"] != "mistral-large-latest" {
t.Errorf("expected model=mistral-large-latest, got %v", body["model"])
}
if body["stream"] != false {
t.Errorf("expected stream=false, got %v", body["stream"])
}
msgs, ok := body["messages"].([]interface{})
if !ok || len(msgs) != 1 {
t.Errorf("expected 1 message, got %v", body["messages"])
return
}
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"choices": []map[string]interface{}{
{"message": map[string]interface{}{"content": "pong"}},
},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
resp, err := m.ChatWithMessages("mistral-large-latest", []Message{
{Role: "user", Content: "ping"},
}, &APIConfig{ApiKey: &apiKey}, 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 != "" {
t.Errorf("expected empty reason content, got %v", resp.ReasonContent)
}
}
func TestMistralChatPropagatesConfig(t *testing.T) {
srv := newMistralServer(t, "/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) {
if body["max_tokens"] != float64(64) {
t.Errorf("max_tokens=%v want 64", body["max_tokens"])
}
if body["temperature"] != 0.3 {
t.Errorf("temperature=%v want 0.3", body["temperature"])
}
if body["top_p"] != 0.9 {
t.Errorf("top_p=%v want 0.9", body["top_p"])
}
stop, ok := body["stop"].([]interface{})
if !ok || len(stop) != 1 || stop[0] != "END" {
t.Errorf("stop=%v want [END]", body["stop"])
}
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"choices": []map[string]interface{}{{"message": map[string]interface{}{"content": "ok"}}},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
mt := 64
temp := 0.3
topP := 0.9
stop := []string{"END"}
_, err := m.ChatWithMessages("mistral-large-latest", []Message{{Role: "user", Content: "ping"}},
&APIConfig{ApiKey: &apiKey},
&ChatConfig{MaxTokens: &mt, Temperature: &temp, TopP: &topP, Stop: &stop},
)
if err != nil {
t.Fatalf("ChatWithMessages: %v", err)
}
}
func TestMistralChatRequiresAPIKey(t *testing.T) {
m := newMistralForTest("http://unused")
_, err := m.ChatWithMessages("mistral-large-latest", []Message{{Role: "user", Content: "x"}}, &APIConfig{}, nil)
if err == nil || !strings.Contains(err.Error(), "api key is required") {
t.Errorf("expected api-key error, got %v", err)
}
emptyKey := ""
_, err = m.ChatWithMessages("mistral-large-latest", []Message{{Role: "user", Content: "x"}}, &APIConfig{ApiKey: &emptyKey}, nil)
if err == nil || !strings.Contains(err.Error(), "api key is required") {
t.Errorf("empty key: expected api-key error, got %v", err)
}
}
func TestMistralChatRequiresMessages(t *testing.T) {
m := newMistralForTest("http://unused")
apiKey := "test-key"
_, err := m.ChatWithMessages("mistral-large-latest", nil, &APIConfig{ApiKey: &apiKey}, nil)
if err == nil || !strings.Contains(err.Error(), "messages is empty") {
t.Errorf("expected messages-empty error, got %v", err)
}
}
func TestMistralChatRejectsHTTPError(t *testing.T) {
srv := newMistralServer(t, "/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) {
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error":"unauthorized"}`))
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
_, err := m.ChatWithMessages("mistral-large-latest", []Message{{Role: "user", Content: "x"}}, &APIConfig{ApiKey: &apiKey}, nil)
if err == nil || !strings.Contains(err.Error(), "401") {
t.Errorf("expected 401 propagated, got %v", err)
}
}
func TestMistralChatFallsBackToDefaultOnEmptyRegion(t *testing.T) {
// Empty *Region pointer must fall back to the "default" entry, not
// be treated as an explicit "" region (which would miss the lookup).
srv := newMistralServer(t, "/chat/completions", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"choices": []map[string]interface{}{{"message": map[string]interface{}{"content": "ok"}}},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
emptyRegion := ""
_, err := m.ChatWithMessages("mistral-large-latest",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey, Region: &emptyRegion}, nil)
if err != nil {
t.Errorf("empty Region: expected fallback to default, got %v", err)
}
}
func TestMistralListModelsFallsBackToDefaultOnEmptyRegion(t *testing.T) {
srv := newMistralServer(t, "/models", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{"data": []map[string]interface{}{{"id": "x"}}})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
emptyRegion := ""
if _, err := m.ListModels(&APIConfig{ApiKey: &apiKey, Region: &emptyRegion}); err != nil {
t.Errorf("empty Region: expected fallback to default, got %v", err)
}
}
func TestMistralStreamRequiresSender(t *testing.T) {
m := newMistralForTest("http://unused")
apiKey := "test-key"
err := m.ChatStreamlyWithSender("mistral-large-latest",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey}, nil, nil)
if err == nil || !strings.Contains(err.Error(), "sender is required") {
t.Errorf("expected sender-required error, got %v", err)
}
}
func TestMistralChatRejectsUnknownRegion(t *testing.T) {
m := newMistralForTest("http://unused")
apiKey := "test-key"
region := "eu"
_, err := m.ChatWithMessages("mistral-large-latest", []Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey, Region: &region}, nil)
if err == nil || !strings.Contains(err.Error(), "no base URL configured for region") {
t.Errorf("expected region error, got %v", err)
}
}
func TestMistralStreamHappyPath(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/chat/completions" {
t.Errorf("path=%s", r.URL.Path)
return
}
raw, _ := io.ReadAll(r.Body)
var body map[string]interface{}
_ = json.Unmarshal(raw, &body)
if body["stream"] != true {
t.Errorf("expected stream=true, got %v", body["stream"])
}
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
// Two content chunks then finish_reason terminator, then [DONE].
_, _ = io.WriteString(w,
`data: {"choices":[{"delta":{"content":"Hello "}}]}`+"\n"+
`data: {"choices":[{"delta":{"content":"world"}}]}`+"\n"+
`data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`+"\n"+
`data: [DONE]`+"\n",
)
}))
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
var chunks []string
var sawDone int32
err := m.ChatStreamlyWithSender("mistral-large-latest",
[]Message{{Role: "user", Content: "hi"}},
&APIConfig{ApiKey: &apiKey}, nil,
func(content *string, _ *string) error {
if content == nil {
return nil
}
if *content == "[DONE]" {
atomic.StoreInt32(&sawDone, 1)
return nil
}
chunks = append(chunks, *content)
return nil
},
)
if err != nil {
t.Fatalf("stream: %v", err)
}
if strings.Join(chunks, "") != "Hello world" {
t.Errorf("chunks=%v want [\"Hello \" \"world\"]", chunks)
}
if atomic.LoadInt32(&sawDone) != 1 {
t.Error("expected sender to receive [DONE] sentinel")
}
}
func TestMistralStreamRejectsExplicitFalse(t *testing.T) {
m := newMistralForTest("http://unused")
apiKey := "test-key"
stream := false
err := m.ChatStreamlyWithSender("mistral-large-latest",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey},
&ChatConfig{Stream: &stream},
func(*string, *string) error { return nil },
)
if err == nil || !strings.Contains(err.Error(), "stream must be true") {
t.Errorf("expected stream-true guard, got %v", err)
}
}
func TestMistralStreamFailsWithoutTerminal(t *testing.T) {
// Body closes before [DONE] or a finish_reason -> driver must complain
// instead of pretending the stream finished cleanly.
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"half"}}]}`+"\n")
}))
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
err := m.ChatStreamlyWithSender("mistral-large-latest",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey}, nil,
func(*string, *string) error { return nil },
)
if err == nil || !strings.Contains(err.Error(), "stream ended before") {
t.Errorf("expected stream-truncation error, got %v", err)
}
}
func TestMistralListModelsHappyPath(t *testing.T) {
srv := newMistralServer(t, "/models", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"data": []map[string]interface{}{
{"id": "mistral-large-latest"},
{"id": "mistral-small-latest"},
{"id": "mistral-embed"},
},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
ids, err := m.ListModels(&APIConfig{ApiKey: &apiKey})
if err != nil {
t.Fatalf("ListModels: %v", err)
}
fix(ci): re-enable Go tests and fix compilation errors after ListModels signature change (#15862) ## Summary This PR re-enables the Go test steps in CI that were previously commented out, and fixes all compilation errors that have accumulated in `internal/entity/models/` since the `ListModels` return type was changed from `[]string` to `[]ListModelResponse`. ## Changes ### CI (`.github/workflows/tests.yml`) - Re-enable **Prepare test resources** step (clones resource repo with WordNet data) - Re-enable **Test Go packages** step (runs `go test ./internal/...`) - Fix resource path race condition by using `/tmp/resource-${GITHUB_RUN_ID}` instead of `/tmp/resource` - Exclude `/cli` package from Go tests (contains `main` redeclarations) ### Test fixes (16 model provider test files) All errors were caused by the upstream change from `[]string` to `[]ListModelResponse` in the `ListModels` interface: - Add `joinModelNames` test helper to extract `.Name` from `[]ListModelResponse` slices - `strings.Join(models, ",")` → `joinModelNames(models, ",")` (11 files) - `ids[i] != "..."` → `ids[i].Name != "..."` (cometapi, mistral) - `got[i] != want[i]` → `got[i].Name != want[i]` (bedrock) - `[]string` return types → `[]ListModelResponse` (google) ### Pre-existing bugs in model_test.go Bugs introduced by the upstream `entity/` → `entity/models/` directory rename: - Add missing `pm := GetProviderManager()` calls in 3 test functions - Fix `InitProviderManager` signature (`_, err :=` → `err :=`) - Fix `MaxTokens` `*int` dereference (6 comparisons) - Fix `readProviderConfig` relative path (3 levels up instead of 2) ### model.go - Add `findRepoRoot()` to make `conf/all_models.json` resolution work from any CWD, fixing `TestSiliconFlowProviderConfigLoadsLatestProModels` ### Test validation ```bash go build ./internal/... # ✅ go test ./internal/entity/models/... -count=1 # ✅ all pass ``` 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-09 21:12:15 +08:00
if len(ids) != 3 || ids[0].Name != "mistral-large-latest" || ids[2].Name != "mistral-embed" {
Go: implement Embed (embeddings) in Mistral driver (#14807) ### What problem does this PR solve? The Mistral Go driver landed in #14805 with chat, list models, and check connection. `Embed` was left as a stub that returns `"not implemented"`. This PR fills the gap. `conf/models/mistral.json` did not list any embedding model out of the box, so a tenant who wanted to use Mistral end to end (chat + embeddings) could not run an embedding call. This PR adds `mistral-embed` to the config and a real `/v1/embeddings` implementation. ### What this PR includes - `conf/models/mistral.json`: add `"embedding": "embeddings"` under `url_suffix` so the driver can build the URL from config (matches the `URLSuffix.Embedding` field already used by openai, siliconflow, zhipu-ai), and add a `mistral-embed` entry under `models` (1024-dimensional vectors, 8192 max input tokens). - `internal/entity/models/mistral.go`: replace the `Embed` stub with a real implementation that POSTs to `/v1/embeddings`. Adds local response types `mistralEmbeddingData` and `mistralEmbeddingResponse`. No factory change. No interface change. ### How the implementation works - Validate `apiConfig`, the API key, and the model name. Use the existing `baseURLForRegion` helper so an unknown region fails fast with a clear error. - Wrap the request with `context.WithTimeout(nonStreamCallTimeout)` so the call has a clear deadline. Same pattern as `ChatWithMessages` and `ListModels` already use in this file. - Send all input texts in one request. The Mistral API accepts the `input` field as an array. - Parse `data[*].embedding` and copy each slice into a `[]EmbeddingData` indexed by `data[*].index` so the output order matches the input order even if the API returns items in a different order. - An empty input slice returns `[]EmbeddingData{}` with no HTTP call. - Non-200 responses propagate the upstream status line and body. - A final pass checks that every input slot got a vector. If any slot is still empty, return a clear error so the caller does not silently use a zero vector. ### Note on stacking This PR builds on #14805 (the Mistral driver). Until #14805 merges, this PR's diff on GitHub will include both that PR's commits and this one. After #14805 lands on `main`, GitHub will auto-reduce this PR to only the `Embed` changes (one commit, ~111 line diff in `mistral.go` plus 8 lines in `mistral.json`). ### Type of change - [x] New Feature (non-breaking change which adds functionality) ### How was this tested? - `go build ./internal/entity/models/...` returns exit 0 on go 1.25 (the `go.mod` minimum). - The full method set on `MistralModel` still matches the `ModelDriver` interface. - Pattern parity with the existing OpenAI Embed implementation (`internal/entity/models/openai.go`). Closes #14806 Depends on #14805 Tracking: #14736 --------- Co-authored-by: Jin Hai <haijin.chn@gmail.com>
2026-05-11 23:45:48 -10:00
t.Errorf("ids=%v, want [mistral-large-latest mistral-small-latest mistral-embed]", ids)
}
}
func TestMistralListModelsRequiresAPIKey(t *testing.T) {
m := newMistralForTest("http://unused")
if _, err := m.ListModels(&APIConfig{}); err == nil || !strings.Contains(err.Error(), "api key is required") {
t.Errorf("expected api-key error, got %v", err)
}
}
func TestMistralCheckConnectionDelegatesToListModels(t *testing.T) {
// 200 -> CheckConnection succeeds; 401 -> CheckConnection propagates.
okSrv := newMistralServer(t, "/models", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{"data": []map[string]interface{}{{"id": "x"}}})
})
defer okSrv.Close()
failSrv := newMistralServer(t, "/models", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
w.WriteHeader(http.StatusUnauthorized)
})
defer failSrv.Close()
apiKey := "test-key"
mOK := newMistralForTest(okSrv.URL)
if err := mOK.CheckConnection(&APIConfig{ApiKey: &apiKey}); err != nil {
t.Errorf("CheckConnection(ok): %v", err)
}
mFail := newMistralForTest(failSrv.URL)
if err := mFail.CheckConnection(&APIConfig{ApiKey: &apiKey}); err == nil {
t.Error("CheckConnection(fail): expected error, got nil")
}
}
func TestMistralBalanceReturnsNoSuchMethod(t *testing.T) {
m := newMistralForTest("http://unused")
_, err := m.Balance(&APIConfig{})
if err == nil || !strings.Contains(err.Error(), "no such method") {
t.Errorf("Balance: expected 'no such method', got %v", err)
}
}
func TestMistralRerankReturnsNoSuchMethod(t *testing.T) {
m := newMistralForTest("http://unused")
q := "mistral-large-latest"
_, err := m.Rerank(&q, "what is rag?", []string{"a", "b"}, &APIConfig{}, &RerankConfig{TopN: 2})
if err == nil || !strings.Contains(err.Error(), "no such method") {
t.Errorf("Rerank: expected 'no such method', got %v", err)
}
}
func TestMistralUnsupportedDefaultsReturnNoSuchMethod(t *testing.T) {
m := newMistralForTest("http://unused")
modelName := "mistral-large-latest"
checks := []struct {
name string
call func() error
}{
{"TranscribeAudio", func() error {
_, err := m.TranscribeAudio(&modelName, nil, &APIConfig{}, nil)
return err
}},
{"TranscribeAudioWithSender", func() error {
return m.TranscribeAudioWithSender(&modelName, nil, &APIConfig{}, nil, nil)
}},
{"AudioSpeech", func() error {
_, err := m.AudioSpeech(&modelName, nil, &APIConfig{}, nil)
return err
}},
{"AudioSpeechWithSender", func() error {
return m.AudioSpeechWithSender(&modelName, nil, &APIConfig{}, nil, nil)
}},
{"ParseFile", func() error {
_, err := m.ParseFile(&modelName, nil, nil, &APIConfig{}, nil)
return err
}},
{"ListTasks", func() error {
_, err := m.ListTasks(&APIConfig{})
return err
}},
{"ShowTask", func() error {
_, err := m.ShowTask("task-id", &APIConfig{})
return err
}},
}
for _, check := range checks {
t.Run(check.name, func(t *testing.T) {
requireNoSuchMethod(t, check.name, check.call())
})
}
}
Go: implement Embed (embeddings) in Mistral driver (#14807) ### What problem does this PR solve? The Mistral Go driver landed in #14805 with chat, list models, and check connection. `Embed` was left as a stub that returns `"not implemented"`. This PR fills the gap. `conf/models/mistral.json` did not list any embedding model out of the box, so a tenant who wanted to use Mistral end to end (chat + embeddings) could not run an embedding call. This PR adds `mistral-embed` to the config and a real `/v1/embeddings` implementation. ### What this PR includes - `conf/models/mistral.json`: add `"embedding": "embeddings"` under `url_suffix` so the driver can build the URL from config (matches the `URLSuffix.Embedding` field already used by openai, siliconflow, zhipu-ai), and add a `mistral-embed` entry under `models` (1024-dimensional vectors, 8192 max input tokens). - `internal/entity/models/mistral.go`: replace the `Embed` stub with a real implementation that POSTs to `/v1/embeddings`. Adds local response types `mistralEmbeddingData` and `mistralEmbeddingResponse`. No factory change. No interface change. ### How the implementation works - Validate `apiConfig`, the API key, and the model name. Use the existing `baseURLForRegion` helper so an unknown region fails fast with a clear error. - Wrap the request with `context.WithTimeout(nonStreamCallTimeout)` so the call has a clear deadline. Same pattern as `ChatWithMessages` and `ListModels` already use in this file. - Send all input texts in one request. The Mistral API accepts the `input` field as an array. - Parse `data[*].embedding` and copy each slice into a `[]EmbeddingData` indexed by `data[*].index` so the output order matches the input order even if the API returns items in a different order. - An empty input slice returns `[]EmbeddingData{}` with no HTTP call. - Non-200 responses propagate the upstream status line and body. - A final pass checks that every input slot got a vector. If any slot is still empty, return a clear error so the caller does not silently use a zero vector. ### Note on stacking This PR builds on #14805 (the Mistral driver). Until #14805 merges, this PR's diff on GitHub will include both that PR's commits and this one. After #14805 lands on `main`, GitHub will auto-reduce this PR to only the `Embed` changes (one commit, ~111 line diff in `mistral.go` plus 8 lines in `mistral.json`). ### Type of change - [x] New Feature (non-breaking change which adds functionality) ### How was this tested? - `go build ./internal/entity/models/...` returns exit 0 on go 1.25 (the `go.mod` minimum). - The full method set on `MistralModel` still matches the `ModelDriver` interface. - Pattern parity with the existing OpenAI Embed implementation (`internal/entity/models/openai.go`). Closes #14806 Depends on #14805 Tracking: #14736 --------- Co-authored-by: Jin Hai <haijin.chn@gmail.com>
2026-05-11 23:45:48 -10:00
func TestMistralEmbedHappyPath(t *testing.T) {
srv := newMistralServer(t, "/embeddings", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) {
if body["model"] != "mistral-embed" {
t.Errorf("model=%v want mistral-embed", body["model"])
}
inputs, ok := body["input"].([]interface{})
if !ok || len(inputs) != 3 {
t.Errorf("input=%v want 3-element array", body["input"])
}
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"data": []map[string]interface{}{
{"embedding": []float64{0.1, 0.2}, "index": 0},
{"embedding": []float64{0.3, 0.4}, "index": 1},
{"embedding": []float64{0.5, 0.6}, "index": 2},
},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
model := "mistral-embed"
vecs, err := m.Embed(&model, []string{"a", "b", "c"}, &APIConfig{ApiKey: &apiKey}, nil)
if err != nil {
t.Fatalf("Embed: %v", err)
}
if len(vecs) != 3 {
t.Fatalf("len(vecs)=%d want 3", len(vecs))
}
if vecs[1].Embedding[0] != 0.3 || vecs[1].Index != 1 {
t.Errorf("vecs[1]=%+v want {Embedding:[0.3 0.4] Index:1}", vecs[1])
}
}
func TestMistralEmbedReordersByIndex(t *testing.T) {
// Upstream returns the three vectors in shuffled order. The driver
// must reorder them so the slot at position i corresponds to input i.
srv := newMistralServer(t, "/embeddings", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"data": []map[string]interface{}{
{"embedding": []float64{2}, "index": 2},
{"embedding": []float64{0}, "index": 0},
{"embedding": []float64{1}, "index": 1},
},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
model := "mistral-embed"
vecs, err := m.Embed(&model, []string{"a", "b", "c"}, &APIConfig{ApiKey: &apiKey}, nil)
if err != nil {
t.Fatalf("Embed: %v", err)
}
for i, v := range vecs {
if v.Index != i || v.Embedding[0] != float64(i) {
t.Errorf("slot %d = %+v, want Embedding=[%d] Index=%d", i, v, i, i)
}
}
}
func TestMistralEmbedEmptyInputShortCircuits(t *testing.T) {
// Empty input must NOT make an HTTP call; the test fails the request
// rather than the assertion if it does.
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
t.Error("Embed([]) made an unexpected HTTP call")
w.WriteHeader(http.StatusInternalServerError)
}))
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
model := "mistral-embed"
vecs, err := m.Embed(&model, []string{}, &APIConfig{ApiKey: &apiKey}, nil)
if err != nil {
t.Fatalf("Embed([]): %v", err)
}
if len(vecs) != 0 {
t.Errorf("len(vecs)=%d want 0", len(vecs))
}
}
func TestMistralEmbedRequiresAPIKey(t *testing.T) {
m := newMistralForTest("http://unused")
model := "mistral-embed"
_, err := m.Embed(&model, []string{"a"}, &APIConfig{}, nil)
if err == nil || !strings.Contains(err.Error(), "api key is required") {
t.Errorf("expected api-key error, got %v", err)
}
}
func TestMistralEmbedRequiresModelName(t *testing.T) {
m := newMistralForTest("http://unused")
apiKey := "test-key"
_, err := m.Embed(nil, []string{"a"}, &APIConfig{ApiKey: &apiKey}, nil)
if err == nil || !strings.Contains(err.Error(), "model name is required") {
t.Errorf("expected model-name error, got %v", err)
}
empty := ""
_, err = m.Embed(&empty, []string{"a"}, &APIConfig{ApiKey: &apiKey}, nil)
if err == nil || !strings.Contains(err.Error(), "model name is required") {
t.Errorf("empty model: expected model-name error, got %v", err)
}
}
func TestMistralEmbedRejectsDuplicateIndex(t *testing.T) {
// A malformed upstream that repeats data[*].index would silently
// overwrite the earlier vector; the driver must fail loudly instead.
srv := newMistralServer(t, "/embeddings", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"data": []map[string]interface{}{
{"embedding": []float64{1}, "index": 0},
{"embedding": []float64{2}, "index": 0},
},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
model := "mistral-embed"
_, err := m.Embed(&model, []string{"a", "b"}, &APIConfig{ApiKey: &apiKey}, nil)
if err == nil || !strings.Contains(err.Error(), "duplicate embedding index 0") {
t.Errorf("expected duplicate-index error, got %v", err)
}
}
func TestMistralEmbedRejectsOutOfRangeIndex(t *testing.T) {
srv := newMistralServer(t, "/embeddings", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"data": []map[string]interface{}{
{"embedding": []float64{1}, "index": 7}, // out of range for 2-input request
},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
model := "mistral-embed"
_, err := m.Embed(&model, []string{"a", "b"}, &APIConfig{ApiKey: &apiKey}, nil)
if err == nil || !strings.Contains(err.Error(), "out of range") {
t.Errorf("expected out-of-range error, got %v", err)
}
}
func TestMistralEmbedRejectsMissingSlot(t *testing.T) {
// Upstream returns only one of the two requested embeddings.
srv := newMistralServer(t, "/embeddings", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"data": []map[string]interface{}{
{"embedding": []float64{1}, "index": 0},
},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
model := "mistral-embed"
_, err := m.Embed(&model, []string{"a", "b"}, &APIConfig{ApiKey: &apiKey}, nil)
if err == nil || !strings.Contains(err.Error(), "missing embedding for input index 1") {
t.Errorf("expected missing-embedding error for slot 1, got %v", err)
}
}
func TestMistralEmbedRejectsHTTPError(t *testing.T) {
srv := newMistralServer(t, "/embeddings", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error":"unauthorized"}`))
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
model := "mistral-embed"
_, err := m.Embed(&model, []string{"a"}, &APIConfig{ApiKey: &apiKey}, nil)
if err == nil || !strings.Contains(err.Error(), "Mistral embeddings API error") {
t.Errorf("expected Mistral embeddings API error, got %v", err)
}
}
fix(mistral): handle structured content from magistral reasoning models (#14805) ### What problem does this PR solve? `MistralModel.ChatWithMessages` (in the driver merged via #14807) assumes that `choices[0].message.content` from `/v1/chat/completions` is always a string and falls through to `return nil, fmt.Errorf("invalid content format")` on anything else. That assumption breaks for the **magistral reasoning family** (`magistral-small-*`, `magistral-medium-*`). When the model needs a chain-of-thought to answer, Mistral returns `content` as a **structured array of typed parts**: ```json "content": [ {"type": "thinking", "thinking": [{"type": "text", "text": "Combined speed is 150 mph. 300 / 150 = 2 hours."}], "closed": true}, {"type": "text", "text": "They will meet after **2 hours**."} ] ``` Concretely, this is what the live API returns today (probed against `api.mistral.ai/v1`): ``` $ curl -H "Authorization: Bearer <key>" -H "Content-Type: application/json" \ -X POST https://api.mistral.ai/v1/chat/completions \ -d '{"model":"magistral-medium-latest", "messages":[{"role":"user","content":"two trains 60mph and 90mph, 300mi apart, when do they meet? step by step."}], "max_tokens":1024}' HTTP 200 { "choices":[{"message":{ "role":"assistant", "content":[ {"type":"thinking","thinking":[{"type":"text","text":"Okay, let's see..."}],"closed":true}, {"type":"text","text":"To determine when the two trains meet..."} ]}}] } ``` With the current driver, every call like that returns the generic `"invalid content format"` error. Trivial prompts that happen to fit in a string answer still succeed, so the breakage is **non-deterministic from the tenant's POV**: same model, same provider, sometimes works, sometimes 500s with no useful error. A secondary issue: `conf/models/mistral.json` does not include any magistral model. The picker hid the broken path, which is why this wasn't caught during #14807's review. ### What this PR includes - New helper `extractMistralContent(raw interface{}) (answer, reasonContent string, err error)` in `internal/entity/models/mistral.go`, which normalizes both shapes Mistral can return: - `string` → historical path. `Answer = content`, `ReasonContent = ""`. Preserves behavior for every non-reasoning model (`mistral-large-*`, `mistral-small-*`, `ministral-*`, `codestral-*`, `pixtral-*`, `open-mistral-nemo`). - `[]interface{}` → walk the parts. Concatenate every `{"type":"text", "text":...}` part into `Answer`; concatenate the inner text inside every `{"type":"thinking", "thinking":[...]}` part into `ReasonContent`. - `ChatWithMessages` now calls the helper instead of doing the raw `.(string)` cast. - Unknown part types are **skipped, not failed**. Mistral has been adding new content variants quickly (audio chunks, citations, etc.); this driver should not 500 every call when a new part type appears. - `conf/models/mistral.json`: add `magistral-medium-latest` and `magistral-small-latest`. Both are visible in `/v1/models` today. No interface change. No factory change. No new dependencies. ### How was this tested? **Unit tests** — 5 new tests in `internal/entity/models/mistral_test.go` on top of the 27 already shipped via #14807: - `TestMistralChatHandlesStringContent` — regression net for the historical path - `TestMistralChatExtractsReasoningFromStructuredContent` — the fixture body is a trimmed copy of the actual `magistral-medium-latest` response captured above; asserts both `Answer` and `ReasonContent` are populated correctly - `TestMistralChatHandlesStructuredContentWithoutThinking` — `magistral-*` with a trivial answer returns a structured shape that has only a `text` part; `ReasonContent` must stay empty - `TestMistralChatIgnoresUnknownContentPartTypes` — `audio_url` and `future_part_type` parts are skipped, `text` parts still flow through - `TestExtractMistralContent` — table-driven unit coverage of the helper for string, empty string, nil, empty array, text-only, thinking+text, unsupported root type ``` $ go test -vet=off -run "TestMistral|TestExtractMistralContent" -count=1 -v ./internal/entity/models/... === RUN TestMistralChatHandlesStringContent --- PASS: TestMistralChatHandlesStringContent (0.00s) === RUN TestMistralChatExtractsReasoningFromStructuredContent --- PASS: TestMistralChatExtractsReasoningFromStructuredContent (0.00s) === RUN TestMistralChatHandlesStructuredContentWithoutThinking --- PASS: TestMistralChatHandlesStructuredContentWithoutThinking (0.00s) === RUN TestMistralChatIgnoresUnknownContentPartTypes --- PASS: TestMistralChatIgnoresUnknownContentPartTypes (0.00s) === RUN TestExtractMistralContent === RUN TestExtractMistralContent/plain_string === RUN TestExtractMistralContent/empty_string === RUN TestExtractMistralContent/nil === RUN TestExtractMistralContent/empty_array === RUN TestExtractMistralContent/text_only === RUN TestExtractMistralContent/thinking_then_text === RUN TestExtractMistralContent/unknown_root_type --- PASS: TestExtractMistralContent (0.00s) PASS ok ragflow/internal/entity/models 0.046s ``` All 32 Mistral tests pass on go 1.25. `go build ./internal/entity/models/...` exits 0. **Live integration test** — driver exercised against `api.mistral.ai/v1` with the patched code: ``` === RUN TestMistralMagistralSmoke [OK] "magistral-small-latest" present upstream [OK] "magistral-medium-latest" present upstream [OK trivial] Answer="7" ReasonContent="" [OK reasoning] Answer len=797 head="To determine when the two trains meet, we can follow these steps:\n\n1. **Identify..." ReasonContent len=1069 head="Okay, let's see. There are two trains, one going 60 mph and the other going 90 mph. They're moving towards each other, s..." MAGISTRAL SMOKE PASSED --- PASS: TestMistralMagistralSmoke (18.09s) PASS ok ragflow/internal/entity/models 18.112s ``` What the live run proves on the wire: - `magistral-small-latest` with a trivial prompt still uses the string-content shape; the regression-net path is exercised against the real server, not just the mock. - `magistral-medium-latest` with a reasoning prompt uses the structured-array shape; the new code path extracts a 1069-character reasoning trace into `ChatResponse.ReasonContent` and a 797-character visible answer into `ChatResponse.Answer`. Before this fix, the same call returned `"invalid content format"` and the caller saw nothing. The smoke-test file itself is not committed (live tests live outside the PR diff, same convention used for prior provider PRs). ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue)
2026-05-20 21:19:04 -10:00
// --- structured-content (magistral reasoning) tests ---
// Regression net: the existing string-content path stays green for every
// non-reasoning Mistral model.
func TestMistralChatHandlesStringContent(t *testing.T) {
srv := newMistralServer(t, "/chat/completions", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"choices": []map[string]interface{}{{"message": map[string]interface{}{
"role": "assistant",
"content": "Pong.",
}}},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
resp, err := m.ChatWithMessages("ministral-3b-latest",
[]Message{{Role: "user", Content: "ping"}},
&APIConfig{ApiKey: &apiKey}, nil)
if err != nil {
t.Fatalf("Chat: %v", err)
}
if *resp.Answer != "Pong." {
t.Errorf("Answer=%q want %q", *resp.Answer, "Pong.")
}
if *resp.ReasonContent != "" {
t.Errorf("ReasonContent=%q want empty", *resp.ReasonContent)
}
}
// New path: magistral with a non-trivial answer returns a structured
// content array. Two part types — `thinking` and `text` — must be
// concatenated into ReasonContent and Answer respectively.
//
// The fixture body is a trimmed copy of the actual response captured
// from api.mistral.ai/v1/chat/completions against magistral-medium-latest
// with the prompt "When do two trains meet?".
func TestMistralChatExtractsReasoningFromStructuredContent(t *testing.T) {
srv := newMistralServer(t, "/chat/completions", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"choices": []map[string]interface{}{{"message": map[string]interface{}{
"role": "assistant",
"content": []map[string]interface{}{
{
"type": "thinking",
"thinking": []map[string]interface{}{
{"type": "text", "text": "Combined speed is 150 mph. "},
{"type": "text", "text": "300 / 150 = 2 hours."},
},
"closed": true,
},
{"type": "text", "text": "They will meet after **2 hours**."},
},
}}},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
resp, err := m.ChatWithMessages("magistral-medium-latest",
[]Message{{Role: "user", Content: "When do they meet?"}},
&APIConfig{ApiKey: &apiKey}, nil)
if err != nil {
t.Fatalf("Chat: %v", err)
}
wantAnswer := "They will meet after **2 hours**."
wantReason := "Combined speed is 150 mph. 300 / 150 = 2 hours."
if *resp.Answer != wantAnswer {
t.Errorf("Answer=%q want %q", *resp.Answer, wantAnswer)
}
if *resp.ReasonContent != wantReason {
t.Errorf("ReasonContent=%q want %q", *resp.ReasonContent, wantReason)
}
}
// magistral with a trivial answer that needed no reasoning returns the
// structured shape with only a `text` part. ReasonContent must be empty.
func TestMistralChatHandlesStructuredContentWithoutThinking(t *testing.T) {
srv := newMistralServer(t, "/chat/completions", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"choices": []map[string]interface{}{{"message": map[string]interface{}{
"role": "assistant",
"content": []map[string]interface{}{
{"type": "text", "text": "12"},
},
}}},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
resp, err := m.ChatWithMessages("magistral-small-latest",
[]Message{{Role: "user", Content: "15% of 80?"}},
&APIConfig{ApiKey: &apiKey}, nil)
if err != nil {
t.Fatalf("Chat: %v", err)
}
if *resp.Answer != "12" {
t.Errorf("Answer=%q want %q", *resp.Answer, "12")
}
if *resp.ReasonContent != "" {
t.Errorf("ReasonContent=%q want empty (no thinking part)", *resp.ReasonContent)
}
}
// Unknown part types must be skipped, not crash the driver. This makes
// the parser forward-compatible with new Mistral content variants
// (audio chunks, citations, etc.) that ragflow doesn't surface yet.
func TestMistralChatIgnoresUnknownContentPartTypes(t *testing.T) {
srv := newMistralServer(t, "/chat/completions", func(t *testing.T, _ map[string]interface{}, w http.ResponseWriter) {
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"choices": []map[string]interface{}{{"message": map[string]interface{}{
"role": "assistant",
"content": []map[string]interface{}{
{"type": "audio_url", "audio_url": "ignored://x"},
{"type": "text", "text": "Hello"},
{"type": "future_part_type", "blob": "?"},
},
}}},
})
})
defer srv.Close()
m := newMistralForTest(srv.URL)
apiKey := "test-key"
resp, err := m.ChatWithMessages("magistral-small-latest",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{ApiKey: &apiKey}, nil)
if err != nil {
t.Fatalf("Chat: %v", err)
}
if *resp.Answer != "Hello" {
t.Errorf("Answer=%q want %q", *resp.Answer, "Hello")
}
}
// Direct unit coverage of the helper, including the nil and bad-type
// edge cases that won't surface in the integration tests above.
func TestExtractMistralContent(t *testing.T) {
tests := []struct {
name string
input interface{}
wantAns string
wantReason string
wantErr bool
}{
{"plain string", "hi", "hi", "", false},
{"empty string", "", "", "", false},
{"nil", nil, "", "", false},
{"empty array", []interface{}{}, "", "", false},
{
"text only",
[]interface{}{
map[string]interface{}{"type": "text", "text": "a"},
map[string]interface{}{"type": "text", "text": "b"},
},
"ab", "", false,
},
{
"thinking then text",
[]interface{}{
map[string]interface{}{
"type": "thinking",
"thinking": []interface{}{
map[string]interface{}{"type": "text", "text": "why "},
map[string]interface{}{"type": "text", "text": "this"},
},
},
map[string]interface{}{"type": "text", "text": "answer"},
},
"answer", "why this", false,
},
{"unknown root type", 42, "", "", true},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
ans, reason, err := extractMistralContent(tc.input)
if tc.wantErr {
if err == nil {
t.Errorf("want error, got nil")
}
return
}
if err != nil {
t.Errorf("unexpected err: %v", err)
}
if ans != tc.wantAns {
t.Errorf("answer=%q want %q", ans, tc.wantAns)
}
if reason != tc.wantReason {
t.Errorf("reason=%q want %q", reason, tc.wantReason)
}
})
}
}