mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-10-12 02:16:39 +00:00
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:
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"`
|
||||
}
|
||||
|
||||
@@ -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"]),
|
||||
|
||||
Reference in New Issue
Block a user