ragflow/internal/service/nlp/retrieval.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

package nlp

import (
	"context"
	"fmt"
	"math"
	"ragflow/internal/common"
	"ragflow/internal/dao"
	"ragflow/internal/engine"
	"ragflow/internal/engine/types"
	"ragflow/internal/entity/models"
	"sort"
	"strconv"
	"strings"

	"ragflow/internal/tokenizer"

	"go.uber.org/zap"
)

// RetrievalService provides retrieval search functionality
type RetrievalService struct {
	docEngine   engine.DocEngine
	documentDAO *dao.DocumentDAO
}

// NewRetrievalService creates a new RetrievalService with the given doc engine
func NewRetrievalService(docEngine engine.DocEngine, documentDAO *dao.DocumentDAO) *RetrievalService {
	return &RetrievalService{docEngine: docEngine, documentDAO: documentDAO}
}

// RetrievalRequest request for retrieval search
type RetrievalRequest struct {
	Question               string
	TenantIDs              []string
	KbIDs                  []string
	DocIDs                 []string
	Page                   int
	PageSize               int
	Top                    *int
	SimilarityThreshold    *float64
	VectorSimilarityWeight *float64
	RankFeature            *map[string]float64
	RerankModel            *models.RerankModel
	EmbeddingModel         *models.EmbeddingModel
	Aggs                   *bool
	Highlight              *bool
	Filter                 map[string]interface{}
}

// RetrievalResult result from retrieval search
type RetrievalResult struct {
	Chunks  []map[string]interface{}
	DocAggs []map[string]interface{} // Aggregated document counts, sorted by count desc
	Total   int64                    // Threshold-valid matches across the retrieval candidate set
}

