1pub mod bm25;
2pub mod extraction;
3mod flat;
4#[cfg(feature = "kuzu")]
5mod kuzu_store;
6mod neo4j;
7pub mod prompts;
8pub mod types;
9
10pub use flat::FlatGraph;
11#[cfg(feature = "kuzu")]
12pub use kuzu_store::KuzuGraphStore;
13pub use neo4j::Neo4jGraphStore;
14pub use types::*;
15
16use std::sync::Arc;
17
18use async_trait::async_trait;
19use mem7_config::GraphConfig;
20use mem7_core::MemoryFilter;
21use mem7_error::{Mem7Error, Result};
22
23#[async_trait]
25pub trait GraphStore: Send + Sync {
26 async fn add_entities(&self, entities: &[Entity], filter: &MemoryFilter) -> Result<()>;
28
29 async fn add_relations(
31 &self,
32 relations: &[Relation],
33 entities: &[Entity],
34 filter: &MemoryFilter,
35 ) -> Result<()>;
36
37 async fn search(
39 &self,
40 query: &str,
41 filter: &MemoryFilter,
42 limit: usize,
43 ) -> Result<Vec<GraphSearchResult>>;
44
45 async fn search_by_embedding(
48 &self,
49 embedding: &[f32],
50 filter: &MemoryFilter,
51 threshold: f32,
52 limit: usize,
53 ) -> Result<Vec<GraphSearchResult>>;
54
55 async fn invalidate_relations(
58 &self,
59 triples: &[(String, String, String)],
60 filter: &MemoryFilter,
61 ) -> Result<()>;
62
63 async fn rehearse_relations(
66 &self,
67 triples: &[(String, String, String)],
68 filter: &MemoryFilter,
69 now: &str,
70 ) -> Result<()>;
71
72 async fn delete_all(&self, filter: &MemoryFilter) -> Result<()>;
74
75 async fn reset(&self) -> Result<()>;
77}
78
79pub async fn create_graph_store(config: &GraphConfig) -> Result<Arc<dyn GraphStore>> {
81 match config.provider.as_str() {
82 "flat" => Ok(Arc::new(FlatGraph::new())),
83 #[cfg(feature = "kuzu")]
84 "kuzu" => {
85 let path = config.kuzu_db_path.as_deref().unwrap_or("mem7_graph.kuzu");
86 Ok(Arc::new(KuzuGraphStore::new(path)?))
87 }
88 #[cfg(not(feature = "kuzu"))]
89 "kuzu" => Err(Mem7Error::Config(
90 "kuzu provider requires the `kuzu` feature: cargo add mem7-graph --features kuzu"
91 .into(),
92 )),
93 "neo4j" => {
94 let url = config
95 .neo4j_url
96 .as_deref()
97 .ok_or_else(|| Mem7Error::Config("neo4j_url is required".into()))?;
98 let username = config
99 .neo4j_username
100 .as_deref()
101 .ok_or_else(|| Mem7Error::Config("neo4j_username is required".into()))?;
102 let password = config
103 .neo4j_password
104 .as_deref()
105 .ok_or_else(|| Mem7Error::Config("neo4j_password is required".into()))?;
106 let database = config.neo4j_database.as_deref();
107 Ok(Arc::new(
108 Neo4jGraphStore::new(url, username, password, database).await?,
109 ))
110 }
111 other => Err(Mem7Error::Config(format!(
112 "unknown graph provider: {other}"
113 ))),
114 }
115}