lc_vector_stores/
neo4j.rs

1// lc-vector-stores/src/neo4j.rs
2//! Neo4j vector store implementation.
3//!
4//! Uses Neo4j's vector index feature (available since Neo4j 5.11) for
5//! similarity search via the Cypher API over HTTP.
6//!
7//! # Example
8//!
9//! ```rust,ignore
10//! use lc_vector_stores::neo4j::{Neo4jVectorStore, Neo4jConfig};
11//!
12//! let config = Neo4jConfig::new("bolt://localhost:7687", "neo4j", "password", "my_index");
13//! let store = Neo4jVectorStore::new(config);
14//! store.add_documents(docs, embeddings).await?;
15//! let results = store.similarity_search(&query_embedding, 5).await?;
16//! ```
17
18use async_trait::async_trait;
19use serde::Deserialize;
20use serde_json::json;
21
22use crate::{Document, SearchResult, VectorStore, VectorStoreError};
23
24/// Neo4j vector store configuration.
25#[derive(Debug, Clone)]
26pub struct Neo4jConfig {
27    /// Neo4j URI (e.g., "bolt://localhost:7687" or "neo4j://localhost:7687").
28    pub uri: String,
29    /// Username.
30    pub username: String,
31    /// Password.
32    pub password: String,
33    /// Database name (default: "neo4j").
34    pub database: String,
35    /// Node label for vector documents (default: "Document").
36    pub node_label: String,
37    /// Vector index name.
38    pub index_name: String,
39    /// Embedding property name on the node (default: "embedding").
40    pub embedding_property: String,
41    /// Content property name on the node (default: "content").
42    pub content_property: String,
43    /// Metadata property name on the node (default: "metadata").
44    pub metadata_property: String,
45    /// ID property name on the node (default: "id").
46    pub id_property: String,
47}
48
49impl Neo4jConfig {
50    /// Creates a new Neo4jConfig.
51    pub fn new(
52        uri: impl Into<String>,
53        username: impl Into<String>,
54        password: impl Into<String>,
55        index_name: impl Into<String>,
56    ) -> Self {
57        Self {
58            uri: uri.into(),
59            username: username.into(),
60            password: password.into(),
61            database: "neo4j".to_string(),
62            node_label: "Document".to_string(),
63            index_name: index_name.into(),
64            embedding_property: "embedding".to_string(),
65            content_property: "content".to_string(),
66            metadata_property: "metadata".to_string(),
67            id_property: "id".to_string(),
68        }
69    }
70
71    /// Creates config from environment variables.
72    pub fn from_env_result() -> Result<Self, String> {
73        let uri = std::env::var("NEO4J_URI")
74            .map_err(|_| "NEO4J_URI environment variable not set".to_string())?;
75        let username = std::env::var("NEO4J_USERNAME")
76            .map_err(|_| "NEO4J_USERNAME environment variable not set".to_string())?;
77        let password = std::env::var("NEO4J_PASSWORD")
78            .map_err(|_| "NEO4J_PASSWORD environment variable not set".to_string())?;
79        let index_name = std::env::var("NEO4J_VECTOR_INDEX_NAME")
80            .map_err(|_| "NEO4J_VECTOR_INDEX_NAME environment variable not set".to_string())?;
81        let database = std::env::var("NEO4J_DATABASE").unwrap_or_else(|_| "neo4j".to_string());
82        Ok(Self {
83            uri,
84            username,
85            password,
86            database,
87            index_name,
88            ..Default::default()
89        })
90    }
91
92    /// Sets the database name.
93    pub fn with_database(mut self, database: impl Into<String>) -> Self {
94        self.database = database.into();
95        self
96    }
97
98    /// Sets the node label.
99    pub fn with_node_label(mut self, label: impl Into<String>) -> Self {
100        self.node_label = label.into();
101        self
102    }
103
104    /// Sets the embedding property name.
105    pub fn with_embedding_property(mut self, prop: impl Into<String>) -> Self {
106        self.embedding_property = prop.into();
107        self
108    }
109
110    /// Sets the content property name.
111    pub fn with_content_property(mut self, prop: impl Into<String>) -> Self {
112        self.content_property = prop.into();
113        self
114    }
115}
116
117impl Default for Neo4jConfig {
118    fn default() -> Self {
119        Self {
120            uri: "bolt://localhost:7687".to_string(),
121            username: "neo4j".to_string(),
122            password: String::new(),
123            database: "neo4j".to_string(),
124            node_label: "Document".to_string(),
125            index_name: "vector_index".to_string(),
126            embedding_property: "embedding".to_string(),
127            content_property: "content".to_string(),
128            metadata_property: "metadata".to_string(),
129            id_property: "id".to_string(),
130        }
131    }
132}
133
134/// Neo4j vector store.
135///
136/// Communicates with Neo4j via the HTTP transaction API.
137pub struct Neo4jVectorStore {
138    config: Neo4jConfig,
139    client: reqwest::Client,
140}
141
142impl std::fmt::Debug for Neo4jVectorStore {
143    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
144        f.debug_struct("Neo4jVectorStore")
145            .field("uri", &self.config.uri)
146            .field("index", &self.config.index_name)
147            .finish()
148    }
149}
150
151impl Neo4jVectorStore {
152    /// Creates a new Neo4jVectorStore with the given configuration.
153    pub fn new(config: Neo4jConfig) -> Self {
154        Self {
155            config,
156            client: reqwest::Client::new(),
157        }
158    }
159
160    /// Creates from environment variables.
161    pub fn from_env_result() -> Result<Self, String> {
162        Ok(Self::new(Neo4jConfig::from_env_result()?))
163    }
164
165    /// Builds the HTTP API URL for the transaction endpoint.
166    fn tx_url(&self) -> String {
167        // Convert bolt:// or neo4j:// to http:// for the REST API
168        let http_uri = self
169            .config
170            .uri
171            .replace("bolt://", "http://")
172            .replace("neo4j://", "http://")
173            .replace("bolt+s://", "https://")
174            .replace("neo4j+s://", "https://");
175        format!(
176            "{}/db/{}/tx/commit",
177            http_uri.trim_end_matches('/'),
178            self.config.database
179        )
180    }
181
182    /// Executes a Cypher query via the HTTP transaction API.
183    async fn run_query(
184        &self,
185        query: &str,
186        params: serde_json::Value,
187    ) -> Result<Neo4jResponse, VectorStoreError> {
188        let body = json!({
189            "statements": [{
190                "statement": query,
191                "parameters": params,
192            }]
193        });
194
195        let response = self
196            .client
197            .post(self.tx_url())
198            .header("Content-Type", "application/json")
199            .header(
200                "Authorization",
201                format!(
202                    "Basic {}",
203                    base64_encode(format!(
204                        "{}:{}",
205                        self.config.username, self.config.password
206                    ))
207                ),
208            )
209            .json(&body)
210            .send()
211            .await
212            .map_err(|e| VectorStoreError::ConnectionError(e.to_string()))?;
213
214        let status = response.status();
215        if !status.is_success() {
216            let error_text = response.text().await.unwrap_or_default();
217            return Err(VectorStoreError::ConnectionError(format!(
218                "HTTP {}: {}",
219                status, error_text
220            )));
221        }
222
223        let neo4j_response: Neo4jResponse = response
224            .json()
225            .await
226            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
227
228        // Check for Neo4j-level errors
229        if let Some(errors) = &neo4j_response.errors {
230            if !errors.is_empty() {
231                let msg = errors
232                    .iter()
233                    .map(|e| e.message.clone())
234                    .collect::<Vec<_>>()
235                    .join("; ");
236                return Err(VectorStoreError::StorageError(msg));
237            }
238        }
239
240        Ok(neo4j_response)
241    }
242}
243
244/// Base64 encoding helper (no external dependency needed).
245fn base64_encode(input: String) -> String {
246    const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
247    let bytes = input.as_bytes();
248    let mut result = String::new();
249    let mut i = 0;
250    while i < bytes.len() {
251        let b0 = bytes[i];
252        let b1 = if i + 1 < bytes.len() { bytes[i + 1] } else { 0 };
253        let b2 = if i + 2 < bytes.len() { bytes[i + 2] } else { 0 };
254
255        result.push(CHARSET[((b0 >> 2) & 0x3F) as usize] as char);
256        result.push(CHARSET[(((b0 << 4) | (b1 >> 4)) & 0x3F) as usize] as char);
257        result.push(if i + 1 < bytes.len() {
258            CHARSET[(((b1 << 2) | (b2 >> 6)) & 0x3F) as usize] as char
259        } else {
260            '='
261        });
262        result.push(if i + 2 < bytes.len() {
263            CHARSET[(b2 & 0x3F) as usize] as char
264        } else {
265            '='
266        });
267
268        i += 3;
269    }
270    result
271}
272
273// ---------------------------------------------------------------------------
274// Neo4j HTTP API response types
275// ---------------------------------------------------------------------------
276
277/// Neo4j transaction commit response.
278#[derive(Debug, Deserialize)]
279struct Neo4jResponse {
280    results: Vec<Neo4jResult>,
281    errors: Option<Vec<Neo4jError>>,
282}
283
284#[derive(Debug, Deserialize)]
285struct Neo4jResult {
286    #[allow(dead_code)]
287    columns: Vec<String>,
288    data: Vec<Neo4jRow>,
289}
290
291#[derive(Debug, Deserialize)]
292struct Neo4jRow {
293    row: Vec<serde_json::Value>,
294}
295
296#[derive(Debug, Deserialize)]
297struct Neo4jError {
298    message: String,
299}
300
301#[async_trait]
302impl VectorStore for Neo4jVectorStore {
303    async fn add_documents(
304        &self,
305        documents: Vec<Document>,
306        embeddings: Vec<Vec<f32>>,
307    ) -> Result<Vec<String>, VectorStoreError> {
308        if documents.len() != embeddings.len() {
309            return Err(VectorStoreError::EmbeddingError(
310                "Number of documents and embeddings must match".to_string(),
311            ));
312        }
313
314        let ids: Vec<String> = documents
315            .iter()
316            .map(|doc| doc.id.clone().unwrap_or_else(|| uuid::Uuid::new_v4().to_string()))
317            .collect();
318
319        // Build UNWIND Cypher for batch insert
320        let rows: Vec<serde_json::Value> = documents
321            .into_iter()
322            .zip(embeddings)
323            .zip(ids.iter())
324            .map(|((doc, vec), id)| {
325                let metadata: serde_json::Value = doc
326                    .metadata
327                    .iter()
328                    .map(|(k, v)| (k.clone(), json!(v)))
329                    .collect();
330                json!({
331                    "id": id,
332                    "content": doc.content,
333                    "embedding": vec,
334                    "metadata": metadata,
335                })
336            })
337            .collect();
338
339        let query = format!(
340            "UNWIND $rows AS row \
341             MERGE (n:{label} {{{id_prop}: row.id}}) \
342             SET n.{content_prop} = row.content, \
343                 n.{embedding_prop} = row.embedding, \
344                 n.{metadata_prop} = row.metadata",
345            label = self.config.node_label,
346            id_prop = self.config.id_property,
347            content_prop = self.config.content_property,
348            embedding_prop = self.config.embedding_property,
349            metadata_prop = self.config.metadata_property,
350        );
351
352        self.run_query(&query, json!({ "rows": rows }))
353            .await?;
354
355        Ok(ids)
356    }
357
358    async fn similarity_search(
359        &self,
360        query_embedding: &[f32],
361        k: usize,
362    ) -> Result<Vec<SearchResult>, VectorStoreError> {
363        // Use Neo4j's db.index.vector.queryNodes procedure
364        let query = format!(
365            "CALL db.index.vector.queryNodes($index_name, $k, $query_vector) \
366             YIELD node, score \
367             RETURN node.{id_prop} AS id, \
368                    node.{content_prop} AS content, \
369                    node.{metadata_prop} AS metadata, \
370                    score \
371             ORDER BY score DESC",
372            id_prop = self.config.id_property,
373            content_prop = self.config.content_property,
374            metadata_prop = self.config.metadata_property,
375        );
376
377        let params = json!({
378            "index_name": self.config.index_name,
379            "k": k,
380            "query_vector": query_embedding,
381        });
382
383        let response = self.run_query(&query, params).await?;
384
385        let result = response.results.first();
386        let Some(neo4j_result) = result else {
387            return Ok(Vec::new());
388        };
389
390        let mut search_results = Vec::new();
391        for row in &neo4j_result.data {
392            if row.row.len() >= 4 {
393                let id = row.row[0].as_str().unwrap_or_default().to_string();
394                let content = row.row[1].as_str().unwrap_or_default().to_string();
395                let score = row.row[3].as_f64().unwrap_or(0.0) as f32;
396
397                let mut doc = Document::new(content).with_id(id);
398
399                // Parse metadata from JSON object
400                if let Some(meta_obj) = row.row[2].as_object() {
401                    for (key, value) in meta_obj {
402                        if let Some(s) = value.as_str() {
403                            doc = doc.with_metadata(key, s);
404                        } else {
405                            doc = doc.with_metadata(key, value.to_string());
406                        }
407                    }
408                }
409
410                search_results.push(SearchResult {
411                    document: doc,
412                    score,
413                });
414            }
415        }
416
417        Ok(search_results)
418    }
419
420    async fn get_document(&self, id: &str) -> Result<Option<Document>, VectorStoreError> {
421        let query = format!(
422            "MATCH (n:{label} {{{id_prop}: $id}}) \
423             RETURN n.{content_prop} AS content, n.{metadata_prop} AS metadata",
424            label = self.config.node_label,
425            id_prop = self.config.id_property,
426            content_prop = self.config.content_property,
427            metadata_prop = self.config.metadata_property,
428        );
429
430        let response = self.run_query(&query, json!({ "id": id })).await?;
431
432        let result = response.results.first();
433        let Some(neo4j_result) = result else {
434            return Ok(None);
435        };
436
437        let row = neo4j_result.data.first();
438        let Some(row) = row else {
439            return Ok(None);
440        };
441
442        if row.row.is_empty() {
443            return Ok(None);
444        }
445
446        let content = row.row[0].as_str().unwrap_or_default().to_string();
447        let mut doc = Document::new(content).with_id(id);
448
449        if row.row.len() > 1 {
450            if let Some(meta_obj) = row.row[1].as_object() {
451                for (key, value) in meta_obj {
452                    if let Some(s) = value.as_str() {
453                        doc = doc.with_metadata(key, s);
454                    } else {
455                        doc = doc.with_metadata(key, value.to_string());
456                    }
457                }
458            }
459        }
460
461        Ok(Some(doc))
462    }
463
464    async fn get_embedding(&self, id: &str) -> Result<Option<Vec<f32>>, VectorStoreError> {
465        let query = format!(
466            "MATCH (n:{label} {{{id_prop}: $id}}) \
467             RETURN n.{embedding_prop} AS embedding",
468            label = self.config.node_label,
469            id_prop = self.config.id_property,
470            embedding_prop = self.config.embedding_property,
471        );
472
473        let response = self.run_query(&query, json!({ "id": id })).await?;
474
475        let result = response.results.first();
476        let Some(neo4j_result) = result else {
477            return Ok(None);
478        };
479
480        let row = neo4j_result.data.first();
481        let Some(row) = row else {
482            return Ok(None);
483        };
484
485        if row.row.is_empty() {
486            return Ok(None);
487        }
488
489        let embedding: Vec<f32> = row.row[0]
490            .as_array()
491            .map(|arr| {
492                arr.iter()
493                    .filter_map(|v| v.as_f64().map(|f| f as f32))
494                    .collect()
495            })
496            .unwrap_or_default();
497
498        if embedding.is_empty() {
499            Ok(None)
500        } else {
501            Ok(Some(embedding))
502        }
503    }
504
505    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError> {
506        let query = format!(
507            "MATCH (n:{label} {{{id_prop}: $id}}) \
508             DETACH DELETE n",
509            label = self.config.node_label,
510            id_prop = self.config.id_property,
511        );
512
513        self.run_query(&query, json!({ "id": id })).await?;
514        Ok(())
515    }
516
517    async fn count(&self) -> usize {
518        let query = format!(
519            "MATCH (n:{label}) RETURN count(n) AS cnt",
520            label = self.config.node_label,
521        );
522
523        let result = self.run_query(&query, json!({})).await;
524        match result {
525            Ok(response) => {
526                if let Some(neo4j_result) = response.results.first() {
527                    if let Some(row) = neo4j_result.data.first() {
528                        if let Some(cnt) = row.row.first() {
529                            return cnt.as_u64().unwrap_or(0) as usize;
530                        }
531                    }
532                }
533                0
534            }
535            Err(_) => 0,
536        }
537    }
538
539    async fn clear(&self) -> Result<(), VectorStoreError> {
540        let query = format!(
541            "MATCH (n:{label}) \
542             DETACH DELETE n",
543            label = self.config.node_label,
544        );
545
546        self.run_query(&query, json!({})).await?;
547        Ok(())
548    }
549}
550
551#[cfg(test)]
552mod tests {
553    use super::*;
554
555    #[test]
556    fn test_config_new() {
557        let config = Neo4jConfig::new("bolt://localhost:7687", "neo4j", "pass", "my_index");
558        assert_eq!(config.uri, "bolt://localhost:7687");
559        assert_eq!(config.username, "neo4j");
560        assert_eq!(config.password, "pass");
561        assert_eq!(config.index_name, "my_index");
562        assert_eq!(config.database, "neo4j");
563    }
564
565    #[test]
566    fn test_config_builder() {
567        let config = Neo4jConfig::new("bolt://localhost:7687", "neo4j", "pass", "idx")
568            .with_database("mydb")
569            .with_node_label("Chunk")
570            .with_embedding_property("vec")
571            .with_content_property("text");
572        assert_eq!(config.database, "mydb");
573        assert_eq!(config.node_label, "Chunk");
574        assert_eq!(config.embedding_property, "vec");
575        assert_eq!(config.content_property, "text");
576    }
577
578    #[test]
579    fn test_config_default() {
580        let config = Neo4jConfig::default();
581        assert_eq!(config.uri, "bolt://localhost:7687");
582        assert_eq!(config.node_label, "Document");
583        assert_eq!(config.embedding_property, "embedding");
584    }
585
586    #[test]
587    fn test_tx_url_bolt() {
588        let config = Neo4jConfig::new("bolt://localhost:7687", "neo4j", "pass", "idx");
589        let store = Neo4jVectorStore::new(config);
590        assert_eq!(
591            store.tx_url(),
592            "http://localhost:7687/db/neo4j/tx/commit"
593        );
594    }
595
596    #[test]
597    fn test_tx_url_neo4j_scheme() {
598        let config = Neo4jConfig::new("neo4j://host:7687", "neo4j", "pass", "idx");
599        let store = Neo4jVectorStore::new(config);
600        assert_eq!(
601            store.tx_url(),
602            "http://host:7687/db/neo4j/tx/commit"
603        );
604    }
605
606    #[test]
607    fn test_tx_url_bolt_s() {
608        let config = Neo4jConfig::new("bolt+s://host:7687", "neo4j", "pass", "idx");
609        let store = Neo4jVectorStore::new(config);
610        assert_eq!(
611            store.tx_url(),
612            "https://host:7687/db/neo4j/tx/commit"
613        );
614    }
615
616    #[test]
617    fn test_tx_url_custom_database() {
618        let config =
619            Neo4jConfig::new("bolt://localhost:7687", "neo4j", "pass", "idx").with_database("mydb");
620        let store = Neo4jVectorStore::new(config);
621        assert_eq!(
622            store.tx_url(),
623            "http://localhost:7687/db/mydb/tx/commit"
624        );
625    }
626
627    #[test]
628    fn test_base64_encode() {
629        // "neo4j:password" in base64
630        let encoded = base64_encode("neo4j:password".to_string());
631        assert_eq!(encoded, "bmVvNGo6cGFzc3dvcmQ=");
632    }
633
634    #[test]
635    fn test_base64_encode_empty() {
636        let encoded = base64_encode(String::new());
637        assert_eq!(encoded, "");
638    }
639
640    #[test]
641    fn test_store_new() {
642        let config = Neo4jConfig::new("bolt://localhost:7687", "neo4j", "pass", "idx");
643        let _store = Neo4jVectorStore::new(config);
644    }
645
646    #[test]
647    fn test_store_debug() {
648        let config = Neo4jConfig::new("bolt://localhost:7687", "neo4j", "pass", "idx");
649        let store = Neo4jVectorStore::new(config);
650        let debug_str = format!("{:?}", store);
651        assert!(debug_str.contains("Neo4jVectorStore"));
652        assert!(debug_str.contains("idx"));
653    }
654}
lc_vector_stores/neo4j.rs

lc_vector_stores/
neo4j.rs