// Retrieval performs hybrid search + reranking + pagination
// - Calculate rerank limit and call Search() to fetch rerankLimit candidates for reranking
// - Perform reranking via Rerank()
// - Sort indices by score descending and filter by threshold
// - Calculate pagination to extract actual page returned from reranked results
// - Build chunks
// - Build document aggregation if specified
func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) (*RetrievalResult, error) {
	common.Info("Retrieval START", zap.String("question", req.Question), zap.Int("page", req.Page), zap.Int("pageSize", req.PageSize))
	if req.Question == "" {
		return &RetrievalResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}, Total: 0}, nil
	}

	// Apply default values
	if req.Top == nil {
		req.Top = func() *int { v := 1024; return &v }()
	}
	if req.SimilarityThreshold == nil {
		req.SimilarityThreshold = func() *float64 { v := 0.2; return &v }()
	}
	if req.VectorSimilarityWeight == nil {
		req.VectorSimilarityWeight = func() *float64 { v := 0.3; return &v }()
	}
	if req.RankFeature == nil {
		req.RankFeature = &map[string]float64{"pagerank_fea": 10.0}
	}
	if req.Aggs == nil {
		req.Aggs = func() *bool { v := true; return &v }()
	}

	if req.Page <= 0 {
		req.Page = 1
	}
	if req.PageSize <= 0 {
		req.PageSize = 1
	}

	// Calculate rerank limit to ensure we get enough results for proper pagination
	pageSize := req.PageSize
	rerankLimit := pageSize
	if pageSize > 1 {
		rerankLimit = int(math.Ceil(64.0/float64(pageSize))) * pageSize
	} else {
		rerankLimit = 1
	}
	if rerankLimit < 30 {
		rerankLimit = 30
	}
	// Cap rerank limit when external rerank model is used
	if req.RerankModel != nil && *req.Top > 0 {
		if rerankLimit > *req.Top {
			rerankLimit = *req.Top
		}
		if rerankLimit > 64 {
			rerankLimit = 64
		}
	}

	page := req.Page
	globalOffset := (page - 1) * pageSize
	searchPage := globalOffset/rerankLimit + 1
	common.Debug("Retrieval rerank params", zap.Int("page", req.Page), zap.Int("pageSize", pageSize),
		zap.Int("searchPage", searchPage), zap.Int("rerankLimit", rerankLimit), zap.Int("globalOffset", globalOffset))

	// Execute search via Search()
	searchReq := &RetrievalSearchRequest{
		TenantIDs:      req.TenantIDs,
		Question:       req.Question,
		KbIDs:          req.KbIDs,
		DocIDs:         req.DocIDs,
		Page:           searchPage,
		PageSize:       rerankLimit,
		Top:            *req.Top,
		RankFeature:    *req.RankFeature,
		EmbeddingModel: req.EmbeddingModel,
		Filter:         req.Filter,
	}
	searchResult, err := s.Search(ctx, searchReq)
	if err != nil {
		return nil, fmt.Errorf("search failed: %w", err)
	}
	searchTotal := searchResult.Total

	// Prune deleted chunks
	searchResult, err = s.PruneDeletedChunks(ctx, searchResult)
	if err != nil {
		return nil, fmt.Errorf("PruneDeletedChunks failed: %w", err)
	}
	if searchResult.Total == 0 {
		return &RetrievalResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}, Total: 0}, nil
	}

	sim, termSimilarity, vectorSimilarity := s.scoreSearchResult(ctx, req, searchResult)
	if len(sim) == 0 {
		return &RetrievalResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}, Total: 0}, nil
	}

	// Sort indices (positions into search results) by score descending
	// After sorting by score descending, we process chunks in relevance order
	type idxScore struct {
		idx   int
		score float64
	}
	idxScores := make([]idxScore, 0, len(sim))
	for i, s := range sim {
		idxScores = append(idxScores, idxScore{idx: i, score: s})
	}
	// Use SliceStable for deterministic ordering when scores are tied
	sort.SliceStable(idxScores, func(i, j int) bool {
		return idxScores[i].score > idxScores[j].score
	})

	// When vector_similarity_weight is 0, similarity_threshold is not meaningful for term-only scores
	postThreshold := *req.SimilarityThreshold
	if *req.VectorSimilarityWeight <= 0 {
		postThreshold = 0.0
	}

	// Get valid indices where score >= postThreshold
	validIdx := make([]int, 0)
	for _, is := range idxScores {
		if is.score >= postThreshold {
			validIdx = append(validIdx, is.idx)
		}
	}
	if len(validIdx) == 0 {
		return &RetrievalResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}, Total: 0}, nil
	}

	// Calculate pagination
	// begin and end define which of validIdx to return as the page
	begin := globalOffset % rerankLimit
	end := begin + pageSize

	// Get page indices
	var pageIdx []int
	if begin < len(validIdx) {
		if end > len(validIdx) {
			end = len(validIdx)
		}
		pageIdx = validIdx[begin:end]
	}
	common.Info("Pagination result info", zap.Int("totalValid", len(validIdx)), zap.Int("begin", begin),
		zap.Int("end", end), zap.Int("chunkCount", len(pageIdx)), zap.Float64("postThreshold", postThreshold))

	//searchTotal > len(searchResult.IDs): The total number reported by the engine exceeds the number actually retrieved this time,
	//indicating that there are still unretrieved chunks in the candidate set; the current count may be an underestimate
	//*req.Top > rerankLimit: The Top value requested by the caller exceeds the re-ranking limit, indicating that we have indeed
	//‘only examined a portion’ rather than having covered all target results
	total := int64(len(validIdx))
	if searchTotal > int64(len(searchResult.IDs)) && *req.Top > rerankLimit {
		total, err = s.countThresholdValidMatches(ctx, req, *req.Top, postThreshold)
		if err != nil {
			return nil, err
		}
	}

	// Build chunks for pageIdx, transforms raw search results into the API response format
	var filteredChunks []map[string]interface{}
	dim := 0
	if searchResult.QueryVector != nil {
		dim = len(searchResult.QueryVector)
	}
	zeroVector := make([]float64, dim)
	for j := 0; j < dim; j++ {
		zeroVector[j] = 0.0
	}

	for _, i := range pageIdx {
		if i < 0 || i >= len(searchResult.IDs) {
			continue
		}
		chunkID := searchResult.IDs[i]
		chunk, exists := searchResult.Field[chunkID]
		if !exists {
			continue
		}

		resultChunk := make(map[string]interface{})
		resultChunk["chunk_id"] = chunkID
		if v, ok := chunk["content_ltks"]; ok {
			resultChunk["content_ltks"] = v
		}
		if v, ok := chunk["content_with_weight"]; ok {
			resultChunk["content_with_weight"] = v
		}
		if v, ok := chunk["doc_id"]; ok {
			resultChunk["doc_id"] = v
		}
		if v, ok := chunk["docnm_kwd"]; ok {
			resultChunk["docnm_kwd"] = v
		}
		if v, ok := chunk["kb_id"]; ok {
			resultChunk["kb_id"] = v
		}
		if v, ok := chunk["important_kwd"]; ok {
			resultChunk["important_kwd"] = v
		}
		if v, ok := chunk["tag_kwd"]; ok {
			resultChunk["tag_kwd"] = v
		}
		if v, ok := chunk["img_id"]; ok {
			resultChunk["image_id"] = v
		} else {
			resultChunk["image_id"] = ""
		}
		if v, ok := chunk["position_int"]; ok && v != nil {
			resultChunk["positions"] = v
		} else {
			resultChunk["positions"] = []interface{}{}
		}
		if v, ok := chunk["doc_type_kwd"]; ok && v != nil {
			if s, ok := v.(string); ok {
				if s == "" {
					// Infinity's whitespace-# analyzer returns empty string as [] in Python SDK
					// but as "" in Go SDK. Both Infinity and Elasticsearch paths normalize
					// to None on the Python side (the test converts [] to None), so use nil
					// here for parity instead of []interface{}{}.
					resultChunk["doc_type_kwd"] = nil
				} else {
					resultChunk["doc_type_kwd"] = s
				}
			} else if sliceVal, ok := v.([]interface{}); ok {
				if len(sliceVal) == 0 {
					resultChunk["doc_type_kwd"] = nil
				} else {
					resultChunk["doc_type_kwd"] = sliceVal
				}
			} else {
				resultChunk["doc_type_kwd"] = nil
			}
		} else {
			resultChunk["doc_type_kwd"] = nil
		}
		// row_id: row identifier (for structured data like tables)
		if v, ok := chunk["row_id()"]; ok {
			resultChunk["row_id"] = v
		}
		resultChunk["similarity"] = sim[i]
		resultChunk["term_similarity"] = termSimilarity[i]
		resultChunk["vector_similarity"] = vectorSimilarity[i]

		// Always set these fields even if empty, to match Python response format
		if v, ok := chunk["important_kwd"]; ok {
			resultChunk["important_kwd"] = v
		} else {
			resultChunk["important_kwd"] = []string{}
		}
		if v, ok := chunk["mom_id"]; ok {
			resultChunk["mom_id"] = v
		} else {
			resultChunk["mom_id"] = ""
		}
		if v, ok := chunk["row_id()"]; ok {
			resultChunk["row_id"] = v
		} else {
			resultChunk["row_id"] = nil
		}
		if v, ok := chunk["tag_kwd"]; ok {
			resultChunk["tag_kwd"] = v
		} else {
			resultChunk["tag_kwd"] = []string{}
		}

		vectorColumn := fmt.Sprintf("q_%d_vec", dim)
		if v, ok := chunk[vectorColumn]; ok {
			resultChunk["vector"] = v
		} else {
			resultChunk["vector"] = zeroVector
		}

		highlightEnabled := false
		if req.Highlight != nil && *req.Highlight {
			highlightEnabled = true
		}
		if highlightEnabled && searchResult.Highlight != nil {
			if highlightText, ok := searchResult.Highlight[chunkID]; ok {
				resultChunk["highlight"] = RemoveRedundantSpaces(highlightText)
			} else if contentWithWeight, ok := chunk["content_with_weight"].(string); ok {
				resultChunk["highlight"] = RemoveRedundantSpaces(contentWithWeight)
			}
		}
		filteredChunks = append(filteredChunks, resultChunk)
	}

	// Build document aggregation, aggregates document-level statistics across all valid chunks
	// This is useful for showing users which documents are most relevant to their query.
	var docAggs []map[string]interface{}
	if req.Aggs != nil && *req.Aggs {
		docAggsMap := make(map[string]struct {
			docID string
			count int
		})
		for _, i := range validIdx {
			if i < 0 || i >= len(searchResult.IDs) {
				continue
			}
			chunkID := searchResult.IDs[i]
			chunk, exists := searchResult.Field[chunkID]
			if !exists {
				continue
			}
			docName := ""
			docID := ""
			if v, ok := chunk["docnm_kwd"].(string); ok {
				docName = v
			}
			if v, ok := chunk["doc_id"].(string); ok {
				docID = v
			}
			if entry, exists := docAggsMap[docName]; exists {
				entry.count++
				docAggsMap[docName] = entry
			} else {
				docAggsMap[docName] = struct {
					docID string
					count int
				}{docID: docID, count: 1}
			}
		}

		// Sort by count descending
		type docAggEntry struct {
			docName string
			docID   string
			count   int
		}
		docAggsList := make([]docAggEntry, 0, len(docAggsMap))
		for docName, entry := range docAggsMap {
			docAggsList = append(docAggsList, docAggEntry{docName: docName, docID: entry.docID, count: entry.count})
		}
		sort.Slice(docAggsList, func(i, j int) bool {
			return docAggsList[i].count > docAggsList[j].count
		})

		docAggs = make([]map[string]interface{}, 0, len(docAggsList))
		for _, entry := range docAggsList {
			docAggs = append(docAggs, map[string]interface{}{
				"doc_name": entry.docName,
				"doc_id":   entry.docID,
				"count":    entry.count,
			})
		}
	} else {
		docAggs = []map[string]interface{}{}
	}

	return &RetrievalResult{
		Chunks:  filteredChunks,
		DocAggs: docAggs,
		Total:   total,
	}, nil
}

