agentdb/
hybrid.rs

1use crate::error::Result;
2use crate::memory::MemoryGraph;
3use crate::memory::TraversalOptions;
4use crate::vectors::collection::{Collection, SearchOptions, SearchResult};
5use crate::vectors::hnsw::DistanceMetric;
6use rusqlite::Connection;
7use serde_json::Value;
8use std::collections::HashMap;
9use std::sync::{Arc, Mutex};
10
11/// A single result from a hybrid graph + vector query.
12#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
13pub struct HybridResult {
14    /// ID of the matched entity.
15    pub id: String,
16    /// Raw cosine distance from the vector ANN search (lower = more similar).
17    pub vector_score: f32,
18    /// Maximum edge weight along any graph path from the anchor to this node.
19    /// `0.0` if the node is not reachable from the anchor.
20    pub graph_weight: f64,
21    /// Final blended rank score: `alpha × vec_similarity + (1 − alpha) × graph_weight`.
22    /// Higher is better.
23    pub rank_score: f64,
24    /// Metadata stored alongside the vector, if any.
25    pub metadata: Option<Value>,
26}
27
28/// Parameters for a hybrid graph + vector query.
29pub struct HybridQuery<'a> {
30    /// The memory-graph node to start traversal from.
31    pub anchor_node: &'a str,
32    /// Query embedding to rank against the vector collection.
33    pub embedding: &'a [f32],
34    /// Name of the vector collection to search.
35    pub collection: &'a str,
36    /// Maximum graph traversal depth from `anchor_node`.
37    pub graph_depth: usize,
38    /// Maximum number of results to return after blending.
39    pub top_k: usize,
40    /// Interpolation factor between vector similarity and graph proximity.
41    /// `0.0` = pure graph weight, `1.0` = pure vector similarity.
42    pub alpha: f64,
43    /// Optional metadata filter applied before vector scoring.
44    pub filter: Option<Value>,
45}
46
47/// Executes hybrid graph + vector queries.
48pub struct HybridStore {
49    conn: Arc<Mutex<Connection>>,
50}
51
52impl HybridStore {
53    pub(crate) fn new(conn: Arc<Mutex<Connection>>) -> Self {
54        Self { conn }
55    }
56
57    /// Run a hybrid graph + vector query.
58    ///
59    /// The algorithm proceeds in three stages:
60    ///
61    /// 1. **Graph traversal** — walks the memory graph from `q.anchor_node` up to
62    ///    `q.graph_depth` hops, recording the maximum edge weight seen for each
63    ///    reachable node.
64    /// 2. **Vector search** — retrieves the top `q.top_k × 20` approximate nearest
65    ///    neighbours from the named collection.
66    /// 3. **Score blending** — for each candidate, computes
67    ///    `rank = q.alpha × vec_similarity + (1 − q.alpha) × graph_weight`,
68    ///    then returns the top `q.top_k` results sorted by rank descending.
69    pub fn query(&self, q: HybridQuery, col: &Collection) -> Result<Vec<HybridResult>> {
70        // Step 1: graph traversal
71        let graph = MemoryGraph::new(Arc::clone(&self.conn));
72        let traversal = graph
73            .neighbors(
74                q.anchor_node,
75                TraversalOptions {
76                    relation: None,
77                    max_depth: q.graph_depth,
78                    min_weight: Some(0.0),
79                },
80            )
81            .unwrap_or_default();
82
83        let mut graph_weights: HashMap<String, f64> = HashMap::new();
84        for t in &traversal {
85            let e = graph_weights.entry(t.node.id.clone()).or_insert(0.0);
86            if t.weight > *e {
87                *e = t.weight;
88            }
89        }
90
91        // Step 2: vector search
92        let fetch_k = (q.top_k * 20).max(100);
93        let vec_results: Vec<SearchResult> = col.search(
94            q.embedding,
95            SearchOptions {
96                top_k: fetch_k,
97                metric: DistanceMetric::Cosine,
98                filter: q.filter.clone(),
99            },
100        )?;
101
102        if vec_results.is_empty() {
103            return Ok(vec![]);
104        }
105
106        // Step 3: normalize vector scores (distance -> similarity)
107        let max_s = vec_results
108            .iter()
109            .map(|r| r.score)
110            .fold(f32::NEG_INFINITY, f32::max);
111        let min_s = vec_results
112            .iter()
113            .map(|r| r.score)
114            .fold(f32::INFINITY, f32::min);
115        let range = (max_s - min_s).max(1e-6);
116
117        // Step 4: blend and rank
118        let mut blended: Vec<HybridResult> = vec_results
119            .into_iter()
120            .map(|r| {
121                let vec_sim = 1.0 - ((r.score - min_s) / range) as f64;
122                let gw = graph_weights.get(&r.id).copied().unwrap_or(0.0);
123                let rank = q.alpha * vec_sim + (1.0 - q.alpha) * gw;
124                HybridResult {
125                    id: r.id,
126                    vector_score: r.score,
127                    graph_weight: gw,
128                    rank_score: rank,
129                    metadata: r.metadata,
130                }
131            })
132            .collect();
133
134        blended.sort_by(|a, b| {
135            b.rank_score
136                .partial_cmp(&a.rank_score)
137                .unwrap_or(std::cmp::Ordering::Equal)
138        });
139        blended.truncate(q.top_k);
140        Ok(blended)
141    }
142}
agentdb/hybrid.rs

agentdb/
hybrid.rs
