fix: return total numbers of the search result (#18071)

### Summary As title
2026-08-15 13:14:28 +08:00 · 2026-08-14 17:44:38 +08:00
parent d77b514d3e
commit 5bb5ba2216
2 changed files with 313 additions and 129 deletions
--- 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
--- a/internal/service/nlp/retrieval_test.go
+++ 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
+}