func (s *RetrievalService) countThresholdValidMatches(ctx context.Context, req *RetrievalRequest, limit int, postThreshold float64) (int64, error) {
	searchReq := &RetrievalSearchRequest{
		TenantIDs:      req.TenantIDs,
		Question:       req.Question,
		KbIDs:          req.KbIDs,
		DocIDs:         req.DocIDs,
		Page:           1,
		PageSize:       limit,
		Top:            *req.Top,
		RankFeature:    *req.RankFeature,
		EmbeddingModel: req.EmbeddingModel,
		Filter:         req.Filter,
	}
	searchResult, err := s.Search(ctx, searchReq)
	if err != nil {
		return 0, fmt.Errorf("search count failed: %w", err)
	}
	searchResult, err = s.PruneDeletedChunks(ctx, searchResult)
	if err != nil {
		return 0, fmt.Errorf("PruneDeletedChunks count failed: %w", err)
	}
	if searchResult.Total == 0 {
		return 0, nil
	}

	sim, _, _ := s.scoreSearchResult(ctx, req, searchResult)
	var total int64
	for _, score := range sim {
		if score >= postThreshold {
			total++
		}
	}
	return total, nil
}

func (s *RetrievalService) scoreSearchResult(ctx context.Context, req *RetrievalRequest, searchResult *RetrievalSearchResult) ([]float64, []float64, []float64) {
	// sim = tkWeight*tsim + vtWeight*vsim
	vtWeight := *req.VectorSimilarityWeight
	tkWeight := 1.0 - vtWeight
	qb := GetQueryBuilder()
	useInfinity := engine.GetEngineType() == "infinity"
	useOceanBase := engine.IsOceanBaseFamily(s.docEngine.GetType())

	if req.RerankModel != nil && searchResult.Total > 0 {
		return RerankByModel(
			ctx,
			req.RerankModel,
			searchResult.Chunks,
			searchResult.IDs,
			searchResult.Field,
			req.Question,
			tkWeight,
			vtWeight,
			"content_ltks",
			qb,
			*req.RankFeature,
		)
	}

	if useInfinity {
		sim := make([]float64, len(searchResult.IDs))
		for i, id := range searchResult.IDs {
			if chunk, ok := searchResult.Field[id]; ok {
				if score, ok := chunk["_score"].(float64); ok {
					sim[i] = score
				} else if score, ok := chunk["SCORE"].(float64); ok {
					sim[i] = score
				} else if score, ok := chunk["SIMILARITY"].(float64); ok {
					sim[i] = score
				}
			}
		}
		return sim, sim, sim
	}

	if useOceanBase {
		return RerankStandard(
			searchResult.Chunks,
			nil,
			searchResult.QueryVector,
			req.Question,
			tkWeight,
			vtWeight,
			"content_ltks",
			qb,
			*req.RankFeature,
		)
	}

	knnResult, err := s.docEngine.KNNScores(ctx, searchResult.Chunks, searchResult.QueryVector, len(searchResult.IDs))
	if err != nil {
		common.Warn("KNNScores failed for ES, falling back to local computation", zap.Error(err))
		return RerankStandard(
			searchResult.Chunks,
			nil,
			searchResult.QueryVector,
			req.Question,
			tkWeight,
			vtWeight,
			"content_ltks",
			qb,
			*req.RankFeature,
		)
	}
	knnScores := s.docEngine.GetScores(knnResult)
	return RerankWithKNN(
		searchResult.Chunks,
		searchResult.IDs,
		searchResult.Field,
		knnScores,
		req.Question,
		tkWeight,
		vtWeight,
		"content_ltks",
		qb,
		*req.RankFeature,
	)
}

