From 5bb5ba22166da4f5db689440d3dd76692890feb9 Mon Sep 17 00:00:00 2001 From: Haruko386 Date: Fri, 14 Aug 2026 17:44:38 +0800 Subject: [PATCH] fix: return total numbers of the search result (#18071) ### Summary As title --- internal/service/nlp/retrieval.go | 264 +++++++++++++------------ internal/service/nlp/retrieval_test.go | 178 +++++++++++++++++ 2 files changed, 313 insertions(+), 129 deletions(-) create mode 100644 internal/service/nlp/retrieval_test.go diff --git a/internal/service/nlp/retrieval.go b/internal/service/nlp/retrieval.go index e6f8b4affe..50ae32467f 100644 --- a/internal/service/nlp/retrieval.go +++ b/internal/service/nlp/retrieval.go @@ -67,7 +67,7 @@ type RetrievalRequest struct { type RetrievalResult struct { Chunks []map[string]interface{} DocAggs []map[string]interface{} // Aggregated document counts, sorted by count desc - Total int64 // Post-pagination chunk count (matches Python's len(ranks["chunks"])) + Total int64 // Threshold-valid matches across the retrieval candidate set } // Retrieval performs hybrid search + reranking + pagination @@ -150,6 +150,7 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) if err != nil { return nil, fmt.Errorf("search failed: %w", err) } + searchTotal := searchResult.Total // Prune deleted chunks searchResult, err = s.PruneDeletedChunks(ctx, searchResult) @@ -160,132 +161,7 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) return &RetrievalResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}, Total: 0}, nil } - // sim = tkWeight*tsim + vtWeight*vsim - vtWeightOrig := *req.VectorSimilarityWeight - tkWeightOrig := 1.0 - vtWeightOrig - tkWeight := tkWeightOrig - vtWeight := vtWeightOrig - qb := GetQueryBuilder() - useInfinity := engine.GetEngineType() == "infinity" - useOceanBase := engine.IsOceanBaseFamily(s.docEngine.GetType()) - - // For ES path: call GetScores() for second-pass KNN to get clean cosine similarity - // For Infinity path: use _score directly (scores already normalized during fusion) - // For OceanBase path: extract vectors and compute locally - var sim []float64 - var termSimilarity []float64 - var vectorSimilarity []float64 - - if req.RerankModel != nil && searchResult.Total > 0 { - // External rerank model path - use RerankByModel - sim, termSimilarity, vectorSimilarity = RerankByModel( - ctx, - req.RerankModel, - searchResult.Chunks, - searchResult.IDs, - searchResult.Field, - req.Question, - tkWeight, - vtWeight, - "content_ltks", - qb, - *req.RankFeature, - ) - } else if useInfinity { - // Infinity: scores already normalized before fusion, just extract _score - 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 - } else { - sim[i] = 0.0 - } - } else { - sim[i] = 0.0 - } - } - termSimilarity = sim - vectorSimilarity = sim - } else if useOceanBase { - // OceanBase returns the selected vector column, so reranking can - // reproduce the Python connector's local cosine calculation. - sim, termSimilarity, vectorSimilarity = RerankStandard( - searchResult.Chunks, - nil, - searchResult.QueryVector, - req.Question, - tkWeight, - vtWeight, - "content_ltks", - qb, - *req.RankFeature, - ) - } else { - // ES PATH: Two-pass KNN approach for clean cosine similarity scores - // - // Python's equivalent flow (rag/nlp/search.py L656-669): - // 1. First search returns text+vector matched chunks (hybrid BM25 + KNN fusion) - // 2. _knn_scores: second KNN-only query filtered by those chunk IDs - // - ES computes cosine similarity between query vector and stored vectors - // - Vectors stay in ES index (not shipped to application) - // - Returns raw KNN result with _id -> _score mappings - // 3. get_scores: extracts doc_id -> score from the KNN result - // 4. rerank_with_knn: combines token similarity + vector similarity + rank features - // - // Go implementation mirrors this exactly: - // KNNScores() -> performs second KNN query (ES-specific, on DocEngine interface) - // GetScores() -> extracts doc_id -> score from result (matches Python's get_scores) - // RerankWithKNN() -> combines tksim + vtsim + rank_features (matches rerank_with_knn) - // - // Why two passes? - // - First search uses fusion (BM25 * vector_similarity_weight + KNN * (1-vector_similarity_weight)) - // - The fusion score is not a clean cosine similarity - it's a weighted combination - // - Second KNN-only query gives us the pure cosine similarity for reranking - // - This keeps vectors in ES (no need to extract them) while getting clean scores - - // PASS 1: Second KNN query to get clean cosine similarities - // KNNScores() performs the ES-specific KNN search and returns raw result - 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)) - // Fallback: RerankStandard computes vector similarity locally (requires shipping vectors) - sim, termSimilarity, vectorSimilarity = RerankStandard( - searchResult.Chunks, - nil, // keywords computed internally - searchResult.QueryVector, - req.Question, - tkWeight, - vtWeight, - "content_ltks", - qb, - *req.RankFeature, - ) - } else { - // PASS 2: Extract scores from KNN result - // GetScores() mirrors Python's get_scores() - maps doc_id -> _score - knnScores := s.docEngine.GetScores(knnResult) - - // RERANK: Combine token + vector + rank feature similarities - // Matches Python's rerank_with_knn(): sim = tkweight * tksim + vtweight * vtsim + rank_fea - sim, termSimilarity, vectorSimilarity = RerankWithKNN( - searchResult.Chunks, - searchResult.IDs, - searchResult.Field, - knnScores, - req.Question, - tkWeight, - vtWeight, - "content_ltks", - qb, - *req.RankFeature, - ) - } - } + 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 } @@ -307,7 +183,7 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) // When vector_similarity_weight is 0, similarity_threshold is not meaningful for term-only scores postThreshold := *req.SimilarityThreshold - if vtWeight <= 0 { + if *req.VectorSimilarityWeight <= 0 { postThreshold = 0.0 } @@ -338,6 +214,18 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) 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 @@ -531,10 +419,128 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) return &RetrievalResult{ Chunks: filteredChunks, DocAggs: docAggs, - Total: int64(len(filteredChunks)), + 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, + } + 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 diff --git a/internal/service/nlp/retrieval_test.go b/internal/service/nlp/retrieval_test.go new file mode 100644 index 0000000000..531616ecb0 --- /dev/null +++ b/internal/service/nlp/retrieval_test.go @@ -0,0 +1,178 @@ +package nlp + +import ( + "context" + "fmt" + "testing" + + "ragflow/internal/dao" + "ragflow/internal/engine/types" + + "gorm.io/gorm" +) + +func TestRetrievalTotalCountsThresholdValidMatchesBeyondRerankWindow(t *testing.T) { + oldQueryBuilder := globalQueryBuilder + globalQueryBuilder = NewQueryBuilder() + defer func() { globalQueryBuilder = oldQueryBuilder }() + + rows := make([]map[string]interface{}, 75) + for i := range rows { + rows[i] = map[string]interface{}{ + "id": fmt.Sprintf("chunk-%02d", i), + "content_ltks": "alpha", + "content_with_weight": "alpha", + "_score": 0.9, + } + } + engine := &retrievalCountEngine{rows: rows} + service := NewRetrievalService(engine, &dao.DocumentDAO{}) + top := 100 + threshold := 0.5 + vectorWeight := 1.0 + aggs := false + + result, err := service.Retrieval(context.Background(), &RetrievalRequest{ + Question: "alpha", + TenantIDs: []string{"tenant-1"}, + Page: 1, + PageSize: 10, + Top: &top, + SimilarityThreshold: &threshold, + VectorSimilarityWeight: &vectorWeight, + Aggs: &aggs, + }) + if err != nil { + t.Fatalf("Retrieval failed: %v", err) + } + if len(result.Chunks) != 10 { + t.Fatalf("page chunk count = %d, want 10", len(result.Chunks)) + } + if result.Total != 75 { + t.Fatalf("total = %d, want 75", result.Total) + } + if len(engine.searchLimits) != 2 || engine.searchLimits[0] != 70 || engine.searchLimits[1] != 100 { + t.Fatalf("search limits = %v, want [70 100]", engine.searchLimits) + } +} + +type retrievalCountEngine struct { + rows []map[string]interface{} + searchLimits []int +} + +func (e *retrievalCountEngine) Search(_ context.Context, req *types.SearchRequest) (*types.SearchResult, error) { + e.searchLimits = append(e.searchLimits, req.Limit) + offset := req.Offset + if offset > len(e.rows) { + offset = len(e.rows) + } + end := offset + req.Limit + if req.Limit <= 0 || end > len(e.rows) { + end = len(e.rows) + } + return &types.SearchResult{Chunks: e.rows[offset:end], Total: int64(len(e.rows))}, nil +} + +func (e *retrievalCountEngine) GetChunkIDs(chunks []map[string]interface{}) []string { + ids := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + if id, ok := chunk["id"].(string); ok { + ids = append(ids, id) + } + } + return ids +} + +func (e *retrievalCountEngine) GetFields(chunks []map[string]interface{}, _ []string) map[string]map[string]interface{} { + fields := make(map[string]map[string]interface{}, len(chunks)) + for _, chunk := range chunks { + if id, ok := chunk["id"].(string); ok { + fields[id] = chunk + } + } + return fields +} + +func (e *retrievalCountEngine) KNNScores(_ context.Context, chunks []map[string]interface{}, _ []float64, _ int) (map[string]interface{}, error) { + scores := make(map[string]interface{}, len(chunks)) + for _, chunk := range chunks { + id, _ := chunk["id"].(string) + score, _ := chunk["_score"].(float64) + scores[id] = score + } + return scores, nil +} + +func (e *retrievalCountEngine) GetScores(result map[string]interface{}) map[string]float64 { + scores := make(map[string]float64, len(result)) + for id, raw := range result { + if score, ok := raw.(float64); ok { + scores[id] = score + } + } + return scores +} + +func (e *retrievalCountEngine) DropChunkStore(context.Context, string, string) error { return nil } +func (e *retrievalCountEngine) ChunkStoreExists(context.Context, string, string) (bool, error) { + return true, nil +} +func (e *retrievalCountEngine) Close() error { return nil } +func (e *retrievalCountEngine) Ping(context.Context) error { return nil } +func (e *retrievalCountEngine) GetType() string { return "elasticsearch" } +func (e *retrievalCountEngine) SupportsPageRank() bool { return false } +func (e *retrievalCountEngine) CreateChunkStore(context.Context, string, string, int, string) error { + return nil +} +func (e *retrievalCountEngine) InsertChunks(context.Context, []map[string]interface{}, string, string) ([]string, error) { + return nil, nil +} +func (e *retrievalCountEngine) UpdateChunks(context.Context, map[string]interface{}, map[string]interface{}, string, string) error { + return nil +} +func (e *retrievalCountEngine) DeleteChunks(context.Context, map[string]interface{}, string, string) (int64, error) { + return 0, nil +} +func (e *retrievalCountEngine) GetChunk(context.Context, string, string, []string) (interface{}, error) { + return nil, nil +} +func (e *retrievalCountEngine) CreateMetadataStore(context.Context, string) error { return nil } +func (e *retrievalCountEngine) InsertMetadata(context.Context, []map[string]interface{}, string) ([]string, error) { + return nil, nil +} +func (e *retrievalCountEngine) UpdateMetadata(context.Context, string, string, map[string]interface{}, string) error { + return nil +} +func (e *retrievalCountEngine) DeleteMetadata(context.Context, map[string]interface{}, string) (int64, error) { + return 0, nil +} +func (e *retrievalCountEngine) DeleteMetadataKeys(context.Context, string, string, []string, string) error { + return nil +} +func (e *retrievalCountEngine) DropMetadataStore(context.Context, string) error { return nil } +func (e *retrievalCountEngine) MetadataStoreExists(context.Context, string) (bool, error) { + return true, nil +} +func (e *retrievalCountEngine) SearchMetadata(context.Context, *types.SearchMetadataRequest) (*types.SearchMetadataResult, error) { + return nil, nil +} +func (e *retrievalCountEngine) IndexDocument(context.Context, string, string, interface{}) error { + return nil +} +func (e *retrievalCountEngine) DeleteDocument(context.Context, string, string) error { return nil } +func (e *retrievalCountEngine) BulkIndex(context.Context, string, []interface{}) (interface{}, error) { + return nil, nil +} +func (e *retrievalCountEngine) GetAggregation([]map[string]interface{}, string) []map[string]interface{} { + return nil +} +func (e *retrievalCountEngine) GetHighlight([]map[string]interface{}, []string, string) map[string]string { + return nil +} +func (e *retrievalCountEngine) RunSQL(context.Context, string, string, []string, string) ([]map[string]interface{}, error) { + return nil, nil +} +func (e *retrievalCountEngine) FilterDocIdsByMetaPushdown(context.Context, *gorm.DB, []string, []map[string]interface{}, string) []string { + return nil +}