Skip to main content

mem7_graph/
lib.rs

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/// Trait for graph storage backends.
24#[async_trait]
25pub trait GraphStore: Send + Sync {
26    /// Store extracted entities (provider may use this for node creation).
27    async fn add_entities(&self, entities: &[Entity], filter: &MemoryFilter) -> Result<()>;
28
29    /// Store extracted relations along with their entity metadata.
30    async fn add_relations(
31        &self,
32        relations: &[Relation],
33        entities: &[Entity],
34        filter: &MemoryFilter,
35    ) -> Result<()>;
36
37    /// Text-based search for relations (legacy fallback).
38    async fn search(
39        &self,
40        query: &str,
41        filter: &MemoryFilter,
42        limit: usize,
43    ) -> Result<Vec<GraphSearchResult>>;
44
45    /// Semantic search: find entities by embedding cosine similarity, then
46    /// return all valid relations touching matched entities (1-hop traversal).
47    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    /// Soft-delete relations by marking them as `valid = false`.
56    /// Each tuple is `(source, relationship, destination)`.
57    async fn invalidate_relations(
58        &self,
59        triples: &[(String, String, String)],
60        filter: &MemoryFilter,
61    ) -> Result<()>;
62
63    /// Strengthen relations by updating `last_accessed_at` and incrementing
64    /// `mentions`. Called asynchronously after successful search retrieval.
65    async fn rehearse_relations(
66        &self,
67        triples: &[(String, String, String)],
68        filter: &MemoryFilter,
69        now: &str,
70    ) -> Result<()>;
71
72    /// Delete all relations matching the filter.
73    async fn delete_all(&self, filter: &MemoryFilter) -> Result<()>;
74
75    /// Remove all data from the graph store.
76    async fn reset(&self) -> Result<()>;
77}
78
79/// Create a graph store from configuration.
80pub 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}