rig_neo4j/
vector_index.rs

1//! A vector index for a Neo4j graph DB.
2//!
3//! This module provides a way to perform vector searches on a Neo4j graph DB.
4//! It uses the [Neo4j vector index](https://neo4j.com/docs/cypher-manual/current/indexes/semantic-indexes/vector-indexes/)
5//! to search for similar nodes based on a query.
6
7use neo4rs::{Graph, Query};
8use rig_core::{
9    Embed,
10    embeddings::{Embedding, EmbeddingModel},
11    vector_store::{
12        InsertDocuments, VectorStoreError, VectorStoreIndex,
13        request::{SearchFilter, VectorSearchRequest},
14    },
15};
16use serde::{Deserialize, Serialize, de::Error};
17
18use crate::{Neo4jClient, Neo4jSearchFilter, ToBoltType};
19
20pub struct Neo4jVectorIndex<M>
21where
22    M: EmbeddingModel,
23{
24    graph: Graph,
25    embedding_model: M,
26    index_config: IndexConfig,
27}
28
29/// The index name must be unique among both indexes and constraints.
30/// A newly created index is not immediately available but is created in the background.
31///
32/// #### Default Values
33/// - `index_name`: "vector_index"
34/// - `embedding_property`: "embedding"
35/// - `similarity_function`: VectorSimilarityFunction::Cosine
36/// - `node_label`: None (inserts default to the `Document` label)
37#[derive(Serialize, Deserialize, Clone)]
38pub struct IndexConfig {
39    pub index_name: String,
40    pub embedding_property: String,
41    pub similarity_function: VectorSimilarityFunction,
42    /// The node label that [`InsertDocuments`] writes to (and that the index
43    /// applies to). Populated from the index's `labelsOrTypes` when loaded via
44    /// [`Neo4jClient::get_index`](crate::Neo4jClient::get_index).
45    pub node_label: Option<String>,
46}
47
48impl Default for IndexConfig {
49    fn default() -> Self {
50        Self {
51            index_name: "vector_index".to_string(),
52            embedding_property: "embedding".to_string(),
53            similarity_function: VectorSimilarityFunction::Cosine,
54            node_label: None,
55        }
56    }
57}
58
59impl IndexConfig {
60    pub fn new(index_name: impl Into<String>) -> Self {
61        Self {
62            index_name: index_name.into(),
63            embedding_property: "embedding".to_string(),
64            similarity_function: VectorSimilarityFunction::Cosine,
65            node_label: None,
66        }
67    }
68
69    pub fn index_name(mut self, index_name: &str) -> Self {
70        self.index_name = index_name.to_string();
71        self
72    }
73
74    pub fn similarity_function(mut self, similarity_function: VectorSimilarityFunction) -> Self {
75        self.similarity_function = similarity_function;
76        self
77    }
78
79    pub fn embedding_property(mut self, embedding_property: &str) -> Self {
80        self.embedding_property = embedding_property.to_string();
81        self
82    }
83
84    /// Sets the node label that [`InsertDocuments`] writes to.
85    pub fn node_label(mut self, node_label: &str) -> Self {
86        self.node_label = Some(node_label.to_string());
87        self
88    }
89}
90
91/// Cosine is most commonly used, but Euclidean is also supported.
92/// See [Neo4j vector similarity functions](https://neo4j.com/docs/cypher-manual/current/indexes/semantic-indexes/vector-indexes/#similarity-functions)
93/// for more information.
94#[derive(Default, Serialize, Deserialize, Clone)]
95#[serde(rename_all = "lowercase")]
96pub enum VectorSimilarityFunction {
97    #[default]
98    Cosine,
99    Euclidean,
100}
101
102use std::str::FromStr;
103
104impl FromStr for VectorSimilarityFunction {
105    type Err = VectorStoreError;
106
107    fn from_str(s: &str) -> Result<Self, VectorStoreError> {
108        match s.to_lowercase().as_str() {
109            "cosine" => Ok(VectorSimilarityFunction::Cosine),
110            "euclidean" => Ok(VectorSimilarityFunction::Euclidean),
111            _ => Err(VectorStoreError::JsonError(serde_json::Error::custom(
112                format!("Invalid similarity function: {s}"),
113            ))),
114        }
115    }
116}
117
118const BASE_VECTOR_SEARCH_QUERY: &str = "
119    CALL db.index.vector.queryNodes($index_name, $num_candidates, $queryVector)
120    YIELD node, score
121";
122
123impl<M> Neo4jVectorIndex<M>
124where
125    M: EmbeddingModel,
126{
127    pub fn new(graph: Graph, embedding_model: M, index_config: IndexConfig) -> Self {
128        Self {
129            graph,
130            embedding_model,
131            index_config,
132        }
133    }
134
135    /// Build a Neo4j query that performs a vector search against an index.
136    /// See [Query vector index](https://neo4j.com/docs/cypher-manual/current/indexes/semantic-indexes/vector-indexes/#query-vector-index) for more information.
137    ///
138    /// Query template:
139    /// ```text
140    /// CALL db.index.vector.queryNodes($index_name, $num_candidates, $queryVector)
141    /// YIELD node, score
142    /// WHERE {where_clause}
143    /// RETURN score, ID(node) as element_id, node {.*, embedding:null } as node
144    /// ```
145    pub fn build_vector_search_query(
146        &self,
147        prompt_embedding: Embedding,
148        return_node: bool,
149        req: &VectorSearchRequest<Neo4jSearchFilter>,
150    ) -> Query {
151        let where_clause = match (req.threshold(), req.filter()) {
152            (Some(thresh), Some(filt)) => Neo4jSearchFilter::gt("distance", thresh.into())
153                .and(filt.clone())
154                .render(),
155            (Some(thresh), _) => Neo4jSearchFilter::gt("distance", thresh.into()).render(),
156            (_, Some(filt)) => filt.clone().render(),
157            _ => String::new(),
158        };
159
160        // Propertiy containing the embedding vectors are excluded from the returned node
161        let query = format!(
162            "\
163            {}\
164            \t{}\n\
165            \tRETURN score, ID(node) as element_id {}
166            ",
167            BASE_VECTOR_SEARCH_QUERY,
168            where_clause,
169            if return_node {
170                format!(
171                    ", node {{.*, {}:null }} as node",
172                    self.index_config.embedding_property
173                )
174            } else {
175                "".to_string()
176            }
177        );
178
179        tracing::debug!("Query before params: {}", query);
180
181        Query::new(query)
182            .param("queryVector", prompt_embedding.vec)
183            .param("num_candidates", req.samples() as i64)
184            .param("index_name", self.index_config.index_name.clone())
185    }
186
187    /// Embeds the query and runs the vector search, deserializing each row as `R`.
188    async fn run_search<R: for<'a> Deserialize<'a>>(
189        &self,
190        req: &VectorSearchRequest<Neo4jSearchFilter>,
191    ) -> Result<Vec<R>, VectorStoreError> {
192        let prompt_embedding = self.embedding_model.embed_text(req.query()).await?;
193        let query = self.build_vector_search_query(prompt_embedding, true, req);
194
195        Neo4jClient::execute_and_collect::<R>(&self.graph, query).await
196    }
197}
198
199#[derive(Debug, Deserialize)]
200pub struct RowResultNode<T> {
201    score: f64,
202    element_id: i64,
203    node: T,
204}
205
206#[derive(Debug, Deserialize)]
207struct RowResult {
208    score: f64,
209    element_id: i64,
210}
211
212impl<M> VectorStoreIndex for Neo4jVectorIndex<M>
213where
214    M: EmbeddingModel + std::marker::Sync + Send,
215{
216    type Filter = Neo4jSearchFilter;
217
218    /// Get the top n nodes and scores matching the query.
219    ///
220    /// #### Generic Type Parameters
221    ///
222    /// - `T`: The type used to deserialize the result from the Neo4j query.
223    ///   It must implement the `serde::Deserialize` trait.
224    ///
225    /// #### Returns
226    ///
227    /// Returns a `Result` containing a vector of tuples. Each tuple contains:
228    /// - A `f64` representing the similarity score
229    /// - A `String` representing the node ID
230    /// - A value of type `T` representing the deserialized node data
231    ///
232    async fn top_n<T: for<'a> Deserialize<'a> + std::marker::Send>(
233        &self,
234        req: VectorSearchRequest<Neo4jSearchFilter>,
235    ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
236        let rows = self.run_search::<RowResultNode<T>>(&req).await?;
237
238        Ok(rows
239            .into_iter()
240            .map(|row| (row.score, row.element_id.to_string(), row.node))
241            .collect())
242    }
243
244    /// Get the top n ids and scores matching the query. Runs faster than top_n since it doesn't need to transfer and parse
245    /// the full nodes and embeddings to the client.
246    async fn top_n_ids(
247        &self,
248        req: VectorSearchRequest<Neo4jSearchFilter>,
249    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
250        let rows = self.run_search::<RowResult>(&req).await?;
251
252        Ok(rows
253            .into_iter()
254            .map(|row| (row.score, row.element_id.to_string()))
255            .collect())
256    }
257}
258
259/// The node label [`InsertDocuments`] writes to when the index config does not
260/// specify one (i.e. `node_label` is `None`).
261const DEFAULT_NODE_LABEL: &str = "Document";
262
263/// The Cypher used to bulk-insert nodes from an `$items` parameter list.
264fn insert_documents_query(node_label: &str) -> String {
265    format!("UNWIND $items AS item CREATE (n:{node_label}) SET n = item")
266}
267
268impl<M> InsertDocuments for Neo4jVectorIndex<M>
269where
270    M: EmbeddingModel + Send + Sync,
271{
272    /// Inserts one node per embedding, flattening the document's JSON fields
273    /// onto the node alongside the embedding (`embedding_property`) and its
274    /// source text (`embedded_text`). Nodes are written under the index's
275    /// `node_label`, defaulting to the `Document` label.
276    async fn insert_documents<Doc: Serialize + Embed + Send>(
277        &self,
278        documents: Vec<(Doc, Vec<Embedding>)>,
279    ) -> Result<(), VectorStoreError> {
280        let node_label = self
281            .index_config
282            .node_label
283            .as_deref()
284            .unwrap_or(DEFAULT_NODE_LABEL);
285        let embedding_property = &self.index_config.embedding_property;
286
287        // Build one parameter map per embedding for a single UNWIND insert.
288        let mut items: Vec<neo4rs::BoltType> = Vec::new();
289        for (document, embeddings) in documents {
290            let json_doc = serde_json::to_value(&document)?;
291
292            for embedding in embeddings {
293                let mut props = neo4rs::BoltMap::new();
294                if let serde_json::Value::Object(map) = &json_doc {
295                    for (key, value) in map {
296                        props.put(neo4rs::BoltString::new(key), value.to_bolt_type());
297                    }
298                } else {
299                    props.put(neo4rs::BoltString::new("document"), json_doc.to_bolt_type());
300                }
301                props.put(
302                    neo4rs::BoltString::new("embedded_text"),
303                    neo4rs::BoltType::String(neo4rs::BoltString::new(&embedding.document)),
304                );
305                props.put(
306                    neo4rs::BoltString::new(embedding_property),
307                    embedding.vec.to_bolt_type(),
308                );
309                items.push(neo4rs::BoltType::Map(props));
310            }
311        }
312
313        self.graph
314            .run(neo4rs::query(&insert_documents_query(node_label)).param("items", items))
315            .await
316            .map_err(VectorStoreError::datastore)?;
317
318        Ok(())
319    }
320}
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325
326    #[test]
327    fn node_label_defaults_to_none_and_builder_sets_it() {
328        assert_eq!(IndexConfig::new("idx").node_label, None);
329        assert_eq!(
330            IndexConfig::new("idx")
331                .node_label("Movie")
332                .node_label
333                .as_deref(),
334            Some("Movie"),
335        );
336    }
337
338    #[test]
339    fn insert_documents_query_uses_label_else_default() {
340        assert_eq!(
341            insert_documents_query("Movie"),
342            "UNWIND $items AS item CREATE (n:Movie) SET n = item",
343        );
344        assert_eq!(
345            insert_documents_query(DEFAULT_NODE_LABEL),
346            "UNWIND $items AS item CREATE (n:Document) SET n = item",
347        );
348    }
349}
rig_neo4j/vector_index.rs

rig_neo4j/
vector_index.rs
