Skip to main content

nodedb_lite/nodedb/graph_rag/
hybrid.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use super::{HybridSearchParams, fusion::search_results_to_ranked};
4use crate::{nodedb::NodeDbLite, storage::engine::StorageEngine};
5use nodedb_client::NodeDb;
6use nodedb_types::{error::NodeDbResult, result::SearchResult};
7use std::collections::{HashMap, HashSet};
8
9impl<S: StorageEngine> NodeDbLite<S> {
10    /// Hybrid search: vector similarity + BM25 text relevance fused via RRF.
11    ///
12    /// 1. Vector search returns `vector_k` candidates by embedding similarity.
13    /// 2. Text search returns `text_k` candidates by BM25 relevance.
14    /// 3. RRF fuses both rankings into a single score per document.
15    /// 4. Top `top_k` results returned.
16    ///
17    /// Documents found by both searches get boosted (RRF scores are additive).
18    pub async fn hybrid_search(
19        &self,
20        params: &HybridSearchParams<'_>,
21    ) -> NodeDbResult<Vec<SearchResult>> {
22        if params.top_k == 0 || params.allowed_ids.is_some_and(HashSet::is_empty) {
23            return Ok(Vec::new());
24        }
25        let vector_results = if params.vector_k == 0 {
26            Vec::new()
27        } else {
28            self.vector_search(
29                params.collection,
30                params.query_embedding,
31                params.vector_k,
32                params.filter,
33                params.allowed_ids,
34            )
35            .await?
36        };
37
38        let text_results = if params.text_k == 0 {
39            Vec::new()
40        } else {
41            self.text_search(
42                params.collection,
43                params.text_field,
44                params.query_text,
45                params.text_k,
46                nodedb_types::TextSearchParams::default(),
47                params.allowed_ids,
48            )
49            .await?
50        };
51
52        // RRF fusion via shared module.
53        let vector_ranked = search_results_to_ranked(&vector_results, "vector");
54        let text_ranked = search_results_to_ranked(&text_results, "text");
55        let fused = nodedb_query::fusion::reciprocal_rank_fusion(
56            &[vector_ranked, text_ranked],
57            None,
58            params.top_k,
59        );
60
61        // Build metadata cache for result materialization.
62        let mut metadata_cache: HashMap<&str, &HashMap<String, nodedb_types::Value>> =
63            HashMap::new();
64        for results in [&vector_results, &text_results] {
65            for result in results.iter() {
66                metadata_cache
67                    .entry(result.id.as_str())
68                    .or_insert(&result.metadata);
69            }
70        }
71
72        Ok(fused
73            .into_iter()
74            .map(|f| {
75                let metadata = metadata_cache
76                    .get(f.document_id.as_str())
77                    .copied()
78                    .cloned()
79                    .unwrap_or_default();
80                SearchResult {
81                    id: f.document_id,
82                    node_id: None,
83                    distance: 1.0 / (1.0 + f.rrf_score as f32),
84                    metadata,
85                }
86            })
87            .collect())
88    }
89}