crabtalk_runtime/memory/
mod.rs

1//! Built-in memory — file-per-entry storage at `{config_dir}/memory/`.
2
3use crate::config::MemoryConfig;
4use std::{
5    collections::HashMap,
6    path::{Path, PathBuf},
7    sync::RwLock,
8};
9use wcore::model::{HistoryEntry, Role};
10
11pub mod bm25;
12pub mod entry;
13pub mod storage;
14pub mod tool;
15
16use entry::MemoryEntry;
17use storage::Storage;
18
19const MEMORY_PROMPT: &str = include_str!("../../prompts/memory.md");
20
21pub const DEFAULT_SOUL: &str = include_str!("../../prompts/crab.md");
22
23pub struct Memory {
24    storage: Box<dyn Storage>,
25    entries: RwLock<HashMap<String, MemoryEntry>>,
26    index: RwLock<String>,
27    index_path: PathBuf,
28    entries_dir: PathBuf,
29    config: MemoryConfig,
30}
31
32impl Memory {
33    /// Open (or create) memory storage at the given directory.
34    pub fn open(dir: PathBuf, config: MemoryConfig, storage: Box<dyn Storage>) -> Self {
35        let entries_dir = dir.join("entries");
36        let index_path = dir.join("MEMORY.md");
37
38        storage.create_dir_all(&entries_dir).ok();
39
40        let mem = Self {
41            storage,
42            entries: RwLock::new(HashMap::new()),
43            index: RwLock::new(String::new()),
44            index_path,
45            entries_dir,
46            config,
47        };
48
49        mem.migrate_legacy(&dir);
50        mem.load_entries();
51        mem.load_index();
52        mem
53    }
54
55    fn load_entries(&self) {
56        let paths = match self.storage.list(&self.entries_dir) {
57            Ok(p) => p,
58            Err(_) => return,
59        };
60
61        let mut entries = self.entries.write().unwrap();
62        for path in paths {
63            if path.extension().and_then(|e| e.to_str()) != Some("md") {
64                continue;
65            }
66            let raw = match self.storage.read(&path) {
67                Ok(r) => r,
68                Err(_) => continue,
69            };
70            match MemoryEntry::parse(path, &raw) {
71                Ok(entry) => {
72                    entries.insert(entry.name.clone(), entry);
73                }
74                Err(e) => {
75                    tracing::warn!("failed to parse memory entry: {e}");
76                }
77            }
78        }
79    }
80
81    fn load_index(&self) {
82        if let Ok(content) = self.storage.read(&self.index_path) {
83            *self.index.write().unwrap() = content;
84        }
85    }
86
87    /// BM25-ranked recall over all entries.
88    pub fn recall(&self, query: &str, limit: usize) -> String {
89        let entries = self.entries.read().unwrap();
90        if entries.is_empty() {
91            return "no memories found".to_owned();
92        }
93
94        let entry_vec: Vec<&MemoryEntry> = entries.values().collect();
95        let docs: Vec<(usize, String)> = entry_vec
96            .iter()
97            .enumerate()
98            .map(|(i, e)| (i, e.search_text()))
99            .collect();
100        let doc_refs: Vec<(usize, &str)> = docs.iter().map(|(i, s)| (*i, s.as_str())).collect();
101
102        let results = bm25::score(&doc_refs, query, limit);
103        if results.is_empty() {
104            return "no memories found".to_owned();
105        }
106
107        results
108            .iter()
109            .map(|(idx, _score)| {
110                let e = &entry_vec[*idx];
111                format!("## {}\n{}\n\n{}", e.name, e.description, e.content)
112            })
113            .collect::<Vec<_>>()
114            .join("\n---\n")
115    }
116
117    /// Create or update a memory entry.
118    pub fn remember(&self, name: String, description: String, content: String) -> String {
119        let entry = MemoryEntry::new(name.clone(), description, content, &self.entries_dir);
120        if let Err(e) = entry.save(self.storage.as_ref()) {
121            return format!("failed to save entry: {e}");
122        }
123        self.entries.write().unwrap().insert(name.clone(), entry);
124        format!("remembered: {name}")
125    }
126
127    /// Delete a memory entry by name.
128    pub fn forget(&self, name: &str) -> String {
129        let mut entries = self.entries.write().unwrap();
130        match entries.remove(name) {
131            Some(entry) => {
132                if let Err(e) = entry.delete(self.storage.as_ref()) {
133                    tracing::warn!("failed to delete entry file: {e}");
134                }
135                format!("forgot: {name}")
136            }
137            None => format!("no entry named: {name}"),
138        }
139    }
140
141    /// Overwrite MEMORY.md (the curated overview).
142    pub fn write_index(&self, content: &str) -> String {
143        if let Err(e) = self.storage.write(&self.index_path, content) {
144            return format!("failed to write MEMORY.md: {e}");
145        }
146        *self.index.write().unwrap() = content.to_owned();
147        "MEMORY.md updated".to_owned()
148    }
149
150    /// Build system prompt block from MEMORY.md content.
151    pub fn build_prompt(&self) -> String {
152        let index = self.index.read().unwrap();
153        if index.is_empty() {
154            return format!("\n\n{MEMORY_PROMPT}");
155        }
156        format!("\n\n<memory>\n{}\n</memory>\n\n{MEMORY_PROMPT}", *index)
157    }
158
159    /// Auto-recall from last user message, injected before each turn.
160    pub fn before_run(&self, history: &[HistoryEntry]) -> Vec<HistoryEntry> {
161        let last_user = history
162            .iter()
163            .rev()
164            .find(|e| *e.role() == Role::User && !e.text().is_empty());
165
166        let Some(entry) = last_user else {
167            return Vec::new();
168        };
169
170        let query: String = entry
171            .text()
172            .split_whitespace()
173            .take(8)
174            .collect::<Vec<_>>()
175            .join(" ");
176
177        if query.is_empty() {
178            return Vec::new();
179        }
180
181        let limit = self.config.recall_limit;
182        let result = self.recall(&query, limit);
183        if result == "no memories found" {
184            return Vec::new();
185        }
186
187        vec![HistoryEntry::user(format!("<recall>\n{result}\n</recall>")).auto_injected()]
188    }
189
190    fn migrate_legacy(&self, dir: &Path) {
191        let existing = self.storage.list(&self.entries_dir).unwrap_or_default();
192        if !existing.is_empty() {
193            return;
194        }
195
196        let memory_path = dir.join("memory.md");
197        let user_path = dir.join("user.md");
198        let facts_path = dir.join("facts.toml");
199
200        let has_legacy = self.storage.exists(&memory_path)
201            || self.storage.exists(&user_path)
202            || self.storage.exists(&facts_path);
203
204        if !has_legacy {
205            return;
206        }
207
208        if let Ok(content) = self.storage.read(&memory_path)
209            && !content.trim().is_empty()
210        {
211            self.storage.write(&self.index_path, &content).ok();
212
213            for (i, chunk) in content.split("\n\n").enumerate() {
214                let chunk = chunk.trim();
215                if chunk.is_empty() {
216                    continue;
217                }
218                let name = format!("migrated-memory-{}", i + 1);
219                let entry = MemoryEntry::new(
220                    name,
221                    "Migrated from memory.md".to_owned(),
222                    chunk.to_owned(),
223                    &self.entries_dir,
224                );
225                entry.save(self.storage.as_ref()).ok();
226            }
227            self.storage
228                .rename(&memory_path, &dir.join("memory.md.bak"))
229                .ok();
230        }
231
232        if let Ok(content) = self.storage.read(&user_path)
233            && !content.trim().is_empty()
234        {
235            let entry = MemoryEntry::new(
236                "user-profile".to_owned(),
237                "User profile migrated from user.md".to_owned(),
238                content,
239                &self.entries_dir,
240            );
241            entry.save(self.storage.as_ref()).ok();
242            self.storage
243                .rename(&user_path, &dir.join("user.md.bak"))
244                .ok();
245        }
246
247        if let Ok(content) = self.storage.read(&facts_path)
248            && !content.trim().is_empty()
249        {
250            let entry = MemoryEntry::new(
251                "known-facts".to_owned(),
252                "Known facts migrated from facts.toml".to_owned(),
253                content,
254                &self.entries_dir,
255            );
256            entry.save(self.storage.as_ref()).ok();
257            self.storage
258                .rename(&facts_path, &dir.join("facts.toml.bak"))
259                .ok();
260        }
261    }
262}
crabtalk_runtime/memory/mod.rs

crabtalk_runtime/memory/
mod.rs
