nodedb_lite/nodedb/graph_rag/
graph.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use super::{GraphRagParams, fusion::search_results_to_ranked};
4use crate::engine::graph::{index::Direction, traversal::DEFAULT_MAX_VISITED};
5use crate::{
6    nodedb::{LockExt, NodeDbLite},
7    storage::engine::StorageEngine,
8};
9use nodedb_client::NodeDb;
10use nodedb_types::{
11    error::{NodeDbError, NodeDbResult},
12    id::NodeId,
13    result::SearchResult,
14};
15use std::collections::{HashMap, HashSet, VecDeque};
16
17fn limit_error(collection: &str) -> NodeDbError {
18    NodeDbError::program_limit_exceeded(format!(
19        "graph retrieval exceeds {DEFAULT_MAX_VISITED} visited nodes in '{collection}': reduce seeds or graph_depth"
20    ))
21}
22
23impl<S: StorageEngine> NodeDbLite<S> {
24    /// Fuse vector retrieval with bounded, outgoing graph expansion.
25    pub async fn graph_rag_search(
26        &self,
27        params: &GraphRagParams<'_>,
28    ) -> NodeDbResult<Vec<SearchResult>> {
29        if params.top_k == 0 || params.allowed_ids.is_some_and(HashSet::is_empty) {
30            return Ok(Vec::new());
31        }
32        if !params.rrf_k.is_finite() || params.rrf_k < 0.0 {
33            return Err(NodeDbError::bad_request(format!(
34                "invalid rrf_k '{}': use a finite nonnegative value",
35                params.rrf_k
36            )));
37        }
38        let vector_results = if params.vector_k == 0 {
39            Vec::new()
40        } else {
41            self.vector_search(
42                params.collection,
43                params.query,
44                params.vector_k,
45                params.filter,
46                params.allowed_ids,
47            )
48            .await?
49        };
50        let mut unique_seeds = HashSet::new();
51        let mut caller_seeds = HashSet::new();
52        for (id, caller) in vector_results.iter().map(|r| (r.id.as_str(), false)).chain(
53            params
54                .seed_nodes
55                .into_iter()
56                .flatten()
57                .map(|id| (id.as_str(), true)),
58        ) {
59            if params
60                .allowed_ids
61                .is_some_and(|allowed| !allowed.contains(id))
62            {
63                continue;
64            }
65            if !unique_seeds.contains(id) {
66                if unique_seeds.len() == DEFAULT_MAX_VISITED {
67                    return Err(limit_error(params.collection));
68                }
69                unique_seeds.insert(id);
70            }
71            if caller {
72                caller_seeds.insert(id);
73            }
74        }
75        let mut seeds: Vec<String> = unique_seeds.into_iter().map(str::to_owned).collect();
76        seeds.sort_unstable();
77        let graph_ranked = self
78            .expand_rag_nodes(params, seeds)?
79            .into_iter()
80            .filter(|(id, depth)| *depth > 0 || caller_seeds.contains(id.as_str()))
81            .enumerate()
82            .map(|(rank, (id, depth))| nodedb_query::fusion::RankedResult {
83                document_id: id,
84                rank,
85                score: depth as f32,
86                source: "graph",
87            })
88            .collect();
89        let fused = nodedb_query::fusion::reciprocal_rank_fusion_weighted(
90            &[
91                search_results_to_ranked(&vector_results, "vector"),
92                graph_ranked,
93            ],
94            &[params.rrf_k, params.rrf_k],
95            params.top_k,
96        );
97        let vector_map: HashMap<&str, &SearchResult> =
98            vector_results.iter().map(|r| (r.id.as_str(), r)).collect();
99        let crdt = self.crdt.lock_or_recover();
100        Ok(fused
101            .into_iter()
102            .map(|f| {
103                let (distance, metadata) = match vector_map.get(f.document_id.as_str()) {
104                    Some(vr) => (vr.distance, vr.metadata.clone()),
105                    None => {
106                        let metadata = crdt
107                            .read(params.collection, &f.document_id)
108                            .map(|val| {
109                                crate::nodedb::convert::loro_value_to_document(&f.document_id, &val)
110                                    .fields
111                            })
112                            .unwrap_or_default();
113                        (f.rrf_score as f32, metadata)
114                    }
115                };
116                SearchResult {
117                    node_id: Some(NodeId::from_validated(f.document_id.clone())),
118                    id: f.document_id,
119                    distance,
120                    metadata,
121                }
122            })
123            .collect())
124    }
125
126    fn expand_rag_nodes(
127        &self,
128        params: &GraphRagParams<'_>,
129        seeds: Vec<String>,
130    ) -> NodeDbResult<Vec<(String, u8)>> {
131        let mut seen: HashSet<String> = seeds.iter().cloned().collect();
132        let mut queue: VecDeque<(String, u8)> = seeds.into_iter().map(|id| (id, 0)).collect();
133        let mut nodes = Vec::new();
134        let csr_map = self.csr.lock_or_recover();
135        let csr = csr_map.get(params.collection);
136        while let Some((id, depth)) = queue.pop_front() {
137            if depth < params.graph_depth
138                && let Some(csr) = csr
139            {
140                let mut neighbors = csr.neighbors(&id, &[], Direction::Out);
141                neighbors.sort_unstable_by(|a, b| a.1.cmp(&b.1));
142                for (_, neighbor) in neighbors {
143                    if params
144                        .allowed_ids
145                        .is_some_and(|allowed| !allowed.contains(&neighbor))
146                        || seen.contains(&neighbor)
147                    {
148                        continue;
149                    }
150                    if seen.len() == DEFAULT_MAX_VISITED {
151                        return Err(limit_error(params.collection));
152                    }
153                    seen.insert(neighbor.clone());
154                    queue.push_back((neighbor, depth + 1));
155                }
156            }
157            nodes.push((id, depth));
158        }
159        nodes.sort_unstable_by(|a, b| a.1.cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
160        Ok(nodes)
161    }
162}
nodedb_lite/nodedb/graph_rag/graph.rs

nodedb_lite/nodedb/graph_rag/
graph.rs
