ceylon_next/memory/
in_memory.rs1use super::{Memory, MemoryEntry};
7use async_trait::async_trait;
8use std::collections::HashMap;
9use std::sync::{Arc, RwLock};
10
11pub struct InMemoryStore {
26 entries: Arc<RwLock<HashMap<String, MemoryEntry>>>,
28
29 agent_index: Arc<RwLock<HashMap<String, Vec<String>>>>,
31}
32
33impl InMemoryStore {
34 pub fn new() -> Self {
44 Self {
45 entries: Arc::new(RwLock::new(HashMap::new())),
46 agent_index: Arc::new(RwLock::new(HashMap::new())),
47 }
48 }
49}
50
51#[async_trait]
52impl Memory for InMemoryStore {
53 async fn store(&self, entry: MemoryEntry) -> Result<String, String> {
54 let entry_id = entry.id.clone();
55 let agent_id = entry.agent_id.clone();
56
57 {
59 let mut entries = self.entries.write().unwrap();
60 entries.insert(entry_id.clone(), entry);
61 }
62
63 {
65 let mut index = self.agent_index.write().unwrap();
66 index
67 .entry(agent_id)
68 .or_insert_with(Vec::new)
69 .push(entry_id.clone());
70 }
71
72 Ok(entry_id)
73 }
74
75 async fn get(&self, id: &str) -> Result<Option<MemoryEntry>, String> {
76 let entries = self.entries.read().unwrap();
77 Ok(entries.get(id).cloned())
78 }
79
80 async fn get_agent_history(&self, agent_id: &str) -> Result<Vec<MemoryEntry>, String> {
81 let index = self.agent_index.read().unwrap();
82 let entries = self.entries.read().unwrap();
83
84 let entry_ids = index.get(agent_id);
86
87 if entry_ids.is_none() {
88 return Ok(Vec::new());
89 }
90
91 let mut result = Vec::new();
93 for entry_id in entry_ids.unwrap() {
94 if let Some(entry) = entries.get(entry_id) {
95 result.push(entry.clone());
96 }
97 }
98
99 result.sort_by(|a, b| b.created_at.cmp(&a.created_at));
101
102 Ok(result)
103 }
104
105 async fn get_recent(&self, agent_id: &str, limit: usize) -> Result<Vec<MemoryEntry>, String> {
106 let mut history = self.get_agent_history(agent_id).await?;
107
108 history.truncate(limit);
110
111 Ok(history)
112 }
113
114 async fn search(&self, agent_id: &str, query: &str) -> Result<Vec<MemoryEntry>, String> {
115 let history = self.get_agent_history(agent_id).await?;
116 let query_lower = query.to_lowercase();
117
118 let results: Vec<MemoryEntry> = history
120 .into_iter()
121 .filter(|entry| {
122 entry
123 .messages
124 .iter()
125 .any(|msg| msg.content.to_lowercase().contains(&query_lower))
126 })
127 .collect();
128
129 Ok(results)
130 }
131
132 async fn clear_agent_memory(&self, agent_id: &str) -> Result<(), String> {
133 let index = self.agent_index.read().unwrap();
134 let entry_ids = index.get(agent_id);
135
136 if let Some(ids) = entry_ids {
137 let mut entries = self.entries.write().unwrap();
138 for id in ids {
139 entries.remove(id);
140 }
141 }
142
143 let mut index = self.agent_index.write().unwrap();
144 index.remove(agent_id);
145
146 Ok(())
147 }
148}
149
150impl Default for InMemoryStore {
151 fn default() -> Self {
152 Self::new()
153 }
154}