rexis_rag/agent/memory/
manager.rs

1//! Agent memory manager - coordinates all memory types
2
3use super::config::MemoryConfig;
4use super::conversation::{generate_session_id, ConversationMemoryStore};
5use super::episodic::EpisodicMemory;
6use super::semantic::SemanticMemory;
7use super::shared::SharedKnowledgeBase;
8use super::working::WorkingMemory;
9use crate::error::RragResult;
10use crate::storage::Memory;
11use rexis_llm::ChatMessage; // Use re-exported rsllm type
12use std::sync::Arc;
13
14/// Manages all memory types for an agent
15pub struct AgentMemoryManager {
16    /// Storage backend
17    storage: Arc<dyn Memory>,
18
19    /// Agent identifier
20    agent_id: String,
21
22    /// Current session identifier
23    session_id: String,
24
25    /// Conversation memory
26    conversation: ConversationMemoryStore,
27
28    /// Working memory (lazy-initialized)
29    working: Option<WorkingMemory>,
30
31    /// Semantic memory (lazy-initialized)
32    semantic: Option<SemanticMemory>,
33
34    /// Episodic memory (lazy-initialized)
35    episodic: Option<EpisodicMemory>,
36
37    /// Shared knowledge base (lazy-initialized)
38    shared: Option<SharedKnowledgeBase>,
39
40    /// Configuration
41    config: MemoryConfig,
42}
43
44impl AgentMemoryManager {
45    /// Create a new agent memory manager
46    pub fn new(mut config: MemoryConfig) -> Self {
47        // Auto-generate session ID if needed
48        if config.session_id.is_none() && config.auto_generate_session_id {
49            config.session_id = Some(generate_session_id());
50        }
51
52        let session_id = config
53            .session_id
54            .clone()
55            .unwrap_or_else(|| "default".to_string());
56
57        let conversation = ConversationMemoryStore::new(
58            config.backend.clone(),
59            session_id.clone(),
60            config.max_conversation_length,
61            config.persist_conversations,
62        );
63
64        Self {
65            storage: config.backend.clone(),
66            agent_id: config.agent_id.clone(),
67            session_id,
68            conversation,
69            working: None,
70            semantic: None,
71            episodic: None,
72            shared: None,
73            config,
74        }
75    }
76
77    /// Get or initialize working memory
78    pub fn working(&mut self) -> &mut WorkingMemory {
79        if self.working.is_none() {
80            self.working = Some(WorkingMemory::new(
81                self.storage.clone(),
82                self.session_id.clone(),
83            ));
84        }
85        self.working.as_mut().unwrap()
86    }
87
88    /// Get or initialize semantic memory
89    pub fn semantic(&mut self) -> &mut SemanticMemory {
90        if self.semantic.is_none() {
91            self.semantic = Some(SemanticMemory::new(
92                self.storage.clone(),
93                self.agent_id.clone(),
94            ));
95        }
96        self.semantic.as_mut().unwrap()
97    }
98
99    /// Get or initialize episodic memory
100    pub fn episodic(&mut self) -> &mut EpisodicMemory {
101        if self.episodic.is_none() {
102            self.episodic = Some(EpisodicMemory::new(
103                self.storage.clone(),
104                self.agent_id.clone(),
105            ));
106        }
107        self.episodic.as_mut().unwrap()
108    }
109
110    /// Get or initialize shared knowledge base
111    pub fn shared(&mut self) -> &mut SharedKnowledgeBase {
112        if self.shared.is_none() {
113            self.shared = Some(SharedKnowledgeBase::new(
114                self.storage.clone(),
115                self.agent_id.clone(),
116            ));
117        }
118        self.shared.as_mut().unwrap()
119    }
120
121    /// Get agent ID
122    pub fn agent_id(&self) -> &str {
123        &self.agent_id
124    }
125
126    /// Get session ID
127    pub fn session_id(&self) -> &str {
128        &self.session_id
129    }
130
131    /// Get conversation memory
132    pub fn conversation(&self) -> &ConversationMemoryStore {
133        &self.conversation
134    }
135
136    /// Add a message to conversation history
137    pub async fn add_conversation_message(&self, message: ChatMessage) -> RragResult<()> {
138        self.conversation.add_message(message).await
139    }
140
141    /// Get all conversation messages
142    pub async fn get_conversation_messages(&self) -> RragResult<Vec<ChatMessage>> {
143        self.conversation.get_messages().await
144    }
145
146    /// Clear conversation (keeps system message)
147    pub async fn clear_conversation(&self) -> RragResult<()> {
148        self.conversation.clear().await
149    }
150
151    /// Get the underlying storage backend
152    pub fn storage(&self) -> Arc<dyn Memory> {
153        self.storage.clone()
154    }
155
156    /// Generate a namespace key for agent-scoped memory
157    pub fn agent_key(&self, key: &str) -> String {
158        format!("agent::{}::{}", self.agent_id, key)
159    }
160
161    /// Generate a namespace key for session-scoped memory
162    pub fn session_key(&self, key: &str) -> String {
163        format!("session::{}::{}", self.session_id, key)
164    }
165
166    /// Generate a namespace key for global memory
167    pub fn global_key(key: &str) -> String {
168        format!("global::{}", key)
169    }
170
171    /// Store a value in agent-scoped memory
172    pub async fn set_agent_memory(
173        &self,
174        key: &str,
175        value: impl Into<crate::storage::MemoryValue>,
176    ) -> RragResult<()> {
177        let full_key = self.agent_key(key);
178        self.storage.set(&full_key, value.into()).await
179    }
180
181    /// Get a value from agent-scoped memory
182    pub async fn get_agent_memory(
183        &self,
184        key: &str,
185    ) -> RragResult<Option<crate::storage::MemoryValue>> {
186        let full_key = self.agent_key(key);
187        self.storage.get(&full_key).await
188    }
189
190    /// Store a value in session-scoped memory
191    pub async fn set_session_memory(
192        &self,
193        key: &str,
194        value: impl Into<crate::storage::MemoryValue>,
195    ) -> RragResult<()> {
196        let full_key = self.session_key(key);
197        self.storage.set(&full_key, value.into()).await
198    }
199
200    /// Get a value from session-scoped memory
201    pub async fn get_session_memory(
202        &self,
203        key: &str,
204    ) -> RragResult<Option<crate::storage::MemoryValue>> {
205        let full_key = self.session_key(key);
206        self.storage.get(&full_key).await
207    }
208
209    /// Store a value in global memory
210    pub async fn set_global_memory(
211        &self,
212        key: &str,
213        value: impl Into<crate::storage::MemoryValue>,
214    ) -> RragResult<()> {
215        let full_key = Self::global_key(key);
216        self.storage.set(&full_key, value.into()).await
217    }
218
219    /// Get a value from global memory
220    pub async fn get_global_memory(
221        &self,
222        key: &str,
223    ) -> RragResult<Option<crate::storage::MemoryValue>> {
224        let full_key = Self::global_key(key);
225        self.storage.get(&full_key).await
226    }
227
228    /// Get configuration
229    pub fn config(&self) -> &MemoryConfig {
230        &self.config
231    }
232}
233
234impl Clone for AgentMemoryManager {
235    fn clone(&self) -> Self {
236        Self {
237            storage: self.storage.clone(),
238            agent_id: self.agent_id.clone(),
239            session_id: self.session_id.clone(),
240            conversation: ConversationMemoryStore::new(
241                self.storage.clone(),
242                self.session_id.clone(),
243                self.config.max_conversation_length,
244                self.config.persist_conversations,
245            ),
246            working: None, // Lazy-initialize on clone
247            semantic: None,
248            episodic: None,
249            shared: None,
250            config: self.config.clone(),
251        }
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258    use crate::storage::{InMemoryStorage, MemoryValue};
259
260    #[tokio::test]
261    async fn test_memory_manager_namespacing() {
262        let storage = Arc::new(InMemoryStorage::new());
263        let config = MemoryConfig::new(storage, "test-agent").with_session_id("test-session");
264
265        let manager = AgentMemoryManager::new(config);
266
267        // Test namespacing
268        assert_eq!(
269            manager.agent_key("preferences"),
270            "agent::test-agent::preferences"
271        );
272        assert_eq!(
273            manager.session_key("working_data"),
274            "session::test-session::working_data"
275        );
276        assert_eq!(
277            AgentMemoryManager::global_key("shared_config"),
278            "global::shared_config"
279        );
280    }
281
282    #[tokio::test]
283    async fn test_memory_manager_scoped_storage() {
284        let storage = Arc::new(InMemoryStorage::new());
285        let config = MemoryConfig::new(storage.clone(), "test-agent");
286
287        let manager = AgentMemoryManager::new(config);
288
289        // Store in different scopes
290        manager
291            .set_agent_memory("profile::name", MemoryValue::from("Alice"))
292            .await
293            .unwrap();
294
295        manager
296            .set_session_memory("temp::data", MemoryValue::from(42i64))
297            .await
298            .unwrap();
299
300        manager
301            .set_global_memory("config::setting", MemoryValue::from(true))
302            .await
303            .unwrap();
304
305        // Retrieve from different scopes
306        let name = manager.get_agent_memory("profile::name").await.unwrap();
307        assert_eq!(name.unwrap().as_string(), Some("Alice"));
308
309        let data = manager.get_session_memory("temp::data").await.unwrap();
310        assert_eq!(data.unwrap().as_integer(), Some(42));
311
312        let setting = manager.get_global_memory("config::setting").await.unwrap();
313        assert_eq!(setting.unwrap().as_boolean(), Some(true));
314    }
315
316    #[tokio::test]
317    async fn test_conversation_integration() {
318        let storage = Arc::new(InMemoryStorage::new());
319        let config = MemoryConfig::new(storage, "test-agent").with_persistence(true);
320
321        let manager = AgentMemoryManager::new(config);
322
323        // Add messages
324        manager
325            .add_conversation_message(ChatMessage::system("System prompt"))
326            .await
327            .unwrap();
328
329        manager
330            .add_conversation_message(ChatMessage::user("Hello"))
331            .await
332            .unwrap();
333
334        // Get messages
335        let messages = manager.get_conversation_messages().await.unwrap();
336        assert_eq!(messages.len(), 2);
337    }
338}
rexis_rag/agent/memory/manager.rs

rexis_rag/agent/memory/
manager.rs
