nodedb_lite/nodedb/graph_rag/
hybrid.rs1use 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 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 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 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}