// RetrievalSearchRequest is the request struct for RetrievalService.Search()
type RetrievalSearchRequest struct {
	Question            string
	TenantIDs           []string
	KbIDs               []string
	DocIDs              []string
	Top                 int
	Page                int
	PageSize            int
	Sort                bool
	Highlight           *bool
	SimilarityThreshold float64
	RankFeature         map[string]float64
	Filter              map[string]interface{}
	EmbeddingModel      *models.EmbeddingModel
}

type RetrievalSearchResult struct {
	Chunks      []map[string]interface{}          // Search results
	Total       int64                             // Total number of matches
	QueryVector []float64                         // Query vector (for hybrid search, used in reranking)
	Highlight   map[string]string                 // Highlighted snippets (chunk_id -> highlighted text)
	Field       map[string]map[string]interface{} // ID -> chunk mapping
	IDs         []string                          // Ordered list of chunk IDs
	Keywords    []string                          // Keywords from query
	Aggregation []map[string]interface{}          // Doc aggregation by field
	Options     map[string]interface{}            // Engine-specific options (e.g., total from get_total)
	IndexNames  []string                          // Index names for second-pass queries (e.g., KNN scores)
}

// Search performs search based on question and EmbeddingModel:
// - Empty question: list data matching filters, optionally sorted
// - Non-empty question, no EmbeddingModel: fulltext search only
// - Non-empty question, with EmbeddingModel: hybrid search (fulltext + vector + fusion)
//
// Hybrid search path retries with lower thresholds if no results found.
func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchRequest) (*RetrievalSearchResult, error) {
	if req.Highlight == nil {
		req.Highlight = func() *bool { v := false; return &v }()
	}
	filters := req.GetFilters()
	if _, ok := filters["available_int"]; !ok {
		filters["available_int"] = 1
	}
	pg := max(req.Page-1, 0)
	topk := req.Top
	if topk <= 0 {
		topk = 1024
	}
	pageSize := req.PageSize
	if pageSize <= 0 {
		pageSize = topk
	}
	limit := pageSize

	// Build Source field list
	src := []string{
		"docnm_kwd", "content_ltks", "kb_id", "img_id", "title_tks", "important_kwd", "position_int",
		"doc_id", "chunk_order_int", "page_num_int", "top_int", "create_timestamp_flt", "knowledge_graph_kwd",
		"question_kwd", "question_tks", "doc_type_kwd",
		"available_int", "content_with_weight", "mom_id", "pagerank_fea", "tag_feas", "row_id()",
		"_score",
	}

	kwds := make(map[string]struct{})

	// Build base engine request with common fields
	// Note: RankFeature is NOT set here, it's set per-call where needed
	searchRequest := &types.SearchRequest{
		IndexNames:   buildIndexNames(req.TenantIDs),
		KbIDs:        req.KbIDs,
		Offset:       pg * pageSize,
		Limit:        limit,
		Filter:       filters,
		SelectFields: src,
	}

	// engineResult holds the result from docEngine.Search() (types.SearchResult)
	// queryVector tracks the query vector for reranking
	var engineResult *types.SearchResult
	var queryVector []float64
	var err error

	if req.Question == "" {
		// Empty question
		if req.Sort {
			searchRequest.OrderBy = &types.OrderByExpr{}
			searchRequest.OrderBy.Asc("chunk_order_int").Asc("page_num_int").Asc("top_int").Desc("create_timestamp_flt")
		}
		searchRequest.MatchExprs = []interface{}{}
		engineResult, err = s.docEngine.Search(ctx, searchRequest)
		if err != nil {
			return nil, fmt.Errorf("search failed: %w", err)
		}
	} else {
		// Non-empty question

		// Compute keywords via QueryBuilder
		matchText, keywords := GetQueryBuilder().Question(req.Question, "", 0.3)
		for _, k := range keywords {
			kwds[k] = struct{}{}
		}

		// Check if EmbeddingModel is available
		if req.EmbeddingModel == nil {
			// Keyword-only search
			searchRequestWithRank := *searchRequest
			searchRequestWithRank.MatchExprs = []interface{}{matchText}
			searchRequestWithRank.RankFeature = req.RankFeature

			engineResult, err = s.docEngine.Search(ctx, &searchRequestWithRank)
			if err != nil {
				return nil, fmt.Errorf("search failed: %w", err)
			}
			queryVector = nil
		} else {
			// Compute question vector via GetVector
			similarityForGetVector := req.SimilarityThreshold
			if similarityForGetVector <= 0 {
				similarityForGetVector = 0.1
			}
			matchDense, err := s.GetVector(ctx, req.Question, req.EmbeddingModel, topk, similarityForGetVector)
			if err != nil {
				return nil, fmt.Errorf("GetVector failed: %w", err)
			}

			// Execute search with fusion
			fusionExpr := &types.FusionExpr{
				Method:       "weighted_sum",
				TopN:         topk,
				FusionParams: map[string]interface{}{"weights": "0.05,0.95"},
			}

			// Build source with vector column for ES
			searchSrc := make([]string, len(searchRequest.SelectFields))
			copy(searchSrc, searchRequest.SelectFields)
			if engine.GetEngineType() == "elasticsearch" || engine.IsOceanBaseFamily(engine.GetEngineType()) {
				searchSrc = append(searchSrc, matchDense.VectorColumnName)
			}

			searchRequest.SelectFields = searchSrc
			searchRequest.MatchExprs = []interface{}{matchText, matchDense, fusionExpr}
			searchRequest.RankFeature = req.RankFeature

			engineResult, err = s.docEngine.Search(ctx, searchRequest)
			if err != nil {
				return nil, fmt.Errorf("search failed: %w", err)
			}
			// If result is empty, retry with relaxed conditions
			if engineResult.Total == 0 {
				_, hasDocIDFilter := filters["doc_id"]
				if hasDocIDFilter {
					// When a doc_id filter is present (e.g. from metadata filter like era=960)
					// and the hybrid search returns no results, fall back to a filter-only
					// search (no text match, no vector match). This ensures that when a
					// metadata filter restricts the search to a specific set of documents
					// that happen to have no relevant content for the query, we still
					// return those documents' chunks (ordered by the request's sort/order).
					//
					// Example: searching "打虎" with metadata filter era≠960 limits the
					// search to Three Kingdoms documents (era=220). Since "打虎" only
					// appears in Water Margin (era=960), the hybrid search returns 0
					// results. This fallback returns all chunks from Three Kingdoms
					// documents instead of returning an empty result.
					searchRequest.SelectFields = src
					searchRequest.MatchExprs = []interface{}{}
					searchRequest.RankFeature = nil

					engineResult, err = s.docEngine.Search(ctx, searchRequest)
					if err != nil {
						return nil, fmt.Errorf("search retry failed: %w", err)
					}
				} else {
					// No doc_id filter — retry with lower min_match (0.1 vs default 0.3)
					// and lower vector similarity threshold (0.17 vs default 0.1-0.2).
					// This provides a second chance for queries that were too strict
					// on the first attempt.
					matchText, _ := GetQueryBuilder().Question(req.Question, "qa", 0.1)
					matchDense.ExtraOptions["similarity"] = 0.17
					searchRequest.MatchExprs = []interface{}{matchText, matchDense, fusionExpr}
					searchRequest.RankFeature = req.RankFeature

					engineResult, err = s.docEngine.Search(ctx, searchRequest)
					if err != nil {
						return nil, fmt.Errorf("search retry failed: %w", err)
					}
				}
			}

			queryVector = matchDense.EmbeddingData
		}

		// Build kwds from keywords with fine-grained tokenization
		for _, k := range keywords {
			kwds[k] = struct{}{}
			fgToken, _ := tokenizer.FineGrainedTokenize(k)
			for _, kk := range strings.Fields(fgToken) {
				if len(kk) < 2 {
					continue
				}
				if _, ok := kwds[kk]; ok {
					continue
				}
				kwds[kk] = struct{}{}
			}
		}
	}

	searchResult := engineResult
	ids := s.docEngine.GetChunkIDs(searchResult.Chunks)
	common.Info("GetChunkIDs result", zap.Int("count", len(ids)), zap.Strings("ids", ids))

	// Build Keywords list from kwds set
	keywordsList := make([]string, 0, len(kwds))
	for k := range kwds {
		keywordsList = append(keywordsList, k)
	}

	fieldMap := s.docEngine.GetFields(searchResult.Chunks, src)
	common.Info("GetFields result", zap.Int("count", len(fieldMap)), zap.Strings("keys", func() []string {
		keys := make([]string, 0, len(fieldMap))
		for k := range fieldMap {
			keys = append(keys, k)
		}
		return keys
	}()), zap.Strings("ids_from_GetDocIDs", ids))

	// Build Aggregation
	aggregation := s.docEngine.GetAggregation(searchResult.Chunks, "docnm_kwd")

	// Build Highlight using GetHighlight
	var highlight map[string]string
	if len(keywordsList) > 0 {
		highlight = s.docEngine.GetHighlight(searchResult.Chunks, keywordsList, "content_with_weight")
	}

	return &RetrievalSearchResult{
		Chunks:      searchResult.Chunks,
		Total:       searchResult.Total,
		QueryVector: queryVector,
		Highlight:   highlight,
		Field:       fieldMap,
		IDs:         ids,
		Keywords:    keywordsList,
		Aggregation: aggregation,
		IndexNames:  searchRequest.IndexNames,
	}, nil
}

