From 83da0903b6f7be5ca02fa3e3fa0ec6e8ded09008 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mart=C3=ADn=20Alcal=C3=A1=20Rub=C3=AD?= Date: Mon, 20 Jul 2026 23:38:27 -0300 Subject: [PATCH] fix(rerank): stop sending vLLM truncate_prompt_tokens by default The generic OpenAI-compatible reranker hardcoded `truncate_prompt_tokens: 511` on every /rerank request. Providers that honor this vLLM-specific field (e.g. SiliconFlow serving Qwen3-Reranker) keep only the last 511 tokens of the templated rerank prompt, cutting the instruction and query off any document longer than ~511 tokens. Every relevance score then collapses to near zero (~0.0067 instead of ~0.998), falls below the score threshold, and knowledge-base search returns empty results. A direct curl without the field returns the correct scores. WeKnora already removed this exact hardcoded parameter from the embedding path in #960; the rerank path was never fixed. - Add `omitempty` to `truncate_prompt_tokens` and `additional_data` so empty optional fields are not sent. - Make truncation opt-in via `extra_config.truncate_prompt_tokens` (validated > 0 at construction), preserving an escape hatch for self-hosted vLLM deployments with small-context rerankers. - Add regression tests reproducing the collapsed-score symptom and the opt-in path. Fixes #2143 --- internal/models/rerank/remote_api.go | 48 +++++-- internal/models/rerank/remote_api_test.go | 157 ++++++++++++++++++++++ 2 files changed, 193 insertions(+), 12 deletions(-) create mode 100644 internal/models/rerank/remote_api_test.go diff --git a/internal/models/rerank/remote_api.go b/internal/models/rerank/remote_api.go index 105223ae..fc0feb14 100644 --- a/internal/models/rerank/remote_api.go +++ b/internal/models/rerank/remote_api.go @@ -7,6 +7,8 @@ import ( "fmt" "io" "net/http" + "strconv" + "strings" "github.com/Tencent/WeKnora/internal/logger" secutils "github.com/Tencent/WeKnora/internal/utils" @@ -20,6 +22,12 @@ type OpenAIReranker struct { baseURL string // Base URL for API requests client *http.Client // HTTP client for making API requests customHeaders map[string]string + // truncatePromptTokens, when > 0, is sent as the vLLM-specific + // truncate_prompt_tokens request field. It must never be sent by default: + // providers that honor it (e.g. SiliconFlow) keep only the LAST N tokens of + // the templated rerank prompt, which cuts the query off long documents and + // collapses every relevance score to near zero (issue #2143). + truncatePromptTokens int } // SetCustomHeaders 设置用户自定义 HTTP 请求头（类似 OpenAI Python SDK 的 extra_headers）。 @@ -29,11 +37,11 @@ func (r *OpenAIReranker) SetCustomHeaders(headers map[string]string) { // RerankRequest represents a request to rerank documents based on relevance to a query type RerankRequest struct { - Model string `json:"model"` // Model to use for reranking - Query string `json:"query"` // Query text to compare documents against - Documents []string `json:"documents"` // List of document texts to rerank - AdditionalData map[string]interface{} `json:"additional_data"` // Optional additional data for the model - TruncatePromptTokens int `json:"truncate_prompt_tokens"` // Maximum prompt tokens to use + Model string `json:"model"` // Model to use for reranking + Query string `json:"query"` // Query text to compare documents against + Documents []string `json:"documents"` // List of document texts to rerank + AdditionalData map[string]interface{} `json:"additional_data,omitempty"` // Optional additional data for the model + TruncatePromptTokens int `json:"truncate_prompt_tokens,omitempty"` // Maximum prompt tokens to use (vLLM-specific, opt-in) } // RerankResponse represents the response from a reranking request @@ -60,23 +68,39 @@ func NewOpenAIReranker(config *RerankerConfig) (*OpenAIReranker, error) { return nil, err } + // Optional opt-in for vLLM-style deployments that need server-side prompt + // truncation. Configured via extra_config; never enabled by default. + truncatePromptTokens := 0 + if config.ExtraConfig != nil { + if raw := strings.TrimSpace(config.ExtraConfig["truncate_prompt_tokens"]); raw != "" { + n, err := strconv.Atoi(raw) + if err != nil || n <= 0 { + return nil, fmt.Errorf("invalid truncate_prompt_tokens in extra_config: %q", raw) + } + truncatePromptTokens = n + } + } + return &OpenAIReranker{ - modelName: config.ModelName, - modelID: config.ModelID, - apiKey: apiKey, - baseURL: baseURL, - client: newRerankHTTPClient(0), + modelName: config.ModelName, + modelID: config.ModelID, + apiKey: apiKey, + baseURL: baseURL, + client: newRerankHTTPClient(0), + truncatePromptTokens: truncatePromptTokens, }, nil } // Rerank performs document reranking based on relevance to the query func (r *OpenAIReranker) Rerank(ctx context.Context, query string, documents []string) ([]RankResult, error) { - // Build the request body + // Build the request body. truncate_prompt_tokens is only included when + // explicitly configured: sending it unconditionally corrupts scores on + // providers that honor it (see OpenAIReranker.truncatePromptTokens). requestBody := &RerankRequest{ Model: r.modelName, Query: query, Documents: documents, - TruncatePromptTokens: 511, + TruncatePromptTokens: r.truncatePromptTokens, } jsonData, err := json.Marshal(requestBody) diff --git a/internal/models/rerank/remote_api_test.go b/internal/models/rerank/remote_api_test.go new file mode 100644 index 00000000..ef4fa548 --- /dev/null +++ b/internal/models/rerank/remote_api_test.go @@ -0,0 +1,157 @@ +package rerank + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" +) + +// newRerankScoreTestServer emulates an OpenAI-compatible /rerank endpoint the +// way a vLLM-backed provider (e.g. SiliconFlow) behaves, and records the last +// decoded request body. +// +// If the request carries truncate_prompt_tokens, the backend keeps only the +// last N tokens of the templated rerank prompt — the query gets cut off long +// documents and every relevance score collapses to near zero (issue #2143). +// Otherwise it returns the real scores, sorted by relevance_score descending, +// with index pointing back at the original input-document position. +func newRerankScoreTestServer(t *testing.T, lastRequest *map[string]interface{}) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req map[string]interface{} + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Errorf("decode rerank request: %v", err) + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + *lastRequest = req + + w.Header().Set("Content-Type", "application/json") + if _, truncated := req["truncate_prompt_tokens"]; truncated { + // Query truncated away from the prompt: scores collapse. + _, _ = w.Write([]byte(`{ + "id": "rerank-collapsed", + "results": [ + {"index": 0, "relevance_score": 0.00670681, "document": {"text": "A"}}, + {"index": 1, "relevance_score": 0.00587342, "document": {"text": "B"}}, + {"index": 2, "relevance_score": 0.00412907, "document": {"text": "C"}} + ], + "usage": {"total_tokens": 42} + }`)) + return + } + // Healthy response: sorted by relevance_score descending, index values + // out of order relative to the input documents [A, B, C]. + _, _ = w.Write([]byte(`{ + "id": "rerank-ok", + "results": [ + {"index": 2, "relevance_score": 0.998, "document": {"text": "C"}}, + {"index": 0, "relevance_score": 0.51, "document": {"text": "A"}}, + {"index": 1, "relevance_score": 0.006, "document": {"text": "B"}} + ], + "usage": {"total_tokens": 42} + }`)) + })) +} + +// TestOpenAIRerankerDoesNotTruncatePromptByDefault is the regression test for +// issue #2143: the generic OpenAI-compatible reranker must not send the +// vLLM-specific truncate_prompt_tokens field unless explicitly configured, and +// each input document must keep its own relevance score via the index field. +func TestOpenAIRerankerDoesNotTruncatePromptByDefault(t *testing.T) { + withRerankSSRFWhitelist(t, "127.0.0.1") + + var lastRequest map[string]interface{} + server := newRerankScoreTestServer(t, &lastRequest) + defer server.Close() + + reranker, err := NewOpenAIReranker(&RerankerConfig{ + BaseURL: server.URL, + ModelName: "Qwen/Qwen3-VL-Reranker-8B", + APIKey: "sk-test", + }) + if err != nil { + t.Fatalf("NewOpenAIReranker: %v", err) + } + + documents := []string{"A", "B", "C"} + results, err := reranker.Rerank(t.Context(), "query", documents) + if err != nil { + t.Fatalf("Rerank: %v", err) + } + + if v, ok := lastRequest["truncate_prompt_tokens"]; ok { + t.Errorf("request contains truncate_prompt_tokens=%v; it must not be sent unless configured", v) + } + if v, ok := lastRequest["additional_data"]; ok { + t.Errorf("request contains additional_data=%v; empty optional fields must be omitted", v) + } + + if len(results) != len(documents) { + t.Fatalf("got %d results, want %d", len(results), len(documents)) + } + + // Results stay in the order returned by the API (sorted by score desc). + if results[0].Index != 2 || results[0].RelevanceScore != 0.998 { + t.Errorf("top result = {index: %d, score: %v}, want {index: 2, score: 0.998}", + results[0].Index, results[0].RelevanceScore) + } + + // Each input document keeps its own score, resolved through the index field. + wantScoreByDoc := map[string]float64{"C": 0.998, "A": 0.51, "B": 0.006} + for _, rr := range results { + if rr.Index < 0 || rr.Index >= len(documents) { + t.Fatalf("result index %d out of range for %d documents", rr.Index, len(documents)) + } + doc := documents[rr.Index] + if want := wantScoreByDoc[doc]; rr.RelevanceScore != want { + t.Errorf("document %q got score %v, want %v", doc, rr.RelevanceScore, want) + } + } +} + +// TestOpenAIRerankerTruncatePromptTokensOptIn verifies that deployments which +// really need server-side prompt truncation (self-hosted vLLM with small- +// context rerankers) can still opt in via extra_config. +func TestOpenAIRerankerTruncatePromptTokensOptIn(t *testing.T) { + withRerankSSRFWhitelist(t, "127.0.0.1") + + var lastRequest map[string]interface{} + server := newRerankScoreTestServer(t, &lastRequest) + defer server.Close() + + reranker, err := NewOpenAIReranker(&RerankerConfig{ + BaseURL: server.URL, + ModelName: "bge-reranker-base", + APIKey: "sk-test", + ExtraConfig: map[string]string{"truncate_prompt_tokens": "511"}, + }) + if err != nil { + t.Fatalf("NewOpenAIReranker: %v", err) + } + + if _, err := reranker.Rerank(t.Context(), "query", []string{"A", "B", "C"}); err != nil { + t.Fatalf("Rerank: %v", err) + } + + v, ok := lastRequest["truncate_prompt_tokens"] + if !ok { + t.Fatal("request is missing truncate_prompt_tokens despite extra_config opt-in") + } + if n, _ := v.(float64); n != 511 { + t.Errorf("truncate_prompt_tokens = %v, want 511", v) + } +} + +func TestNewOpenAIRerankerRejectsInvalidTruncatePromptTokens(t *testing.T) { + for _, raw := range []string{"abc", "-1", "0"} { + _, err := NewOpenAIReranker(&RerankerConfig{ + ModelName: "rerank-test", + ExtraConfig: map[string]string{"truncate_prompt_tokens": raw}, + }) + if err == nil { + t.Errorf("NewOpenAIReranker accepted invalid truncate_prompt_tokens %q", raw) + } + } +}