mem7-graph 0.3.3

Graph store trait and providers for mem7
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

pub mod bm25;
pub mod extraction;
mod flat;
#[cfg(feature = "kuzu")]
mod kuzu_store;
mod neo4j;
pub mod prompts;
pub mod types;

pub use flat::FlatGraph;
#[cfg(feature = "kuzu")]
pub use kuzu_store::KuzuGraphStore;
pub use neo4j::Neo4jGraphStore;
pub use types::*;

use std::sync::Arc;

use async_trait::async_trait;
use mem7_config::GraphConfig;
use mem7_core::MemoryFilter;
use mem7_error::{Mem7Error, Result};

/// Trait for graph storage backends.
#[async_trait]
pub trait GraphStore: Send + Sync {
    /// Store extracted entities (provider may use this for node creation).
    async fn add_entities(&self, entities: &[Entity], filter: &MemoryFilter) -> Result<()>;

    /// Store extracted relations along with their entity metadata.
    async fn add_relations(
        &self,
        relations: &[Relation],
        entities: &[Entity],
        filter: &MemoryFilter,
    ) -> Result<()>;

    /// Text-based search for relations (legacy fallback).
    async fn search(
        &self,
        query: &str,
        filter: &MemoryFilter,
        limit: usize,
    ) -> Result<Vec<GraphSearchResult>>;

    /// Semantic search: find entities by embedding cosine similarity, then
    /// return all valid relations touching matched entities (1-hop traversal).
    async fn search_by_embedding(
        &self,
        embedding: &[f32],
        filter: &MemoryFilter,
        threshold: f32,
        limit: usize,
    ) -> Result<Vec<GraphSearchResult>>;

    /// Soft-delete relations by marking them as `valid = false`.
    /// Each tuple is `(source, relationship, destination)`.
    async fn invalidate_relations(
        &self,
        triples: &[(String, String, String)],
        filter: &MemoryFilter,
    ) -> Result<()>;

    /// Strengthen relations by updating `last_accessed_at` and incrementing
    /// `mentions`. Called asynchronously after successful search retrieval.
    async fn rehearse_relations(
        &self,
        triples: &[(String, String, String)],
        filter: &MemoryFilter,
        now: &str,
    ) -> Result<()>;

    /// Delete all relations matching the filter.
    async fn delete_all(&self, filter: &MemoryFilter) -> Result<()>;

    /// Remove all data from the graph store.
    async fn reset(&self) -> Result<()>;
}

/// Create a graph store from configuration.
pub async fn create_graph_store(config: &GraphConfig) -> Result<Arc<dyn GraphStore>> {
    match config.provider.as_str() {
        "flat" => Ok(Arc::new(FlatGraph::new())),
        #[cfg(feature = "kuzu")]
        "kuzu" => {
            let path = config.kuzu_db_path.as_deref().unwrap_or("mem7_graph.kuzu");
            Ok(Arc::new(KuzuGraphStore::new(path)?))
        }
        #[cfg(not(feature = "kuzu"))]
        "kuzu" => Err(Mem7Error::Config(
            "kuzu provider requires the `kuzu` feature: cargo add mem7-graph --features kuzu"
                .into(),
        )),
        "neo4j" => {
            let url = config
                .neo4j_url
                .as_deref()
                .ok_or_else(|| Mem7Error::Config("neo4j_url is required".into()))?;
            let username = config
                .neo4j_username
                .as_deref()
                .ok_or_else(|| Mem7Error::Config("neo4j_username is required".into()))?;
            let password = config
                .neo4j_password
                .as_deref()
                .ok_or_else(|| Mem7Error::Config("neo4j_password is required".into()))?;
            let database = config.neo4j_database.as_deref();
            Ok(Arc::new(
                Neo4jGraphStore::new(url, username, password, database).await?,
            ))
        }
        other => Err(Mem7Error::Config(format!(
            "unknown graph provider: {other}"
        ))),
    }
}