Skip to main content

telos_agent/agent/context/
memory_injection.rs

1use std::sync::{Arc, Mutex};
2use std::{
3    collections::hash_map::DefaultHasher,
4    hash::{Hash, Hasher},
5};
6
7use crate::knowledge::memory::MemoryStore;
8use crate::model::message::SystemReminder;
9
10pub struct MemoryInjector {
11    store: Arc<Mutex<MemoryStore>>,
12    max_memories: usize,
13    min_relevance: f64,
14}
15
16pub struct MemoryInjection {
17    pub reminder: SystemReminder,
18    pub fingerprint: u64,
19}
20
21impl MemoryInjector {
22    pub fn new(store: Arc<Mutex<MemoryStore>>) -> Self {
23        Self { store, max_memories: 5, min_relevance: 0.10 }
24    }
25
26    pub fn with_max_memories(mut self, max: usize) -> Self {
27        self.max_memories = max;
28        self
29    }
30
31    pub fn with_min_relevance(mut self, threshold: f64) -> Self {
32        self.min_relevance = threshold.clamp(0.0, 1.0);
33        self
34    }
35
36    pub fn inject_for_query(&self, query: &str) -> Option<MemoryInjection> {
37        let store = self.store.lock().ok()?;
38        let memories = store.search_relevant(query, self.max_memories, self.min_relevance);
39        if memories.is_empty() {
40            return None;
41        }
42
43        let mut lines = vec!["## Relevant Memories".to_string(), String::new()];
44        for entry in &memories {
45            let status_tag = match entry.status {
46                crate::knowledge::memory::MemoryStatus::NeedsFix => " [needs_fix]",
47                _ => "",
48            };
49            lines.push(format!(
50                "- **{}** ({:?}{}): {}",
51                entry.name, entry.category, status_tag, entry.description,
52            ));
53            let preview: String = entry.body.chars().take(80).collect();
54            if !preview.is_empty() {
55                lines.push(format!("  {}", preview));
56            }
57        }
58        let content = lines.join("\n");
59        let mut hasher = DefaultHasher::new();
60        content.hash(&mut hasher);
61        Some(MemoryInjection {
62            reminder: SystemReminder::MemoryInjection { content },
63            fingerprint: hasher.finish(),
64        })
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71    use crate::knowledge::memory::format::{MemoryCategory, MemoryEntry, MemoryStatus};
72
73    fn test_entry(name: &str, body: &str) -> MemoryEntry {
74        MemoryEntry {
75            name: name.into(),
76            description: "Relevant memory".into(),
77            category: MemoryCategory::Fact,
78            tags: vec!["cache".into(), "prompt".into()],
79            created: "2026-06-23".into(),
80            updated: "2026-06-23".into(),
81            status: MemoryStatus::Working,
82            times_used: 3,
83            confidence: None,
84            related: vec![],
85            source_session: None,
86            body: body.into(),
87        }
88    }
89
90    #[test]
91    fn identical_relevant_memories_produce_stable_fingerprint() {
92        let dir = tempfile::tempdir().unwrap();
93        let mut store = MemoryStore::new(dir.path().to_path_buf());
94        store
95            .write(test_entry("cache-hit-guide", "Prompt cache hit guidance for desktop runtime."))
96            .unwrap();
97        let store = Arc::new(Mutex::new(store));
98        let injector = MemoryInjector::new(Arc::clone(&store));
99
100        let first = injector
101            .inject_for_query("How do I improve prompt cache hit rate in desktop runtime?")
102            .expect("first injection should exist");
103        let second = injector
104            .inject_for_query("How do I improve prompt cache hit rate in desktop runtime?")
105            .expect("second injection should exist");
106
107        assert_eq!(first.fingerprint, second.fingerprint);
108        assert_eq!(first.reminder, second.reminder);
109    }
110
111    #[test]
112    fn injection_preview_is_truncated_to_eighty_chars() {
113        let dir = tempfile::tempdir().unwrap();
114        let mut store = MemoryStore::new(dir.path().to_path_buf());
115        let long_body = "a".repeat(140);
116        store.write(test_entry("long-memory", &long_body)).unwrap();
117        let store = Arc::new(Mutex::new(store));
118        let injector = MemoryInjector::new(store);
119
120        let injection =
121            injector.inject_for_query("long memory prompt cache").expect("injection should exist");
122
123        let rendered = injection.reminder.render();
124        assert!(rendered.contains(&"a".repeat(80)));
125        assert!(!rendered.contains(&"a".repeat(81)));
126    }
127}