// GetVector computes query vector and returns MatchDenseExpr for hybrid search
func (s *RetrievalService) GetVector(ctx context.Context, txt string, embModel *models.EmbeddingModel, topk int, similarity float64) (*types.MatchDenseExpr, error) {
	embeddingConfig := &models.EmbeddingConfig{
		Dimension: 0,
	}
	embeddings, err := embModel.ModelDriver.Embed(ctx, embModel.ModelName, models.EmbedRequest{Texts: []string{txt}}, embModel.APIConfig, embeddingConfig, nil)
	if err != nil {
		return nil, err
	}

	vector := embeddings[0].Embedding
	vectorSize := len(vector)
	vectorColumnName := fmt.Sprintf("q_%d_vec", vectorSize)

	return &types.MatchDenseExpr{
		VectorColumnName:  vectorColumnName,
		EmbeddingData:     vector,
		EmbeddingDataType: "float",
		DistanceType:      "cosine",
		TopN:              topk,
		ExtraOptions:      map[string]interface{}{"similarity": similarity},
	}, nil
}

// GetFilters builds metadata filter map from RetrievalSearchRequest
func (r *RetrievalSearchRequest) GetFilters() map[string]interface{} {
	filters := make(map[string]interface{})

	if len(r.KbIDs) > 0 {
		filters["kb_id"] = r.KbIDs
	}
	if len(r.DocIDs) > 0 {
		filters["doc_id"] = r.DocIDs
	}
	for _, key := range []string{"knowledge_graph_kwd", "available_int", "entity_kwd", "from_entity_kwd", "to_entity_kwd", "removed_kwd"} {
		if val, ok := r.Filter[key]; ok && val != nil {
			filters[key] = val
		}
	}
	for key, val := range r.Filter {
		if _, exists := filters[key]; !exists && val != nil {
			filters[key] = val
		}
	}
	return filters
}

