ceylon_next/memory/advanced/
episodic.rs1use super::{EnhancedMemoryEntry, ImportanceLevel, MemoryConfig, MemoryType};
14use crate::memory::{Memory, MemoryEntry};
15use std::collections::HashMap;
16use std::sync::Arc;
17use tokio::sync::RwLock;
18
19#[derive(Debug, Clone)]
21pub enum TimeRange {
22 LastHours(u64),
24 LastDays(u64),
26 LastWeeks(u64),
28 Between(u64, u64),
30 All,
32}
33
34pub struct EpisodicMemory {
36 backend: Arc<dyn Memory>,
38 index: Arc<RwLock<HashMap<String, EnhancedMemoryEntry>>>,
40 config: MemoryConfig,
42}
43
44impl EpisodicMemory {
45 pub async fn new(backend: Arc<dyn Memory>, config: MemoryConfig) -> Result<Self, String> {
47 Ok(Self {
48 backend,
49 index: Arc::new(RwLock::new(HashMap::new())),
50 config,
51 })
52 }
53
54 pub async fn store(&self, mut entry: EnhancedMemoryEntry) -> Result<String, String> {
56 entry.memory_type = MemoryType::Episodic;
57
58 let id = self.backend.store(entry.entry.clone()).await?;
60
61 let mut index = self.index.write().await;
63 index.insert(id.clone(), entry);
64
65 Ok(id)
66 }
67
68 pub async fn get(&self, id: &str) -> Result<Option<EnhancedMemoryEntry>, String> {
70 let mut index = self.index.write().await;
71
72 if let Some(mut entry) = index.get_mut(id) {
73 entry.mark_accessed();
74 Ok(Some(entry.clone()))
75 } else {
76 if let Some(base_entry) = self.backend.get(id).await? {
78 let enhanced = EnhancedMemoryEntry::new(base_entry, MemoryType::Episodic);
79 index.insert(id.to_string(), enhanced.clone());
80 Ok(Some(enhanced))
81 } else {
82 Ok(None)
83 }
84 }
85 }
86
87 pub async fn get_by_time_range(
89 &self,
90 agent_id: &str,
91 range: TimeRange,
92 ) -> Result<Vec<EnhancedMemoryEntry>, String> {
93 let all_memories = self.load_agent_memories(agent_id).await?;
94
95 let current_time = Self::current_timestamp();
96 let filtered: Vec<EnhancedMemoryEntry> = all_memories
97 .into_iter()
98 .filter(|m| {
99 let created = m.entry.created_at;
100 match range {
101 TimeRange::LastHours(hours) => {
102 current_time.saturating_sub(created) <= hours * 3600
103 }
104 TimeRange::LastDays(days) => {
105 current_time.saturating_sub(created) <= days * 86400
106 }
107 TimeRange::LastWeeks(weeks) => {
108 current_time.saturating_sub(created) <= weeks * 604800
109 }
110 TimeRange::Between(start, end) => created >= start && created <= end,
111 TimeRange::All => true,
112 }
113 })
114 .collect();
115
116 Ok(filtered)
117 }
118
119 pub async fn get_by_relevance(
121 &self,
122 agent_id: &str,
123 min_score: f32,
124 limit: usize,
125 ) -> Result<Vec<EnhancedMemoryEntry>, String> {
126 let mut memories = self.load_agent_memories(agent_id).await?;
127
128 memories.sort_by(|a, b| {
130 b.relevance_score()
131 .partial_cmp(&a.relevance_score())
132 .unwrap_or(std::cmp::Ordering::Equal)
133 });
134
135 let result: Vec<EnhancedMemoryEntry> = memories
137 .into_iter()
138 .filter(|m| m.relevance_score() >= min_score)
139 .take(limit)
140 .collect();
141
142 Ok(result)
143 }
144
145 pub async fn get_by_importance(
147 &self,
148 agent_id: &str,
149 min_importance: ImportanceLevel,
150 ) -> Result<Vec<EnhancedMemoryEntry>, String> {
151 let memories = self.load_agent_memories(agent_id).await?;
152
153 let result: Vec<EnhancedMemoryEntry> = memories
154 .into_iter()
155 .filter(|m| m.importance >= min_importance)
156 .collect();
157
158 Ok(result)
159 }
160
161 pub async fn search(
163 &self,
164 agent_id: &str,
165 query: &str,
166 ) -> Result<Vec<EnhancedMemoryEntry>, String> {
167 let base_entries = self.backend.search(agent_id, query).await?;
169
170 let mut result = Vec::new();
171 for entry in base_entries {
172 if let Some(enhanced) = self.get(&entry.id).await? {
173 result.push(enhanced);
174 }
175 }
176
177 Ok(result)
178 }
179
180 pub async fn get_temporal_clusters(
182 &self,
183 agent_id: &str,
184 cluster_gap_hours: u64,
185 ) -> Result<Vec<Vec<EnhancedMemoryEntry>>, String> {
186 let mut memories = self.load_agent_memories(agent_id).await?;
187
188 memories.sort_by_key(|m| m.entry.created_at);
190
191 let mut clusters: Vec<Vec<EnhancedMemoryEntry>> = Vec::new();
192 let mut current_cluster: Vec<EnhancedMemoryEntry> = Vec::new();
193 let gap_seconds = cluster_gap_hours * 3600;
194
195 for memory in memories {
196 if let Some(last) = current_cluster.last() {
197 let time_diff = memory.entry.created_at.saturating_sub(last.entry.created_at);
198 if time_diff > gap_seconds {
199 clusters.push(current_cluster.clone());
201 current_cluster.clear();
202 }
203 }
204 current_cluster.push(memory);
205 }
206
207 if !current_cluster.is_empty() {
208 clusters.push(current_cluster);
209 }
210
211 Ok(clusters)
212 }
213
214 pub async fn get_activity_summary(
216 &self,
217 agent_id: &str,
218 days: u64,
219 ) -> Result<ActivitySummary, String> {
220 let memories = self
221 .get_by_time_range(agent_id, TimeRange::LastDays(days))
222 .await?;
223
224 let total_conversations = memories.len();
225 let total_messages: usize = memories.iter().map(|m| m.entry.messages.len()).sum();
226
227 let avg_importance = if !memories.is_empty() {
228 memories.iter().map(|m| m.importance as u32).sum::<u32>() as f32
229 / memories.len() as f32
230 } else {
231 0.0
232 };
233
234 let critical_count = memories
235 .iter()
236 .filter(|m| m.importance == ImportanceLevel::Critical)
237 .count();
238
239 Ok(ActivitySummary {
240 time_range_days: days,
241 total_conversations,
242 total_messages,
243 average_importance: avg_importance,
244 critical_conversations: critical_count,
245 })
246 }
247
248 async fn load_agent_memories(&self, agent_id: &str) -> Result<Vec<EnhancedMemoryEntry>, String> {
250 let base_entries = self.backend.get_agent_history(agent_id).await?;
251
252 let mut index = self.index.write().await;
253 let mut result = Vec::new();
254
255 for base_entry in base_entries {
256 let id = base_entry.id.clone();
257 if let Some(enhanced) = index.get(&id) {
258 result.push(enhanced.clone());
259 } else {
260 let enhanced = EnhancedMemoryEntry::new(base_entry, MemoryType::Episodic);
261 index.insert(id, enhanced.clone());
262 result.push(enhanced);
263 }
264 }
265
266 Ok(result)
267 }
268
269 pub async fn clear(&self, agent_id: &str) -> Result<(), String> {
271 self.backend.clear_agent_memory(agent_id).await?;
273
274 let mut index = self.index.write().await;
276 index.retain(|_, v| v.entry.agent_id != agent_id);
277
278 Ok(())
279 }
280
281 fn current_timestamp() -> u64 {
282 std::time::SystemTime::now()
283 .duration_since(std::time::UNIX_EPOCH)
284 .unwrap()
285 .as_secs()
286 }
287}
288
289#[derive(Debug, Clone)]
291pub struct ActivitySummary {
292 pub time_range_days: u64,
293 pub total_conversations: usize,
294 pub total_messages: usize,
295 pub average_importance: f32,
296 pub critical_conversations: usize,
297}
298
299#[cfg(test)]
300mod tests {
301 use super::*;
302 use crate::memory::InMemoryStore;
303
304 #[tokio::test]
305 async fn test_episodic_memory_time_range() {
306 let backend = Arc::new(InMemoryStore::new());
307 let config = MemoryConfig::default();
308 let em = EpisodicMemory::new(backend, config).await.unwrap();
309
310 for i in 0..5 {
312 let entry = MemoryEntry::new(
313 "agent-1".to_string(),
314 format!("task-{}", i),
315 vec![],
316 );
317 let enhanced = EnhancedMemoryEntry::new(entry, MemoryType::Episodic);
318 em.store(enhanced).await.unwrap();
319 }
320
321 let all = em
323 .get_by_time_range("agent-1", TimeRange::All)
324 .await
325 .unwrap();
326 assert_eq!(all.len(), 5);
327
328 let recent = em
330 .get_by_time_range("agent-1", TimeRange::LastDays(1))
331 .await
332 .unwrap();
333 assert_eq!(recent.len(), 5); }
335
336 #[tokio::test]
337 async fn test_episodic_memory_relevance() {
338 let backend = Arc::new(InMemoryStore::new());
339 let config = MemoryConfig::default();
340 let em = EpisodicMemory::new(backend, config).await.unwrap();
341
342 for i in 0..3 {
344 let entry = MemoryEntry::new(
345 "agent-1".to_string(),
346 format!("task-{}", i),
347 vec![],
348 );
349 let mut enhanced = EnhancedMemoryEntry::new(entry, MemoryType::Episodic);
350 enhanced.importance = if i == 0 {
351 ImportanceLevel::Critical
352 } else {
353 ImportanceLevel::Low
354 };
355 em.store(enhanced).await.unwrap();
356 }
357
358 let relevant = em.get_by_relevance("agent-1", 0.0, 10).await.unwrap();
360 assert_eq!(relevant.len(), 3);
361 assert_eq!(relevant[0].importance, ImportanceLevel::Critical);
363 }
364}