Skip to main content

mem7_store/
builder.rs

1use std::sync::Arc;
2
3use mem7_config::MemoryEngineConfig;
4use mem7_embedding::EmbeddingClient;
5use mem7_error::{Mem7Error, Result};
6use mem7_history::{HistoryStore, SqliteHistory};
7use mem7_llm::LlmClient;
8use mem7_reranker::RerankerClient;
9use mem7_vector::VectorIndex;
10use tracing::info;
11
12use crate::engine::MemoryEngine;
13use crate::graph_pipeline::GraphPipeline;
14
15/// Builder for `MemoryEngine` with dependency injection support.
16///
17/// Use this when you need to supply custom or mock implementations for
18/// LLM, embedding, vector, history, reranker, or graph components.
19///
20/// ```rust,ignore
21/// let engine = MemoryEngineBuilder::new(config)
22///     .llm(my_llm)
23///     .history(my_history)
24///     .build()
25///     .await?;
26/// ```
27pub struct MemoryEngineBuilder {
28    config: MemoryEngineConfig,
29    llm: Option<Arc<dyn LlmClient>>,
30    embedder: Option<Arc<dyn EmbeddingClient>>,
31    vector_index: Option<Arc<dyn VectorIndex>>,
32    history: Option<Arc<dyn HistoryStore>>,
33    reranker: Option<Arc<dyn RerankerClient>>,
34    graph_pipeline: Option<GraphPipeline>,
35}
36
37impl MemoryEngineBuilder {
38    pub fn new(config: MemoryEngineConfig) -> Self {
39        Self {
40            config,
41            llm: None,
42            embedder: None,
43            vector_index: None,
44            history: None,
45            reranker: None,
46            graph_pipeline: None,
47        }
48    }
49
50    pub fn llm(mut self, llm: Arc<dyn LlmClient>) -> Self {
51        self.llm = Some(llm);
52        self
53    }
54
55    pub fn embedder(mut self, embedder: Arc<dyn EmbeddingClient>) -> Self {
56        self.embedder = Some(embedder);
57        self
58    }
59
60    pub fn vector_index(mut self, vi: Arc<dyn VectorIndex>) -> Self {
61        self.vector_index = Some(vi);
62        self
63    }
64
65    pub fn history(mut self, history: Arc<dyn HistoryStore>) -> Self {
66        self.history = Some(history);
67        self
68    }
69
70    pub fn reranker(mut self, reranker: Arc<dyn RerankerClient>) -> Self {
71        self.reranker = Some(reranker);
72        self
73    }
74
75    /// Build the `MemoryEngine`, creating default implementations for any
76    /// components not explicitly provided.
77    pub async fn build(self) -> Result<MemoryEngine> {
78        let problems = self.config.validate();
79        if !problems.is_empty() {
80            return Err(Mem7Error::Config(problems.join("; ")));
81        }
82
83        let llm = match self.llm {
84            Some(l) => l,
85            None => mem7_llm::create_llm(&self.config.llm)?,
86        };
87        let embedder = match self.embedder {
88            Some(e) => e,
89            None => mem7_embedding::create_embedding(&self.config.embedding)?,
90        };
91        let vector_index = match self.vector_index {
92            Some(v) => v,
93            None => mem7_vector::create_vector_index(&self.config.vector)?,
94        };
95        let history: Arc<dyn HistoryStore> = match self.history {
96            Some(h) => h,
97            None => Arc::new(SqliteHistory::new(&self.config.history.db_path).await?),
98        };
99        let reranker = match self.reranker {
100            Some(r) => Some(r),
101            None => self
102                .config
103                .reranker
104                .as_ref()
105                .map(mem7_reranker::create_reranker)
106                .transpose()?,
107        };
108
109        let graph_pipeline = match self.graph_pipeline {
110            Some(gp) => Some(gp),
111            None => {
112                if let Some(graph_cfg) = &self.config.graph {
113                    let store = mem7_graph::create_graph_store(graph_cfg).await?;
114                    let g_llm = graph_cfg
115                        .llm
116                        .as_ref()
117                        .map(mem7_llm::create_llm)
118                        .transpose()?
119                        .unwrap_or_else(|| llm.clone());
120                    let custom_prompt = graph_cfg.custom_prompt.clone();
121                    Some(GraphPipeline::new(
122                        store,
123                        g_llm,
124                        embedder.clone(),
125                        custom_prompt,
126                    ))
127                } else {
128                    None
129                }
130            }
131        };
132
133        info!(
134            reranker = reranker.is_some(),
135            graph = graph_pipeline.is_some(),
136            "MemoryEngine initialized (builder)"
137        );
138
139        Ok(MemoryEngine {
140            llm,
141            embedder,
142            vector_index,
143            history,
144            reranker,
145            graph_pipeline,
146            config: self.config,
147        })
148    }
149}