zeroclaw_runtime/agent/
memory_loader.rs

1use async_trait::async_trait;
2use std::fmt::Write;
3use zeroclaw_memory::{self, MEMORY_CONTEXT_CLOSE, MEMORY_CONTEXT_OPEN, Memory, decay};
4
5#[async_trait]
6pub trait MemoryLoader: Send + Sync {
7    async fn load_context(
8        &self,
9        memory: &dyn Memory,
10        user_message: &str,
11        session_id: Option<&str>,
12    ) -> anyhow::Result<String>;
13}
14
15pub struct DefaultMemoryLoader {
16    limit: usize,
17    min_relevance_score: f64,
18}
19
20impl Default for DefaultMemoryLoader {
21    fn default() -> Self {
22        Self {
23            limit: 5,
24            min_relevance_score: 0.4,
25        }
26    }
27}
28
29impl DefaultMemoryLoader {
30    pub fn new(limit: usize, min_relevance_score: f64) -> Self {
31        Self {
32            limit: limit.max(1),
33            min_relevance_score,
34        }
35    }
36}
37
38#[async_trait]
39impl MemoryLoader for DefaultMemoryLoader {
40    async fn load_context(
41        &self,
42        memory: &dyn Memory,
43        user_message: &str,
44        session_id: Option<&str>,
45    ) -> anyhow::Result<String> {
46        let mut entries = memory
47            .recall(user_message, self.limit, session_id, None, None)
48            .await?;
49        if entries.is_empty() {
50            return Ok(String::new());
51        }
52
53        // Apply time decay: older non-Core memories score lower
54        decay::apply_time_decay(&mut entries, decay::DEFAULT_HALF_LIFE_DAYS);
55
56        let mut context = String::new();
57        let mut included = false;
58        for entry in entries {
59            if zeroclaw_memory::is_assistant_autosave_key(&entry.key) {
60                continue;
61            }
62            if zeroclaw_memory::is_user_autosave_key(&entry.key) {
63                continue;
64            }
65            if zeroclaw_memory::should_skip_autosave_content(&entry.content) {
66                continue;
67            }
68            if let Some(score) = entry.score
69                && score < self.min_relevance_score
70            {
71                continue;
72            }
73            if !included {
74                context.push_str(MEMORY_CONTEXT_OPEN);
75                context.push('\n');
76                included = true;
77            }
78            let _ = writeln!(context, "- {}: {}", entry.key, entry.content);
79        }
80
81        // If all entries were below threshold, return empty
82        if !included {
83            return Ok(String::new());
84        }
85
86        context.push_str(MEMORY_CONTEXT_CLOSE);
87        context.push_str("\n\n");
88        Ok(context)
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use std::sync::Arc;
96    use zeroclaw_memory::{
97        MEMORY_CONTEXT_CLOSE, MEMORY_CONTEXT_OPEN, Memory, MemoryCategory, MemoryEntry,
98    };
99
100    struct MockMemory;
101    struct MockMemoryWithEntries {
102        entries: Arc<Vec<MemoryEntry>>,
103    }
104
105    #[async_trait]
106    impl Memory for MockMemory {
107        async fn store(
108            &self,
109            _key: &str,
110            _content: &str,
111            _category: MemoryCategory,
112            _session_id: Option<&str>,
113        ) -> anyhow::Result<()> {
114            Ok(())
115        }
116
117        async fn recall(
118            &self,
119            _query: &str,
120            limit: usize,
121            _session_id: Option<&str>,
122            _since: Option<&str>,
123            _until: Option<&str>,
124        ) -> anyhow::Result<Vec<MemoryEntry>> {
125            if limit == 0 {
126                return Ok(vec![]);
127            }
128            Ok(vec![MemoryEntry {
129                id: "1".into(),
130                key: "k".into(),
131                content: "v".into(),
132                category: MemoryCategory::Conversation,
133                timestamp: "now".into(),
134                session_id: None,
135                score: None,
136                namespace: "default".into(),
137                importance: None,
138                superseded_by: None,
139                agent_alias: None,
140                agent_id: None,
141            }])
142        }
143
144        async fn get(&self, _key: &str) -> anyhow::Result<Option<MemoryEntry>> {
145            Ok(None)
146        }
147
148        async fn list(
149            &self,
150            _category: Option<&MemoryCategory>,
151            _session_id: Option<&str>,
152        ) -> anyhow::Result<Vec<MemoryEntry>> {
153            Ok(vec![])
154        }
155
156        async fn forget(&self, _key: &str) -> anyhow::Result<bool> {
157            Ok(true)
158        }
159
160        async fn forget_for_agent(&self, _key: &str, _agent_id: &str) -> anyhow::Result<bool> {
161            Ok(true)
162        }
163
164        async fn count(&self) -> anyhow::Result<usize> {
165            Ok(0)
166        }
167
168        async fn health_check(&self) -> bool {
169            true
170        }
171
172        fn name(&self) -> &str {
173            "mock"
174        }
175
176        async fn store_with_agent(
177            &self,
178            _key: &str,
179            _content: &str,
180            _category: MemoryCategory,
181            _session_id: Option<&str>,
182            _namespace: Option<&str>,
183            _importance: Option<f64>,
184            _agent_id: Option<&str>,
185        ) -> anyhow::Result<()> {
186            Ok(())
187        }
188
189        async fn recall_for_agents(
190            &self,
191            _allowed_agent_ids: &[&str],
192            query: &str,
193            limit: usize,
194            session_id: Option<&str>,
195            since: Option<&str>,
196            until: Option<&str>,
197        ) -> anyhow::Result<Vec<MemoryEntry>> {
198            self.recall(query, limit, session_id, since, until).await
199        }
200    }
201    impl ::zeroclaw_api::attribution::Attributable for MockMemory {
202        fn role(&self) -> ::zeroclaw_api::attribution::Role {
203            ::zeroclaw_api::attribution::Role::Memory(
204                ::zeroclaw_api::attribution::MemoryKind::InMemory,
205            )
206        }
207        fn alias(&self) -> &str {
208            "MockMemory"
209        }
210    }
211
212    #[async_trait]
213    impl Memory for MockMemoryWithEntries {
214        async fn store(
215            &self,
216            _key: &str,
217            _content: &str,
218            _category: MemoryCategory,
219            _session_id: Option<&str>,
220        ) -> anyhow::Result<()> {
221            Ok(())
222        }
223
224        async fn recall(
225            &self,
226            _query: &str,
227            _limit: usize,
228            _session_id: Option<&str>,
229            _since: Option<&str>,
230            _until: Option<&str>,
231        ) -> anyhow::Result<Vec<MemoryEntry>> {
232            Ok(self.entries.as_ref().clone())
233        }
234
235        async fn get(&self, _key: &str) -> anyhow::Result<Option<MemoryEntry>> {
236            Ok(None)
237        }
238
239        async fn list(
240            &self,
241            _category: Option<&MemoryCategory>,
242            _session_id: Option<&str>,
243        ) -> anyhow::Result<Vec<MemoryEntry>> {
244            Ok(vec![])
245        }
246
247        async fn forget(&self, _key: &str) -> anyhow::Result<bool> {
248            Ok(true)
249        }
250
251        async fn forget_for_agent(&self, _key: &str, _agent_id: &str) -> anyhow::Result<bool> {
252            Ok(true)
253        }
254
255        async fn count(&self) -> anyhow::Result<usize> {
256            Ok(self.entries.len())
257        }
258
259        async fn health_check(&self) -> bool {
260            true
261        }
262
263        fn name(&self) -> &str {
264            "mock-with-entries"
265        }
266
267        async fn store_with_agent(
268            &self,
269            _key: &str,
270            _content: &str,
271            _category: MemoryCategory,
272            _session_id: Option<&str>,
273            _namespace: Option<&str>,
274            _importance: Option<f64>,
275            _agent_id: Option<&str>,
276        ) -> anyhow::Result<()> {
277            Ok(())
278        }
279
280        async fn recall_for_agents(
281            &self,
282            _allowed_agent_ids: &[&str],
283            query: &str,
284            limit: usize,
285            session_id: Option<&str>,
286            since: Option<&str>,
287            until: Option<&str>,
288        ) -> anyhow::Result<Vec<MemoryEntry>> {
289            self.recall(query, limit, session_id, since, until).await
290        }
291    }
292    impl ::zeroclaw_api::attribution::Attributable for MockMemoryWithEntries {
293        fn role(&self) -> ::zeroclaw_api::attribution::Role {
294            ::zeroclaw_api::attribution::Role::Memory(
295                ::zeroclaw_api::attribution::MemoryKind::InMemory,
296            )
297        }
298        fn alias(&self) -> &str {
299            "MockMemoryWithEntries"
300        }
301    }
302
303    #[tokio::test]
304    async fn default_loader_formats_context() {
305        let loader = DefaultMemoryLoader::default();
306        let context = loader
307            .load_context(&MockMemory, "hello", None)
308            .await
309            .unwrap();
310        assert_eq!(
311            context,
312            format!("{MEMORY_CONTEXT_OPEN}\n- k: v\n{MEMORY_CONTEXT_CLOSE}\n\n")
313        );
314    }
315
316    #[tokio::test]
317    async fn default_loader_skips_legacy_assistant_autosave_entries() {
318        let loader = DefaultMemoryLoader::new(5, 0.0);
319        let memory = MockMemoryWithEntries {
320            entries: Arc::new(vec![
321                MemoryEntry {
322                    id: "1".into(),
323                    key: "assistant_resp_legacy".into(),
324                    content: "fabricated detail".into(),
325                    category: MemoryCategory::Daily,
326                    timestamp: "now".into(),
327                    session_id: None,
328                    score: Some(0.95),
329                    namespace: "default".into(),
330                    importance: None,
331                    superseded_by: None,
332                    agent_alias: None,
333                    agent_id: None,
334                },
335                MemoryEntry {
336                    id: "2".into(),
337                    key: "user_fact".into(),
338                    content: "User prefers concise answers".into(),
339                    category: MemoryCategory::Conversation,
340                    timestamp: "now".into(),
341                    session_id: None,
342                    score: Some(0.9),
343                    namespace: "default".into(),
344                    importance: None,
345                    superseded_by: None,
346                    agent_alias: None,
347                    agent_id: None,
348                },
349            ]),
350        };
351
352        let context = loader
353            .load_context(&memory, "answer style", None)
354            .await
355            .unwrap();
356        assert!(context.contains("user_fact"));
357        assert!(!context.contains("assistant_resp_legacy"));
358        assert!(!context.contains("fabricated detail"));
359    }
360
361    #[tokio::test]
362    async fn default_loader_skips_user_autosave_entries() {
363        let loader = DefaultMemoryLoader::new(5, 0.0);
364        let memory = MockMemoryWithEntries {
365            entries: Arc::new(vec![
366                MemoryEntry {
367                    id: "1".into(),
368                    key: "user_msg_e5f6g7h8".into(),
369                    content: "User message embedding prior context verbatim".into(),
370                    category: MemoryCategory::Conversation,
371                    timestamp: "now".into(),
372                    session_id: None,
373                    score: Some(0.95),
374                    namespace: "default".into(),
375                    importance: None,
376                    superseded_by: None,
377                    agent_alias: None,
378                    agent_id: None,
379                },
380                MemoryEntry {
381                    id: "2".into(),
382                    key: "user_fact".into(),
383                    content: "User prefers concise answers".into(),
384                    category: MemoryCategory::Conversation,
385                    timestamp: "now".into(),
386                    session_id: None,
387                    score: Some(0.9),
388                    namespace: "default".into(),
389                    importance: None,
390                    superseded_by: None,
391                    agent_alias: None,
392                    agent_id: None,
393                },
394            ]),
395        };
396
397        let context = loader
398            .load_context(&memory, "answer style", None)
399            .await
400            .unwrap();
401        assert!(context.contains("user_fact"));
402        assert!(!context.contains("user_msg_e5f6g7h8"));
403        assert!(!context.contains("embedding prior context"));
404    }
405}
zeroclaw_runtime/agent/memory_loader.rs

zeroclaw_runtime/agent/
memory_loader.rs
