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
15pub 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 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}