// RetrievalByChildren aggregates child chunks into parent chunks
func RetrievalByChildren(chunks []map[string]interface{}, tenantIDs []string, docEngine engine.DocEngine, ctx context.Context) []map[string]interface{} {
	common.Info("RetrievalByChildren started", zap.Int("chunks", len(chunks)), zap.Strings("tenantIDs", tenantIDs))

	indexNames := buildIndexNames(tenantIDs)
	if len(chunks) == 0 || len(indexNames) == 0 {
		return chunks
	}

	// Group child chunks by mom_id
	type childChunk struct {
		chunk map[string]interface{}
		kbID  string
	}
	momChunks := make(map[string][]childChunk)
	remainingChunks := make([]map[string]interface{}, 0, len(chunks))

	for _, ck := range chunks {
		momID, ok := ck["mom_id"].(string)
		if !ok || momID == "" {
			remainingChunks = append(remainingChunks, ck)
			continue
		}
		kbID, _ := ck["kb_id"].(string)
		momChunks[momID] = append(momChunks[momID], childChunk{chunk: ck, kbID: kbID})
	}

	if len(momChunks) == 0 {
		common.Info("RetrievalByChildren finished", zap.Int("momChunks", len(momChunks)), zap.Int("resultChunks", len(chunks)))
		return chunks
	}

	// Fetch parent chunks and aggregate
	vectorSize := 1024
	for momID, childList := range momChunks {
		kbIDs := make([]string, 0, len(childList))
		for _, c := range childList {
			if c.kbID != "" {
				kbIDs = append(kbIDs, c.kbID)
			}
		}
		if len(kbIDs) == 0 {
			kbIDs = append(kbIDs, "")
		}

		parent, err := docEngine.GetChunk(ctx, indexNames[0], momID, kbIDs)
		if err != nil {
			common.Warn("Failed to get parent chunk", zap.String("momID", momID), zap.Error(err))
			continue
		}
		parentMap, ok := parent.(map[string]interface{})
		if !ok {
			continue
		}

		// Calculate average similarity
		simBuf := make([]float64, 0, len(childList))
		for _, c := range childList {
			if sim, ok := c.chunk["similarity"].(float64); ok {
				simBuf = append(simBuf, sim)
			}
		}
		totalSim := common.PairwiseSum(simBuf)
		avgSim := totalSim / float64(len(childList))

		// Collect content_ltks from children
		var contentParts []string
		for _, c := range childList {
			if ltks, ok := c.chunk["content_ltks"].(string); ok {
				contentParts = append(contentParts, ltks)
			}
		}
		contentLTKS := strings.Join(contentParts, " ")

		// Collect important_kwd from children
		allImportantKwd := []string{}
		for _, c := range childList {
			if kwd, ok := c.chunk["important_kwd"].([]interface{}); ok {
				for _, k := range kwd {
					if ks, ok := k.(string); ok {
						allImportantKwd = append(allImportantKwd, ks)
					}
				}
			}
		}

		// Build aggregated chunk
		docTypeKwd := ""
		if v, ok := parentMap["doc_type_kwd"].(string); ok {
			docTypeKwd = v
		}
		imgID := parentMap["img_id"]
		if imgID == nil || imgID == "" {
			imgID = ""
		}
		aggregated := map[string]interface{}{
			"chunk_id":            momID,
			"content_ltks":        contentLTKS,
			"content_with_weight": parentMap["content_with_weight"],
			"doc_id":              parentMap["doc_id"],
			"docnm_kwd":           parentMap["docnm_kwd"],
			"kb_id":               parentMap["kb_id"],
			"important_kwd":       allImportantKwd,
			"image_id":            imgID,
			"similarity":          avgSim,
			"vector_similarity":   avgSim,
			"term_similarity":     avgSim,
			"vector":              make([]float64, vectorSize),
			"positions":           parentMap["position_int"],
			"doc_type_kwd":        docTypeKwd,
		}

		// Get vector from first child if available
	childVecLoop:
		for _, c := range childList {
			for k := range c.chunk {
				if strings.HasSuffix(k, "_vec") {
					if vec, ok := c.chunk[k].([]float64); ok {
						aggregated["vector"] = vec
						vectorSize = len(vec)
						break childVecLoop
					}
				}
			}
		}

		remainingChunks = append(remainingChunks, aggregated)
	}

	// Sort by similarity descending
	for i := 0; i < len(remainingChunks); i++ {
		for j := i + 1; j < len(remainingChunks); j++ {
			simI, _ := remainingChunks[i]["similarity"].(float64)
			simJ, _ := remainingChunks[j]["similarity"].(float64)
			if simJ > simI {
				remainingChunks[i], remainingChunks[j] = remainingChunks[j], remainingChunks[i]
			}
		}
	}

	common.Info("RetrievalByChildren finished", zap.Int("momChunks", len(momChunks)), zap.Int("resultChunks", len(remainingChunks)))
	return remainingChunks
}

