1use std::sync::Arc;
2
3use mem7_config::MemoryEngineConfig;
4use mem7_core::GraphRelation;
5use mem7_embedding::EmbeddingClient;
6use mem7_error::{Mem7Error, Result};
7use mem7_history::{HistoryStore, SqliteHistory};
8use mem7_llm::LlmClient;
9use mem7_reranker::RerankerClient;
10use mem7_vector::VectorIndex;
11use tracing::info;
12
13use crate::graph_pipeline::GraphPipeline;
14
15pub(crate) fn graph_result_to_relation(r: &mem7_graph::GraphSearchResult) -> GraphRelation {
16 GraphRelation {
17 source: r.source.clone(),
18 relationship: r.relationship.clone(),
19 destination: r.destination.clone(),
20 score: r.score,
21 }
22}
23
24pub struct MemoryEngine {
26 pub(crate) llm: Arc<dyn LlmClient>,
27 pub(crate) embedder: Arc<dyn EmbeddingClient>,
28 pub(crate) vector_index: Arc<dyn VectorIndex>,
29 pub(crate) history: Arc<dyn HistoryStore>,
30 pub(crate) reranker: Option<Arc<dyn RerankerClient>>,
31 pub(crate) graph_pipeline: Option<GraphPipeline>,
32 pub(crate) config: MemoryEngineConfig,
33}
34
35impl MemoryEngine {
36 pub async fn new(config: MemoryEngineConfig) -> Result<Self> {
37 let problems = config.validate();
38 if !problems.is_empty() {
39 return Err(Mem7Error::Config(problems.join("; ")));
40 }
41
42 let llm = mem7_llm::create_llm(&config.llm)?;
43 let embedder = mem7_embedding::create_embedding(&config.embedding)?;
44 let vector_index = mem7_vector::create_vector_index(&config.vector)?;
45 let history: Arc<dyn HistoryStore> =
46 Arc::new(SqliteHistory::new(&config.history.db_path).await?);
47 let reranker = config
48 .reranker
49 .as_ref()
50 .map(mem7_reranker::create_reranker)
51 .transpose()?;
52
53 let graph_pipeline = if let Some(graph_cfg) = &config.graph {
54 let store = mem7_graph::create_graph_store(graph_cfg).await?;
55 let g_llm = graph_cfg
56 .llm
57 .as_ref()
58 .map(mem7_llm::create_llm)
59 .transpose()?
60 .unwrap_or_else(|| llm.clone());
61 let custom_prompt = graph_cfg.custom_prompt.clone();
62 Some(GraphPipeline::new(
63 store,
64 g_llm,
65 embedder.clone(),
66 custom_prompt,
67 ))
68 } else {
69 None
70 };
71
72 info!(
73 reranker = reranker.is_some(),
74 graph = graph_pipeline.is_some(),
75 "MemoryEngine initialized"
76 );
77
78 Ok(Self {
79 llm,
80 embedder,
81 vector_index,
82 history,
83 reranker,
84 graph_pipeline,
85 config,
86 })
87 }
88}