diff --git a/conf/models/vllm.json b/conf/models/vllm.json index 9c6a440a87..d4f074330b 100644 --- a/conf/models/vllm.json +++ b/conf/models/vllm.json @@ -3,7 +3,8 @@ "url_suffix": { "chat": "chat/completions", "models": "models", - "embedding": "embeddings" + "embedding": "embeddings", + "rerank": "rerank" }, "class": "local" } \ No newline at end of file diff --git a/internal/entity/models/vllm.go b/internal/entity/models/vllm.go index 4780fee939..61452c9bca 100644 --- a/internal/entity/models/vllm.go +++ b/internal/entity/models/vllm.go @@ -547,9 +547,120 @@ func (z *VllmModel) CheckConnection(apiConfig *APIConfig) error { return err } -// Rerank calculates similarity scores between query and documents +// vllmRerankRequest mirrors vLLM's Jina/Cohere-compatible /v1/rerank +// payload. Unlike NVIDIA NIM (which wraps each passage as {text: "..."}), +// vLLM accepts documents as a flat []string. +type vllmRerankRequest struct { + Model string `json:"model"` + Query string `json:"query"` + Documents []string `json:"documents"` + TopN int `json:"top_n"` +} + +// vllmRerankResponse maps the Jina-style results array. The `document` +// field is intentionally ignored — callers reconstruct text from the +// original input via Index. +type vllmRerankResponse struct { + Results []struct { + Index int `json:"index"` + RelevanceScore float64 `json:"relevance_score"` + } `json:"results"` +} + +// Rerank scores documents against the query using a vLLM rerank model +// served at /v1/rerank (stable since vLLM v0.7). Mirrors the contract +// of NvidiaModel.Rerank: defaults top_n to len(documents) so callers +// get a score per input, shrinks to RerankConfig.TopN only when set +// and smaller. Returned RerankResult entries are in the API's ranking +// order; callers that need original-input order sort by Index. The +// Authorization header is sent only when APIConfig.ApiKey is non-empty, +// matching the existing Embed/ListModels behaviour for this local +// driver. func (z *VllmModel) Rerank(modelName *string, query string, documents []string, apiConfig *APIConfig, rerankConfig *RerankConfig) (*RerankResponse, error) { - return nil, fmt.Errorf("%s, Rerank not implemented", z.Name()) + if len(documents) == 0 { + return &RerankResponse{}, nil + } + if modelName == nil || *modelName == "" { + return nil, fmt.Errorf("model name is required") + } + + region := "default" + if apiConfig != nil && apiConfig.Region != nil && *apiConfig.Region != "" { + region = *apiConfig.Region + } + + baseURL := z.BaseURL[region] + if baseURL == "" { + baseURL = z.BaseURL["default"] + } + if baseURL == "" { + return nil, fmt.Errorf("missing base URL: please configure the local access address for vLLM (e.g., http://127.0.0.1:8000/v1)") + } + + url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), z.URLSuffix.Rerank) + + topN := len(documents) + if rerankConfig != nil && rerankConfig.TopN > 0 && rerankConfig.TopN < topN { + topN = rerankConfig.TopN + } + + reqBody := vllmRerankRequest{ + Model: *modelName, + Query: query, + Documents: documents, + TopN: topN, + } + + jsonData, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal request: %w", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + req.Header.Set("Content-Type", "application/json") + if apiConfig != nil && apiConfig.ApiKey != nil && *apiConfig.ApiKey != "" { + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) + } + + resp, err := z.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("vLLM rerank API error: %s, body: %s", resp.Status, string(body)) + } + + var parsed vllmRerankResponse + if err = json.Unmarshal(body, &parsed); err != nil { + return nil, fmt.Errorf("failed to parse response: %w", err) + } + + rerankResponse := RerankResponse{Data: make([]RerankResult, 0, len(parsed.Results))} + for _, r := range parsed.Results { + if r.Index < 0 || r.Index >= len(documents) { + return nil, fmt.Errorf("unexpected rerank index %d for %d inputs", r.Index, len(documents)) + } + rerankResponse.Data = append(rerankResponse.Data, RerankResult{ + Index: r.Index, + RelevanceScore: r.RelevanceScore, + }) + } + + return &rerankResponse, nil } // TranscribeAudio transcribe audio diff --git a/internal/entity/models/vllm_rerank_test.go b/internal/entity/models/vllm_rerank_test.go new file mode 100644 index 0000000000..42fda948c2 --- /dev/null +++ b/internal/entity/models/vllm_rerank_test.go @@ -0,0 +1,209 @@ +package models + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func newVllmRerankServer(t *testing.T, expectAuth 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.Method != http.MethodPost { + t.Errorf("expected POST, got %s", r.Method) + return + } + if r.URL.Path != "/rerank" { + t.Errorf("expected path=/rerank, got %s", r.URL.Path) + return + } + if got := r.Header.Get("Authorization"); got != expectAuth { + t.Errorf("expected Authorization=%q, got %q", expectAuth, got) + return + } + 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) + })) +} + +func newVllmModelForTest(baseURL string) *VllmModel { + return NewVllmModel( + map[string]string{"default": baseURL}, + URLSuffix{Rerank: "rerank"}, + ) +} + +func TestVllmRerankHappyPath(t *testing.T) { + srv := newVllmRerankServer(t, "Bearer test-key", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + if body["model"] != "BAAI/bge-reranker-v2-m3" { + t.Errorf("expected model=BAAI/bge-reranker-v2-m3, got %v", body["model"]) + } + if body["query"] != "What is RAPTOR?" { + t.Errorf("expected query=What is RAPTOR?, got %v", body["query"]) + } + // vLLM differs from NVIDIA: documents is a flat []string, not [{text}]. + docs, ok := body["documents"].([]interface{}) + if !ok || len(docs) != 3 { + t.Errorf("expected 3 documents, got %v", body["documents"]) + return + } + for i, want := range []string{"doc-zero", "doc-one", "doc-two"} { + if docs[i] != want { + t.Errorf("documents[%d]=%v, want %s", i, docs[i], want) + } + } + if body["top_n"] != float64(3) { + t.Errorf("expected top_n=3 (matching len(documents)), got %v", body["top_n"]) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "results": []map[string]interface{}{ + {"index": 2, "relevance_score": 0.95}, + {"index": 0, "relevance_score": 0.42}, + {"index": 1, "relevance_score": 0.78}, + }, + }) + }) + defer srv.Close() + + model := newVllmModelForTest(srv.URL) + apiKey := "test-key" + modelName := "BAAI/bge-reranker-v2-m3" + resp, err := model.Rerank( + &modelName, + "What is RAPTOR?", + []string{"doc-zero", "doc-one", "doc-two"}, + &APIConfig{ApiKey: &apiKey}, + &RerankConfig{}, + ) + if err != nil { + t.Fatalf("Rerank failed: %v", err) + } + if len(resp.Data) != 3 { + t.Fatalf("expected 3 results, got %d", len(resp.Data)) + } + want := map[int]float64{0: 0.42, 1: 0.78, 2: 0.95} + for _, r := range resp.Data { + if got, ok := want[r.Index]; !ok || got != r.RelevanceScore { + t.Errorf("unexpected result Index=%d RelevanceScore=%v", r.Index, r.RelevanceScore) + } + } +} + +func TestVllmRerankTopNClamp(t *testing.T) { + srv := newVllmRerankServer(t, "Bearer test-key", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + if body["top_n"] != float64(2) { + t.Errorf("expected top_n clamp to RerankConfig.TopN=2, got %v", body["top_n"]) + } + _ = json.NewEncoder(w).Encode(map[string]interface{}{"results": []map[string]interface{}{}}) + }) + defer srv.Close() + + model := newVllmModelForTest(srv.URL) + apiKey := "test-key" + modelName := "BAAI/bge-reranker-v2-m3" + if _, err := model.Rerank( + &modelName, "q", + []string{"a", "b", "c", "d"}, + &APIConfig{ApiKey: &apiKey}, + &RerankConfig{TopN: 2}, + ); err != nil { + t.Fatalf("Rerank failed: %v", err) + } +} + +func TestVllmRerankEmptyDocuments(t *testing.T) { + model := newVllmModelForTest("http://unused") + apiKey := "test-key" + modelName := "BAAI/bge-reranker-v2-m3" + resp, err := model.Rerank(&modelName, "q", nil, &APIConfig{ApiKey: &apiKey}, &RerankConfig{}) + if err != nil { + t.Fatalf("expected nil error for empty documents, got %v", err) + } + if len(resp.Data) != 0 { + t.Errorf("expected empty Data, got %d entries", len(resp.Data)) + } +} + +// vLLM is a local driver; the Authorization header must be omitted when +// no APIConfig.ApiKey is configured. This diverges from the NVIDIA driver +// which requires an API key. +func TestVllmRerankWithoutAPIKey(t *testing.T) { + srv := newVllmRerankServer(t, "", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "results": []map[string]interface{}{ + {"index": 0, "relevance_score": 0.5}, + }, + }) + }) + defer srv.Close() + + model := newVllmModelForTest(srv.URL) + modelName := "BAAI/bge-reranker-v2-m3" + resp, err := model.Rerank(&modelName, "q", []string{"a"}, &APIConfig{}, &RerankConfig{}) + if err != nil { + t.Fatalf("Rerank failed without api key: %v", err) + } + if len(resp.Data) != 1 || resp.Data[0].Index != 0 { + t.Errorf("unexpected response: %+v", resp) + } +} + +func TestVllmRerankRequiresModelName(t *testing.T) { + model := newVllmModelForTest("http://unused") + apiKey := "test-key" + _, err := model.Rerank(nil, "q", []string{"a"}, &APIConfig{ApiKey: &apiKey}, &RerankConfig{}) + if err == nil || !strings.Contains(err.Error(), "model name is required") { + t.Errorf("expected model-name error, got %v", err) + } +} + +func TestVllmRerankRejectsHTTPError(t *testing.T) { + srv := newVllmRerankServer(t, "Bearer test-key", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error":"boom"}`)) + }) + defer srv.Close() + + model := newVllmModelForTest(srv.URL) + apiKey := "test-key" + modelName := "BAAI/bge-reranker-v2-m3" + _, err := model.Rerank(&modelName, "q", []string{"a"}, &APIConfig{ApiKey: &apiKey}, &RerankConfig{}) + if err == nil || !strings.Contains(err.Error(), "vLLM rerank API error") { + t.Errorf("expected API error, got %v", err) + } +} + +func TestVllmRerankRejectsOutOfRangeIndex(t *testing.T) { + srv := newVllmRerankServer(t, "Bearer test-key", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "results": []map[string]interface{}{ + {"index": 5, "relevance_score": 0.9}, + }, + }) + }) + defer srv.Close() + + model := newVllmModelForTest(srv.URL) + apiKey := "test-key" + modelName := "BAAI/bge-reranker-v2-m3" + _, err := model.Rerank(&modelName, "q", []string{"a", "b"}, &APIConfig{ApiKey: &apiKey}, &RerankConfig{}) + if err == nil || !strings.Contains(err.Error(), "unexpected rerank index") { + t.Errorf("expected out-of-range error, got %v", err) + } +}