// PruneDeletedChunks removes chunks whose documents no longer exist
func (s *RetrievalService) PruneDeletedChunks(ctx context.Context, result *RetrievalSearchResult) (*RetrievalSearchResult, error) {
	if s.documentDAO == nil {
		return nil, fmt.Errorf("documentDAO is not initialized")
	}
	// Collect all doc_ids from chunks
	chunkDocIDs := make([]string, 0, len(result.Field))
	for _, chunk := range result.Field {
		if docID, ok := chunk["doc_id"].(string); ok && docID != "" {
			chunkDocIDs = append(chunkDocIDs, docID)
		}
	}

	if len(chunkDocIDs) == 0 {
		return result, nil
	}

	// Deduplicate chunkDocIDs for correct comparison with existingDocIDs
	uniqueDocIDs := make([]string, 0, len(chunkDocIDs))
	seen := make(map[string]struct{}, len(chunkDocIDs))
	for _, id := range chunkDocIDs {
		if _, exists := seen[id]; !exists {
			seen[id] = struct{}{}
			uniqueDocIDs = append(uniqueDocIDs, id)
		}
	}

	// Get existing document IDs
	docs, err := s.documentDAO.GetByIDs(ctx, dao.DB, uniqueDocIDs)
	if err != nil {
		return nil, fmt.Errorf("GetByIDs failed: %w", err)
	}

	existingDocIDs := make(map[string]struct{}, len(docs))
	for _, doc := range docs {
		existingDocIDs[doc.ID] = struct{}{}
	}

	// Early return if all docs exist
	if len(existingDocIDs) == len(uniqueDocIDs) {
		return result, nil
	}

	// Filter out chunks with deleted documents
	filteredIDs := make([]string, 0, len(result.IDs))
	filteredChunks := make([]map[string]interface{}, 0, len(result.IDs))
	filteredField := make(map[string]map[string]interface{}, len(result.IDs))
	filteredHighlight := make(map[string]string)
	removed := 0

	for _, chunkID := range result.IDs {
		chunk, exists := result.Field[chunkID]
		if !exists {
			continue
		}
		docID, ok := chunk["doc_id"].(string)
		if !ok || docID == "" {
			// Keep chunks without doc_id
			filteredIDs = append(filteredIDs, chunkID)
			filteredChunks = append(filteredChunks, chunk)
			filteredField[chunkID] = chunk
			if result.Highlight != nil {
				if hl, ok := result.Highlight[chunkID]; ok {
					filteredHighlight[chunkID] = hl
				}
			}
			continue
		}
		if _, docExists := existingDocIDs[docID]; !docExists {
			removed++
			continue
		}
		filteredIDs = append(filteredIDs, chunkID)
		filteredChunks = append(filteredChunks, chunk)
		filteredField[chunkID] = chunk
		if result.Highlight != nil {
			if hl, ok := result.Highlight[chunkID]; ok {
				filteredHighlight[chunkID] = hl
			}
		}
	}

	if removed > 0 {
		common.Warn("Pruned stale chunks whose documents no longer exist", zap.Int("removed", removed))
	}

	return &RetrievalSearchResult{
		Chunks:      filteredChunks,
		Total:       int64(len(filteredIDs)),
		QueryVector: result.QueryVector,
		Highlight:   filteredHighlight,
		Field:       filteredField,
		IDs:         filteredIDs,
		Keywords:    result.Keywords,
		Aggregation: result.Aggregation,
		Options:     result.Options,
	}, nil
}

