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
This commit is contained in:
Martín Alcalá Rubí
committed by
lyingbug
lyingbug
parent
e289595054
commit
83da0903b6
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user