Fix rerank model does not take effect when search in searchbot (#20158)

Support rerrank for searchbot; add highlight for searchbot.
This commit is contained in:
Wang Qi
2026-09-24 11:59:21 +08:00
committed by GitHub
parent 6fec6bba59
commit 243b9f7c51
5 changed files with 34 additions and 22 deletions
+8 -10
View File
@@ -64,13 +64,9 @@ type SearchBotRetrievalTestRequest struct {
TenantRerankID *string `json:"tenant_rerank_id,omitempty"`
RerankID *string `json:"rerank_id,omitempty"`
Keyword *bool `json:"keyword,omitempty"`
Highlight *bool `json:"highlight,omitempty"`
SimilarityThreshold *float64 `json:"similarity_threshold,omitempty"`
VectorSimilarityWeight *float64 `json:"vector_similarity_weight,omitempty"`
// TODO: wire highlight to nlp Retrieval when engine supports highlightFields
// Python: bot_api.py → retrieval(highlight=req.get("highlight"))
// → search.py highlightFields → ES get_highlight()
// Issue: https://github.com/infiniflow/ragflow/issues/15712
// Highlight *bool `json:"highlight,omitempty"`
}
// UnmarshalJSON accepts both kb_id (Python API) and kb_ids (Go compatibility).
@@ -261,13 +257,14 @@ func (h *SearchBotHandler) Ask(c *gin.Context) {
// Resolve chat model ID.
modelID := ""
options := service.AskStreamOptions{}
if req.SearchID != "" && h.searchSvc != nil {
ctx := c.Request.Context()
if detail, err := h.searchSvc.GetDetail(ctx, req.SearchID); err == nil {
if sc, ok := detail["search_config"].(map[string]interface{}); ok {
if cid, ok := sc["chat_id"].(string); ok && cid != "" {
modelID = cid
}
searchConfig := searchConfigFromDetail(detail)
options = service.BuildAskStreamOptions(req.SearchID, searchConfig)
if chatID, ok := searchConfig["chat_id"].(string); ok && chatID != "" {
modelID = chatID
}
}
}
@@ -298,7 +295,7 @@ func (h *SearchBotHandler) Ask(c *gin.Context) {
}
ctx := c.Request.Context()
adapter := &service.TenantStreamAdapter{LLM: h.streamLLM, TenantID: user.ID, ModelID: modelID}
for delta := range h.askSvc.Stream(ctx, adapter, user.ID, req.Question, filtered) {
for delta := range h.askSvc.StreamWithOptions(ctx, adapter, user.ID, req.Question, filtered, options) {
switch delta.Kind {
case service.AskDeltaAnswer:
h.sseWriter.Write(c, sseAnswer(delta.Value, nil, false))
@@ -529,6 +526,7 @@ func toRetrievalServiceRequest(h *SearchBotRetrievalTestRequest) *service.Retrie
TenantRerankID: h.TenantRerankID,
RerankID: h.RerankID,
Keyword: h.Keyword,
Highlight: h.Highlight,
SimilarityThreshold: h.SimilarityThreshold,
VectorSimilarityWeight: h.VectorSimilarityWeight,
}
+3 -3
View File
@@ -203,7 +203,7 @@ func TestAskService_TemperatureFromLLMSetting(t *testing.T) {
llm := &capturingStreamLLM{fakeStreamLLM: fakeStreamLLM{chunks: []string{"answer"}}}
svc := NewAskService(ret, nil, 0, 0)
opts := askOptionsFromSearchConfig("s1", map[string]interface{}{
opts := BuildAskStreamOptions("s1", map[string]interface{}{
"llm_setting": map[string]interface{}{
"temperature": 0.6,
"temperature_enabled": true,
@@ -246,7 +246,7 @@ func TestAskService_TemperatureDisabledUsesDefault(t *testing.T) {
llm := &capturingStreamLLM{fakeStreamLLM: fakeStreamLLM{chunks: []string{"answer"}}}
svc := NewAskService(ret, nil, 0, 0)
opts := askOptionsFromSearchConfig("s1", map[string]interface{}{
opts := BuildAskStreamOptions("s1", map[string]interface{}{
"llm_setting": map[string]interface{}{
"temperature": 0.6,
"temperature_enabled": false,
@@ -280,7 +280,7 @@ func TestAskService_TemperatureMissingUsesLLMSettingDefaults(t *testing.T) {
// Flag enabled but value absent (and a completely absent flag) must
// substitute the Python LLM_SETTING_DEFAULTS, matching resolve_llm_setting.
opts := askOptionsFromSearchConfig("s1", map[string]interface{}{
opts := BuildAskStreamOptions("s1", map[string]interface{}{
"llm_setting": map[string]interface{}{
"temperature_enabled": true,
},
+19 -7
View File
@@ -216,6 +216,8 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval
var chatID string
var chatModelForFilter *models.ChatModel
filter := req.Filter
tenantRerankID := req.TenantRerankID
rerankID := req.RerankID
rerankCandidatesCount := 64
if req.RerankCandidatesCount != nil {
rerankCandidatesCount = *req.RerankCandidatesCount
@@ -236,6 +238,16 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval
if searchMetaFilter, ok := searchConfigMap(searchConfig["meta_data_filter"]); ok {
filter = searchMetaFilter
}
if tenantRerankID == nil || strings.TrimSpace(*tenantRerankID) == "" {
if configuredTenantRerankID, ok := searchConfig["tenant_rerank_id"].(string); ok && strings.TrimSpace(configuredTenantRerankID) != "" {
tenantRerankID = &configuredTenantRerankID
}
}
if rerankID == nil || strings.TrimSpace(*rerankID) == "" {
if configuredRerankID, ok := searchConfig["rerank_id"].(string); ok && strings.TrimSpace(configuredRerankID) != "" {
rerankID = &configuredRerankID
}
}
chatID, _ = searchConfig["chat_id"].(string)
} else {
common.Warn("No search_config found in search detail", zap.String("searchID", *req.SearchID))
@@ -398,20 +410,19 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval
// Get rerank model if RerankID is specified
var rerankModel *models.RerankModel
if req.TenantRerankID != nil && *req.TenantRerankID != "" {
target, getErr := modelSolver.ResolveModelConfig(ctx, tenantIDs[0], entity.ModelTypeRerank, *req.TenantRerankID)
if tenantRerankID != nil && *tenantRerankID != "" {
target, getErr := modelSolver.ResolveModelConfig(ctx, tenantIDs[0], entity.ModelTypeRerank, *tenantRerankID)
if getErr != nil {
return nil, fmt.Errorf("failed to get rerank model by tenant_rerank_id: %w", getErr)
}
rerankModel = models.NewRerankModel(target.Driver, &target.ModelName, target.APIConfig, target.MaxTokens)
} else if req.RerankID != nil && *req.RerankID != "" {
rerankCompositeName := *req.RerankID
} else if rerankID != nil && *rerankID != "" {
rerankCompositeName := *rerankID
target, getErr := modelSolver.ResolveModelConfig(ctx, tenantIDs[0], entity.ModelTypeRerank, rerankCompositeName)
if getErr != nil {
rerankModel = nil
} else {
rerankModel = models.NewRerankModel(target.Driver, &target.ModelName, target.APIConfig, target.MaxTokens)
return nil, fmt.Errorf("failed to get rerank model by rerank_id: %w", getErr)
}
rerankModel = models.NewRerankModel(target.Driver, &target.ModelName, target.APIConfig, target.MaxTokens)
}
retrievalReq := &nlp.RetrievalRequest{
@@ -428,6 +439,7 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval
RerankModel: rerankModel,
RankFeature: &labels,
EmbeddingModel: embeddingModel,
Highlight: req.Highlight,
}
// Call RetrievalService to perform retrieval
+1
View File
@@ -39,6 +39,7 @@ type RetrievalTestRequest struct {
TenantRerankID *string `json:"tenant_rerank_id,omitempty"`
RerankID *string `json:"rerank_id,omitempty"`
Keyword *bool `json:"keyword,omitempty"`
Highlight *bool `json:"highlight,omitempty"`
SimilarityThreshold *float64 `json:"similarity_threshold,omitempty"`
VectorSimilarityWeight *float64 `json:"vector_similarity_weight,omitempty"`
}
+3 -2
View File
@@ -423,11 +423,12 @@ func (s *SearchService) PrepareCompletion(ctx context.Context, userID, searchID
Question: question,
DatasetIDs: datasetIDs,
ModelID: modelID,
Options: askOptionsFromSearchConfig(searchID, searchConfig),
Options: BuildAskStreamOptions(searchID, searchConfig),
}, common.CodeSuccess, nil
}
func askOptionsFromSearchConfig(searchID string, searchConfig map[string]interface{}) AskStreamOptions {
// BuildAskStreamOptions maps a saved search configuration to Ask retrieval options.
func BuildAskStreamOptions(searchID string, searchConfig map[string]interface{}) AskStreamOptions {
opts := AskStreamOptions{
SearchID: searchID,
DocIDs: stringSliceFromSearchConfig(searchConfig["doc_ids"]),