1pub mod cache;
27pub mod agentdb;
28pub mod reasoningbank;
29pub mod coordination;
30
31pub use cache::{HotCache, CacheConfig, CacheEntry};
33pub use agentdb::{VectorStore, EmbeddingProvider};
34pub use reasoningbank::{
35 TrajectoryTracker,
36 VerdictJudge,
37 MemoryDistiller,
38 Trajectory,
39 Verdict,
40};
41pub use coordination::{
42 PubSubBroker,
43 DistributedLock,
44 ConsensusEngine,
45 Namespace,
46};
47
48use thiserror::Error;
49
50#[derive(Debug, Error)]
52pub enum MemoryError {
53 #[error("Cache error: {0}")]
54 Cache(String),
55
56 #[error("Vector database error: {0}")]
57 VectorDB(String),
58
59 #[error("Storage error: {0}")]
60 Storage(String),
61
62 #[error("Coordination error: {0}")]
63 Coordination(String),
64
65 #[error("Serialization error: {0}")]
66 Serialization(String),
67
68 #[error("Lock acquisition timeout")]
69 LockTimeout,
70
71 #[error("Consensus not reached")]
72 ConsensusFailure,
73
74 #[error("Invalid namespace: {0}")]
75 InvalidNamespace(String),
76}
77
78pub type Result<T> = std::result::Result<T, MemoryError>;
79
80#[derive(Debug, Clone)]
82pub struct MemoryConfig {
83 pub cache_config: CacheConfig,
85
86 pub agentdb_url: String,
88
89 pub storage_path: String,
91
92 pub enable_compression: bool,
94
95 pub max_memory_bytes: usize,
97}
98
99impl Default for MemoryConfig {
100 fn default() -> Self {
101 Self {
102 cache_config: CacheConfig::default(),
103 agentdb_url: "http://localhost:3000".to_string(),
104 storage_path: "./data/memory".to_string(),
105 enable_compression: true,
106 max_memory_bytes: 1_073_741_824, }
108 }
109}
110
111pub struct MemorySystem {
113 cache: HotCache,
115
116 vector_store: VectorStore,
118
119 cold_storage: sled::Db,
121
122 trajectory_tracker: TrajectoryTracker,
124 verdict_judge: VerdictJudge,
125 distiller: MemoryDistiller,
126
127 pubsub: PubSubBroker,
129 locks: DistributedLock,
130 consensus: ConsensusEngine,
131
132 config: MemoryConfig,
134}
135
136impl MemorySystem {
137 pub async fn new(config: MemoryConfig) -> Result<Self> {
139 let cache = HotCache::new(config.cache_config.clone());
141
142 let vector_store = VectorStore::new(&config.agentdb_url)
144 .await
145 .map_err(|e| MemoryError::VectorDB(e.to_string()))?;
146
147 let cold_storage = sled::open(&config.storage_path)
149 .map_err(|e| MemoryError::Storage(e.to_string()))?;
150
151 let trajectory_tracker = TrajectoryTracker::new();
153 let verdict_judge = VerdictJudge::new();
154 let distiller = MemoryDistiller::new(config.enable_compression);
155
156 let pubsub = PubSubBroker::new();
158 let locks = DistributedLock::new();
159 let consensus = ConsensusEngine::new();
160
161 Ok(Self {
162 cache,
163 vector_store,
164 cold_storage,
165 trajectory_tracker,
166 verdict_judge,
167 distiller,
168 pubsub,
169 locks,
170 consensus,
171 config,
172 })
173 }
174
175 pub async fn get(&self, namespace: &str, key: &str) -> Result<Option<Vec<u8>>> {
177 let full_key = format!("{}/{}", namespace, key);
178
179 if let Some(entry) = self.cache.get(&full_key) {
181 tracing::debug!("L1 cache hit: {}", full_key);
182 return Ok(Some(entry.data));
183 }
184
185 if let Some(data) = self.cold_storage
187 .get(full_key.as_bytes())
188 .map_err(|e| MemoryError::Storage(e.to_string()))?
189 {
190 tracing::debug!("L3 storage hit: {}", full_key);
191
192 self.cache.insert(&full_key, data.to_vec());
194
195 return Ok(Some(data.to_vec()));
196 }
197
198 tracing::debug!("Cache miss: {}", full_key);
199 Ok(None)
200 }
201
202 pub async fn put(&self, namespace: &str, key: &str, value: Vec<u8>) -> Result<()> {
204 let full_key = format!("{}/{}", namespace, key);
205
206 self.cache.insert(&full_key, value.clone());
208
209 self.cold_storage
211 .insert(full_key.as_bytes(), value.as_slice())
212 .map_err(|e| MemoryError::Storage(e.to_string()))?;
213
214 tracing::debug!("Stored: {}", full_key);
215 Ok(())
216 }
217
218 pub async fn search_similar(
220 &self,
221 namespace: &str,
222 query_embedding: Vec<f32>,
223 top_k: usize,
224 ) -> Result<Vec<(String, f32)>> {
225 self.vector_store
226 .search(namespace, query_embedding, top_k)
227 .await
228 .map_err(|e| MemoryError::VectorDB(e.to_string()))
229 }
230
231 pub async fn track_trajectory(&self, trajectory: Trajectory) -> Result<()> {
233 self.trajectory_tracker
234 .track(trajectory)
235 .await
236 .map_err(|e| MemoryError::Storage(e.to_string()))
237 }
238
239 pub async fn subscribe(&self, topic: &str) -> Result<tokio::sync::mpsc::Receiver<Vec<u8>>> {
241 self.pubsub
242 .subscribe(topic)
243 .await
244 .map_err(|e| MemoryError::Coordination(e.to_string()))
245 }
246
247 pub async fn publish(&self, topic: &str, message: Vec<u8>) -> Result<()> {
249 self.pubsub
250 .publish(topic, message)
251 .await
252 .map_err(|e| MemoryError::Coordination(e.to_string()))
253 }
254
255 pub async fn acquire_lock(
257 &self,
258 resource: &str,
259 timeout: std::time::Duration,
260 ) -> Result<String> {
261 self.locks
262 .acquire(resource, timeout)
263 .await
264 .map_err(|e| MemoryError::Coordination(e.to_string()))
265 }
266
267 pub async fn release_lock(&self, token: &str) -> Result<()> {
269 self.locks
270 .release(token)
271 .await
272 .map_err(|e| MemoryError::Coordination(e.to_string()))
273 }
274
275 pub fn stats(&self) -> MemoryStats {
277 MemoryStats {
278 l1_entries: self.cache.len(),
279 l1_hit_rate: self.cache.hit_rate(),
280 l3_size_bytes: self.cold_storage.size_on_disk().unwrap_or(0),
281 total_trajectories: self.trajectory_tracker.count(),
282 }
283 }
284}
285
286#[derive(Debug, Clone)]
288pub struct MemoryStats {
289 pub l1_entries: usize,
290 pub l1_hit_rate: f64,
291 pub l3_size_bytes: u64,
292 pub total_trajectories: usize,
293}
294
295#[cfg(test)]
296mod tests {
297 use super::*;
298
299 #[tokio::test]
300 async fn test_memory_system_creation() {
301 let mem_config = MemoryConfig {
302 storage_path: tempfile::tempdir().unwrap().path().to_str().unwrap().to_string(),
303 ..Default::default()
304 };
305
306 let memory = MemorySystem::new(mem_config).await;
307 assert!(memory.is_ok());
308 }
309}