Skip to main content

nodedb_lite/nodedb/graph_rag/
params.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use nodedb_types::{filter::MetadataFilter, id::NodeId};
4use std::collections::HashSet;
5
6/// GraphRAG fusion parameters.
7pub struct GraphRagParams<'a> {
8    /// Collection to search vectors in.
9    pub collection: &'a str,
10    /// Query embedding.
11    pub query: &'a [f32],
12    /// Number of initial vector candidates.
13    pub vector_k: usize,
14    /// Graph expansion depth from caller and vector seeds.
15    pub graph_depth: u8,
16    /// Final number of results to return after fusion.
17    pub top_k: usize,
18    /// Optional metadata filter for vector search.
19    pub filter: Option<&'a MetadataFilter>,
20    /// IDs permitted in every retrieval source.
21    pub allowed_ids: Option<&'a HashSet<String>>,
22    /// Finite nonnegative rank smoothing constant. Default: 60. Lower values emphasize top ranks.
23    pub rrf_k: f64,
24    /// Caller seeds added to vector-derived seeds.
25    pub seed_nodes: Option<&'a [NodeId]>,
26}
27
28impl Default for GraphRagParams<'_> {
29    fn default() -> Self {
30        Self {
31            collection: "",
32            query: &[],
33            vector_k: 10,
34            graph_depth: 2,
35            top_k: 10,
36            filter: None,
37            allowed_ids: None,
38            seed_nodes: None,
39            rrf_k: 60.0,
40        }
41    }
42}
43
44/// Hybrid search parameters.
45pub struct HybridSearchParams<'a> {
46    /// Collection to search.
47    pub collection: &'a str,
48    /// Query embedding for vector similarity.
49    pub query_embedding: &'a [f32],
50    /// Query text for BM25 relevance.
51    pub query_text: &'a str,
52    /// Field the text query is scoped to. Empty searches every string field.
53    pub text_field: &'a str,
54    /// Number of vector candidates.
55    pub vector_k: usize,
56    /// Number of text candidates.
57    pub text_k: usize,
58    /// Final number of results to return after fusion.
59    pub top_k: usize,
60    /// Optional metadata filter for vector search.
61    pub filter: Option<&'a MetadataFilter>,
62    /// IDs permitted in every retrieval source.
63    pub allowed_ids: Option<&'a HashSet<String>>,
64}
65
66impl Default for HybridSearchParams<'_> {
67    fn default() -> Self {
68        Self {
69            collection: "",
70            query_embedding: &[],
71            query_text: "",
72            text_field: "",
73            vector_k: 10,
74            text_k: 10,
75            top_k: 10,
76            filter: None,
77            allowed_ids: None,
78        }
79    }
80}