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

lc_vector_stores/
neo4j.rs