// buildIndexNames creates index names for the given tenant IDs.
// Each tenantID may be a comma-separated list.
func buildIndexNames(tenantIDs []string) []string {
	var indexNames []string
	for _, tid := range tenantIDs {
		for _, part := range strings.Split(tid, ",") {
			part = strings.TrimSpace(part)
			if part != "" {
				indexNames = append(indexNames, fmt.Sprintf("ragflow_%s", part))
			}
		}
	}
	return indexNames
}

// FetchChunkVectors returns q_{dim}_vec for the given chunk IDs.
// Missing or wrong-dimension chunks get a zero vector.
func (s *RetrievalService) FetchChunkVectors(ctx context.Context, chunkIDs []string, tenantIDs []string, kbIDs []string, dim int) (map[string][]float64, error) {
	if dim <= 0 {
		return nil, fmt.Errorf("FetchChunkVectors: dim must be > 0, got %d", dim)
	}
	if len(chunkIDs) == 0 {
		return map[string][]float64{}, nil
	}

	vecField := fmt.Sprintf("q_%d_vec", dim)
	idxNames := buildIndexNames(tenantIDs)

	req := &types.SearchRequest{
		IndexNames:   idxNames,
		KbIDs:        kbIDs,
		Limit:        len(chunkIDs),
		Offset:       0,
		SelectFields: []string{"id", vecField},
		Filter:       map[string]interface{}{"id": chunkIDs},
		MatchExprs:   []interface{}{},
	}

	result, err := s.docEngine.Search(ctx, req)
	if err != nil {
		return nil, fmt.Errorf("FetchChunkVectors: engine search failed: %w", err)
	}

	out := make(map[string][]float64, len(chunkIDs))
	for _, cid := range chunkIDs {
		out[cid] = make([]float64, dim)
	}

	for _, chunk := range result.Chunks {
		cid, _ := chunk["id"].(string)
		if cid == "" {
			continue
		}
		var vec []float64
		switch v := chunk[vecField].(type) {
		case []float64:
			vec = v
		case []interface{}:
			vec = make([]float64, len(v))
			for i, val := range v {
				if f, ok := val.(float64); ok {
					vec[i] = f
				} else if f32, ok := val.(float32); ok {
					vec[i] = float64(f32)
				}
			}
		case string:
			// Tab-separated floats (mirrors Python's split("\t") in
			// search.py:435-437 when Infinity returns vectors as a string).
			parts := strings.Split(v, "\t")
			vec = make([]float64, 0, len(parts))
			for _, p := range parts {
				if f, err := strconv.ParseFloat(strings.TrimSpace(p), 64); err == nil {
					vec = append(vec, f)
				}
			}
		}
		if len(vec) != dim {
			vec = make([]float64, dim)
		}
		out[cid] = vec
	}

	return out, nil
}