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::{Message, 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: &[Message]) -> Vec<Message> {
161        let last_user = history
162            .iter()
163            .rev()
164            .find(|m| m.role == Role::User && !m.content.is_empty());
165
166        let Some(msg) = last_user else {
167            return Vec::new();
168        };
169
170        let query: String = msg
171            .content
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![Message {
188            role: Role::User,
189            content: format!("<recall>\n{result}\n</recall>"),
190            auto_injected: true,
191            ..Default::default()
192        }]
193    }
194
195    fn migrate_legacy(&self, dir: &Path) {
196        let existing = self.storage.list(&self.entries_dir).unwrap_or_default();
197        if !existing.is_empty() {
198            return;
199        }
200
201        let memory_path = dir.join("memory.md");
202        let user_path = dir.join("user.md");
203        let facts_path = dir.join("facts.toml");
204
205        let has_legacy = self.storage.exists(&memory_path)
206            || self.storage.exists(&user_path)
207            || self.storage.exists(&facts_path);
208
209        if !has_legacy {
210            return;
211        }
212
213        if let Ok(content) = self.storage.read(&memory_path)
214            && !content.trim().is_empty()
215        {
216            self.storage.write(&self.index_path, &content).ok();
217
218            for (i, chunk) in content.split("\n\n").enumerate() {
219                let chunk = chunk.trim();
220                if chunk.is_empty() {
221                    continue;
222                }
223                let name = format!("migrated-memory-{}", i + 1);
224                let entry = MemoryEntry::new(
225                    name,
226                    "Migrated from memory.md".to_owned(),
227                    chunk.to_owned(),
228                    &self.entries_dir,
229                );
230                entry.save(self.storage.as_ref()).ok();
231            }
232            self.storage
233                .rename(&memory_path, &dir.join("memory.md.bak"))
234                .ok();
235        }
236
237        if let Ok(content) = self.storage.read(&user_path)
238            && !content.trim().is_empty()
239        {
240            let entry = MemoryEntry::new(
241                "user-profile".to_owned(),
242                "User profile migrated from user.md".to_owned(),
243                content,
244                &self.entries_dir,
245            );
246            entry.save(self.storage.as_ref()).ok();
247            self.storage
248                .rename(&user_path, &dir.join("user.md.bak"))
249                .ok();
250        }
251
252        if let Ok(content) = self.storage.read(&facts_path)
253            && !content.trim().is_empty()
254        {
255            let entry = MemoryEntry::new(
256                "known-facts".to_owned(),
257                "Known facts migrated from facts.toml".to_owned(),
258                content,
259                &self.entries_dir,
260            );
261            entry.save(self.storage.as_ref()).ok();
262            self.storage
263                .rename(&facts_path, &dir.join("facts.toml.bak"))
264                .ok();
265        }
266    }
267}
crabtalk_runtime/memory/mod.rs

crabtalk_runtime/memory/
mod.rs
