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, OneOrMany,
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
188/// Search parameters for a vector search. Neo4j currently only supports post-vector-search filtering.
189pub struct SearchParams {
190    /// Sets the **post-filter** field of the search params. Uses a WHERE clause.
191    /// See [Neo4j WHERE clause](https://neo4j.com/docs/cypher-manual/current/clauses/where/) for more information.
192    post_vector_search_filter: Option<String>,
193}
194
195impl SearchParams {
196    /// Initializes a new `SearchParams` with default values.
197    pub fn new(filter: Option<String>) -> Self {
198        Self {
199            post_vector_search_filter: filter,
200        }
201    }
202
203    pub fn filter(mut self, filter: String) -> Self {
204        self.post_vector_search_filter = Some(filter);
205        self
206    }
207}
208
209impl Default for SearchParams {
210    fn default() -> Self {
211        Self::new(None)
212    }
213}
214
215#[derive(Debug, Deserialize)]
216pub struct RowResultNode<T> {
217    score: f64,
218    element_id: i64,
219    node: T,
220}
221
222#[derive(Debug, Deserialize)]
223struct RowResult {
224    score: f64,
225    element_id: i64,
226}
227
228impl<M> VectorStoreIndex for Neo4jVectorIndex<M>
229where
230    M: EmbeddingModel + std::marker::Sync + Send,
231{
232    type Filter = Neo4jSearchFilter;
233
234    /// Get the top n nodes and scores matching the query.
235    ///
236    /// #### Generic Type Parameters
237    ///
238    /// - `T`: The type used to deserialize the result from the Neo4j query.
239    ///   It must implement the `serde::Deserialize` trait.
240    ///
241    /// #### Returns
242    ///
243    /// Returns a `Result` containing a vector of tuples. Each tuple contains:
244    /// - A `f64` representing the similarity score
245    /// - A `String` representing the node ID
246    /// - A value of type `T` representing the deserialized node data
247    ///
248    async fn top_n<T: for<'a> Deserialize<'a> + std::marker::Send>(
249        &self,
250        req: VectorSearchRequest<Neo4jSearchFilter>,
251    ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
252        let prompt_embedding = self.embedding_model.embed_text(req.query()).await?;
253        let query = self.build_vector_search_query(prompt_embedding, true, &req);
254
255        let rows = Neo4jClient::execute_and_collect::<RowResultNode<T>>(&self.graph, query).await?;
256
257        let results = rows
258            .into_iter()
259            .map(|row| (row.score, row.element_id.to_string(), row.node))
260            .collect::<Vec<_>>();
261
262        Ok(results)
263    }
264
265    /// Get the top n ids and scores matching the query. Runs faster than top_n since it doesn't need to transfer and parse
266    /// the full nodes and embeddings to the client.
267    async fn top_n_ids(
268        &self,
269        req: VectorSearchRequest<Neo4jSearchFilter>,
270    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
271        let prompt_embedding = self.embedding_model.embed_text(req.query()).await?;
272
273        let query = self.build_vector_search_query(prompt_embedding, true, &req);
274
275        let rows = Neo4jClient::execute_and_collect::<RowResult>(&self.graph, query).await?;
276
277        let results = rows
278            .into_iter()
279            .map(|row| (row.score, row.element_id.to_string()))
280            .collect::<Vec<_>>();
281
282        Ok(results)
283    }
284}
285
286/// The node label [`InsertDocuments`] writes to when the index config does not
287/// specify one (i.e. `node_label` is `None`).
288const DEFAULT_NODE_LABEL: &str = "Document";
289
290/// The Cypher used to bulk-insert nodes from an `$items` parameter list.
291fn insert_documents_query(node_label: &str) -> String {
292    format!("UNWIND $items AS item CREATE (n:{node_label}) SET n = item")
293}
294
295impl<M> InsertDocuments for Neo4jVectorIndex<M>
296where
297    M: EmbeddingModel + Send + Sync,
298{
299    /// Inserts one node per embedding, flattening the document's JSON fields
300    /// onto the node alongside the embedding (`embedding_property`) and its
301    /// source text (`embedded_text`). Nodes are written under the index's
302    /// `node_label`, defaulting to [`DEFAULT_NODE_LABEL`].
303    async fn insert_documents<Doc: Serialize + Embed + Send>(
304        &self,
305        documents: Vec<(Doc, OneOrMany<Embedding>)>,
306    ) -> Result<(), VectorStoreError> {
307        let node_label = self
308            .index_config
309            .node_label
310            .as_deref()
311            .unwrap_or(DEFAULT_NODE_LABEL);
312        let embedding_property = &self.index_config.embedding_property;
313
314        // Build one parameter map per embedding for a single UNWIND insert.
315        let mut items: Vec<neo4rs::BoltType> = Vec::new();
316        for (document, embeddings) in documents {
317            let json_doc = serde_json::to_value(&document)?;
318
319            for embedding in embeddings {
320                let mut props = neo4rs::BoltMap::new();
321                if let serde_json::Value::Object(map) = &json_doc {
322                    for (key, value) in map {
323                        props.put(neo4rs::BoltString::new(key), value.to_bolt_type());
324                    }
325                } else {
326                    props.put(neo4rs::BoltString::new("document"), json_doc.to_bolt_type());
327                }
328                props.put(
329                    neo4rs::BoltString::new("embedded_text"),
330                    neo4rs::BoltType::String(neo4rs::BoltString::new(&embedding.document)),
331                );
332                props.put(
333                    neo4rs::BoltString::new(embedding_property),
334                    embedding.vec.to_bolt_type(),
335                );
336                items.push(neo4rs::BoltType::Map(props));
337            }
338        }
339
340        self.graph
341            .run(neo4rs::query(&insert_documents_query(node_label)).param("items", items))
342            .await
343            .map_err(|e| VectorStoreError::DatastoreError(Box::new(e)))?;
344
345        Ok(())
346    }
347}
348
349#[cfg(test)]
350mod tests {
351    use super::*;
352
353    #[test]
354    fn node_label_defaults_to_none_and_builder_sets_it() {
355        assert_eq!(IndexConfig::new("idx").node_label, None);
356        assert_eq!(
357            IndexConfig::new("idx")
358                .node_label("Movie")
359                .node_label
360                .as_deref(),
361            Some("Movie"),
362        );
363    }
364
365    #[test]
366    fn insert_documents_query_uses_label_else_default() {
367        assert_eq!(
368            insert_documents_query("Movie"),
369            "UNWIND $items AS item CREATE (n:Movie) SET n = item",
370        );
371        assert_eq!(
372            insert_documents_query(DEFAULT_NODE_LABEL),
373            "UNWIND $items AS item CREATE (n:Document) SET n = item",
374        );
375    }
376}
rig_neo4j/vector_index.rs

rig_neo4j/
vector_index.rs
