mw_memory/
engine.rs

1//! The pluggable retrieval backend.
2//!
3//! MemoryWhale owns the [`MemoryEngine`] interface. By default it's backed by the
4//! [`BuiltinEngine`] (the explainable scorer over a local memory set). With the
5//! off-by-default `mempalace` feature, `MemPalaceEngine` sits behind the same
6//! interface and talks to `mempalace-mcp` as an MCP *client* — so MemPalace stays
7//! an optional, swappable backend rather than a hard dependency. Callers never
8//! change.
9
10use std::sync::Arc;
11
12use crate::embed::Embedder;
13use crate::scorer::score;
14use crate::{Memory, Query, ScoredMemory, Weights};
15
16pub trait MemoryEngine {
17    fn name(&self) -> &str;
18    /// Return the top-`k` memories for the query, each with its score + reasons.
19    fn retrieve(&self, query: &Query, k: usize) -> Vec<ScoredMemory>;
20    /// Full explanation for a single memory id (the `memory explain <id>` view).
21    fn explain(&self, id: i64, query: &Query) -> Option<ScoredMemory>;
22}
23
24/// The default, zero-setup engine: scores an in-memory set with the explainable
25/// scorer. (A SQLite-backed variant just loads `memories` from the DB first.)
26///
27/// With an [`Embedder`] attached, similarity becomes semantic (cosine over
28/// embeddings); without one, it falls back to lexical term overlap.
29pub struct BuiltinEngine {
30    pub memories: Vec<Memory>,
31    pub weights: Weights,
32    embedder: Option<Arc<dyn Embedder>>,
33}
34
35impl BuiltinEngine {
36    pub fn new(memories: Vec<Memory>) -> Self {
37        Self {
38            memories,
39            weights: Weights::default(),
40            embedder: None,
41        }
42    }
43
44    pub fn with_weights(mut self, weights: Weights) -> Self {
45        self.weights = weights;
46        self
47    }
48
49    /// Attach an embedder and precompute embeddings for every memory that lacks
50    /// one. Returns an error if embedding fails (e.g. Ollama not running).
51    pub fn with_embedder(mut self, embedder: Arc<dyn Embedder>) -> anyhow::Result<Self> {
52        for m in &mut self.memories {
53            if m.embedding.is_none() {
54                m.embedding = Some(embedder.embed(&m.text)?);
55            }
56        }
57        self.embedder = Some(embedder);
58        Ok(self)
59    }
60
61    /// Embed the query text if an embedder is attached (best-effort).
62    fn query_embedding(&self, query: &Query) -> Option<Vec<f32>> {
63        self.embedder.as_ref().and_then(|e| e.embed(&query.text).ok())
64    }
65}
66
67impl MemoryEngine for BuiltinEngine {
68    fn name(&self) -> &str {
69        if self.embedder.is_some() {
70            "builtin+embeddings"
71        } else {
72            "builtin"
73        }
74    }
75
76    fn retrieve(&self, query: &Query, k: usize) -> Vec<ScoredMemory> {
77        let qe = self.query_embedding(query);
78        let mut scored: Vec<ScoredMemory> = self
79            .memories
80            .iter()
81            .map(|m| score(m, query, &self.weights, qe.as_deref()))
82            .collect();
83        scored.sort_by(|a, b| {
84            b.score
85                .partial_cmp(&a.score)
86                .unwrap_or(std::cmp::Ordering::Equal)
87        });
88        scored.truncate(k);
89        scored
90    }
91
92    fn explain(&self, id: i64, query: &Query) -> Option<ScoredMemory> {
93        let qe = self.query_embedding(query);
94        self.memories
95            .iter()
96            .find(|m| m.id == id)
97            .map(|m| score(m, query, &self.weights, qe.as_deref()))
98    }
99}
100
101/// MemPalace as a retrieval backend, spoken to as an **MCP client** over stdio.
102///
103/// Behind the off-by-default `mempalace` feature: the default build (and
104/// `cargo install` of the CLI) pulls in nothing extra.
105///
106/// Each call spawns the configured command (default `mempalace-mcp`), does the
107/// MCP handshake, calls the search tool, and maps its hits into [`ScoredMemory`].
108/// MemPalace does its own ranking, so we carry its relevance across as a single
109/// `similarity` [`Signal`] rather than re-scoring — the reason string names the
110/// source ("mempalace semantic score 0.87").
111///
112/// Expected search-tool result: JSON text content holding an array (or
113/// `{"results": [...]}`) of objects with `text` and `score`, optionally `id`,
114/// `tags`, `created_at`, `last_used`, `mentions`, `importance`.
115#[cfg(feature = "mempalace")]
116pub struct MemPalaceEngine {
117    /// The MCP server command (default: `mempalace-mcp`).
118    pub command: String,
119    /// Extra argv for that command.
120    pub args: Vec<String>,
121    /// The search tool to call (default: `search`).
122    pub tool: String,
123}
124
125#[cfg(feature = "mempalace")]
126impl Default for MemPalaceEngine {
127    fn default() -> Self {
128        Self {
129            command: "mempalace-mcp".into(),
130            args: Vec::new(),
131            tool: "search".into(),
132        }
133    }
134}
135
136#[cfg(feature = "mempalace")]
137impl MemPalaceEngine {
138    pub fn new(command: impl Into<String>, args: Vec<String>) -> Self {
139        Self {
140            command: command.into(),
141            args,
142            tool: "search".into(),
143        }
144    }
145
146    pub fn with_tool(mut self, tool: impl Into<String>) -> Self {
147        self.tool = tool.into();
148        self
149    }
150
151    /// The fallible retrieval. [`MemoryEngine::retrieve`] can't return an error
152    /// (the trait yields a `Vec`), so callers who need to *handle* a missing
153    /// server or a failed handshake should call this instead.
154    pub fn try_retrieve(&self, query: &Query, k: usize) -> anyhow::Result<Vec<ScoredMemory>> {
155        let mut client = crate::mcp::McpClient::spawn(&self.command, &self.args)?;
156        let tools = client.list_tools()?;
157        if let Some(names) = tools.get("tools").and_then(serde_json::Value::as_array) {
158            if !names
159                .iter()
160                .any(|t| t.get("name").and_then(serde_json::Value::as_str) == Some(&self.tool))
161            {
162                anyhow::bail!(
163                    "`{}` does not expose a `{}` tool",
164                    self.command,
165                    self.tool
166                );
167            }
168        }
169        let text = client.call_tool(
170            &self.tool,
171            serde_json::json!({"query": query.text, "limit": k}),
172        )?;
173        let mut out = map_hits(&text, query)?;
174        out.truncate(k);
175        Ok(out)
176    }
177}
178
179/// Map a MemPalace search payload into scored memories. Split out from the
180/// process plumbing so it is testable against captured JSON.
181#[cfg(feature = "mempalace")]
182fn map_hits(payload: &str, query: &Query) -> anyhow::Result<Vec<ScoredMemory>> {
183    use anyhow::Context;
184    use serde_json::Value;
185
186    let parsed: Value = serde_json::from_str(payload)
187        .context("mempalace search result was not JSON")?;
188    let hits = match &parsed {
189        Value::Array(a) => a.clone(),
190        Value::Object(o) => o
191            .get("results")
192            .and_then(Value::as_array)
193            .cloned()
194            .ok_or_else(|| anyhow::anyhow!("mempalace search result has no `results` array"))?,
195        _ => anyhow::bail!("unexpected mempalace search result shape"),
196    };
197
198    let ts = |h: &Value, key: &str| {
199        h.get(key)
200            .and_then(Value::as_str)
201            .and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok())
202            .map(|d| d.with_timezone(&chrono::Utc))
203            .unwrap_or(query.now)
204    };
205
206    Ok(hits
207        .iter()
208        .enumerate()
209        .map(|(i, h)| {
210            let score = h.get("score").and_then(Value::as_f64).unwrap_or(0.0) as f32;
211            let score = score.clamp(0.0, 1.0);
212            let memory = Memory {
213                id: h.get("id").and_then(Value::as_i64).unwrap_or(i as i64),
214                text: h
215                    .get("text")
216                    .and_then(Value::as_str)
217                    .unwrap_or_default()
218                    .to_string(),
219                created_at: ts(h, "created_at"),
220                last_used: ts(h, "last_used"),
221                mentions: h.get("mentions").and_then(Value::as_u64).unwrap_or(0) as u32,
222                importance: h.get("importance").and_then(Value::as_f64).unwrap_or(0.0) as f32,
223                tags: h
224                    .get("tags")
225                    .and_then(Value::as_array)
226                    .map(|t| {
227                        t.iter()
228                            .filter_map(Value::as_str)
229                            .map(str::to_string)
230                            .collect()
231                    })
232                    .unwrap_or_default(),
233                embedding: None,
234            };
235            ScoredMemory {
236                memory,
237                score,
238                signals: vec![crate::Signal {
239                    name: "similarity".into(),
240                    weight: 1.0,
241                    score,
242                    applicable: true,
243                    detail: format!("mempalace semantic score {score:.2}"),
244                }],
245            }
246        })
247        .collect())
248}
249
250#[cfg(feature = "mempalace")]
251impl MemoryEngine for MemPalaceEngine {
252    fn name(&self) -> &str {
253        "mempalace"
254    }
255
256    fn retrieve(&self, query: &Query, k: usize) -> Vec<ScoredMemory> {
257        match self.try_retrieve(query, k) {
258            Ok(hits) => hits,
259            Err(e) => {
260                eprintln!("[mw-memory] mempalace retrieval failed: {e:#}");
261                Vec::new()
262            }
263        }
264    }
265
266    fn explain(&self, id: i64, query: &Query) -> Option<ScoredMemory> {
267        // MemPalace ranks server-side; re-run the query and pick the hit out.
268        self.try_retrieve(query, 50)
269            .ok()?
270            .into_iter()
271            .find(|s| s.memory.id == id)
272    }
273}
274
275/// One item to file into MemPalace: a verbatim `content` string placed under a
276/// `wing` (project) and `room` (aspect). Mirrors `mempalace_add_drawer` inputs.
277#[cfg(feature = "mempalace")]
278pub struct Drawer {
279    pub wing: String,
280    pub room: String,
281    pub content: String,
282}
283
284/// Outcome of a checkpoint push, parsed from the tool's summary.
285#[cfg(feature = "mempalace")]
286#[derive(Debug, Default, PartialEq, Eq)]
287pub struct CheckpointOutcome {
288    pub added: usize,
289    pub duplicates: usize,
290    pub errors: usize,
291}
292
293/// Push memories into a running MemPalace server in one batch via its
294/// `mempalace_checkpoint` tool (which semantic-dedups, then files the
295/// non-duplicates). Spawns `command args…`, does the MCP handshake, and calls
296/// `tool` once with all items. Returns the server's added/duplicate/error tally.
297#[cfg(feature = "mempalace")]
298pub fn checkpoint(
299    command: &str,
300    args: &[String],
301    tool: &str,
302    items: &[Drawer],
303) -> anyhow::Result<CheckpointOutcome> {
304    use serde_json::{json, Value};
305    let mut client = crate::mcp::McpClient::spawn(command, args)?;
306    let payload = json!({
307        "items": items
308            .iter()
309            .map(|d| json!({"wing": d.wing, "room": d.room, "content": d.content}))
310            .collect::<Vec<_>>(),
311        "added_by": "memorywhale",
312    });
313    let text = client.call_tool(tool, payload)?;
314    // The summary is JSON with added/duplicates/errors arrays; tolerate shape
315    // drift by counting whatever arrays are present rather than requiring them.
316    let v: Value = serde_json::from_str(&text).unwrap_or(Value::Null);
317    let count = |k: &str| v.get(k).and_then(Value::as_array).map(Vec::len).unwrap_or(0);
318    Ok(CheckpointOutcome {
319        added: count("added"),
320        duplicates: count("duplicates"),
321        errors: count("errors"),
322    })
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328    use chrono::{Duration, TimeZone, Utc};
329
330    fn now() -> chrono::DateTime<Utc> {
331        Utc.with_ymd_and_hms(2026, 6, 27, 12, 0, 0).unwrap()
332    }
333
334    fn sample() -> Vec<Memory> {
335        let n = now();
336        vec![
337            Memory { id: 143, text: "I use Rust for systems software.".into(), created_at: n - Duration::days(20), last_used: n, mentions: 27, importance: 0.98, tags: vec!["rust".into()], embedding: None },
338            Memory { id: 7, text: "I ate pizza.".into(), created_at: n - Duration::days(40), last_used: n - Duration::days(40), mentions: 1, importance: 0.01, tags: vec![], embedding: None },
339            Memory { id: 22, text: "Use Tokio for async runtime.".into(), created_at: n - Duration::days(5), last_used: n - Duration::days(3), mentions: 6, importance: 0.6, tags: vec!["rust".into(), "tokio".into()], embedding: None },
340        ]
341    }
342
343    #[test]
344    fn retrieve_ranks_and_truncates() {
345        let eng = BuiltinEngine::new(sample());
346        let q = Query::new("which language for systems programming?", now());
347        let top = eng.retrieve(&q, 2);
348        assert_eq!(top.len(), 2);
349        assert_eq!(top[0].memory.id, 143); // Rust wins
350    }
351
352    #[test]
353    fn explain_returns_breakdown() {
354        let eng = BuiltinEngine::new(sample());
355        let q = Query::new("rust", now());
356        let e = eng.explain(143, &q).unwrap();
357        assert!(e.explain().contains("memory explain 143"));
358        assert!(eng.explain(9999, &q).is_none());
359    }
360
361    #[cfg(feature = "mempalace")]
362    #[test]
363    fn mempalace_maps_hits_to_signals() {
364        let payload = include_str!("../tests/fixtures/mempalace_search.json");
365        let hits = map_hits(payload, &Query::new("rust", now())).unwrap();
366        assert_eq!(hits.len(), 2);
367        assert_eq!(hits[0].memory.id, 143);
368        assert_eq!(hits[0].memory.tags, vec!["rust".to_string()]);
369        assert_eq!(hits[0].percent(), 87);
370        let sig = &hits[0].signals[0];
371        assert_eq!(sig.name, "similarity");
372        assert_eq!(sig.detail, "mempalace semantic score 0.87");
373        // No timestamps in the second hit -> falls back to the query's "now".
374        assert_eq!(hits[1].memory.created_at, now());
375    }
376
377    #[cfg(feature = "mempalace")]
378    #[test]
379    fn mempalace_rejects_garbage() {
380        assert!(map_hits("not json", &Query::new("x", now())).is_err());
381        assert!(map_hits("{\"oops\": 1}", &Query::new("x", now())).is_err());
382    }
383
384    /// End-to-end over a fake MCP server (a shell script replaying fixture
385    /// JSON) — no network, no real mempalace binary.
386    #[cfg(all(feature = "mempalace", unix))]
387    #[test]
388    fn mempalace_talks_to_a_fake_server() {
389        let script = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/fake-mempalace-mcp.sh");
390        let eng = MemPalaceEngine::new("sh", vec![script.to_string()]);
391        let hits = eng.try_retrieve(&Query::new("rust", now()), 10).unwrap();
392        assert_eq!(hits.len(), 2);
393        assert_eq!(hits[0].memory.id, 143);
394        assert!(hits[0].reasons()[0].starts_with("mempalace semantic score"));
395    }
396
397    #[cfg(all(feature = "mempalace", unix))]
398    #[test]
399    fn checkpoint_pushes_and_parses_the_summary() {
400        let script =
401            concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/fake-mempalace-checkpoint.sh");
402        let items = vec![
403            Drawer { wing: "memorywhale".into(), room: "command".into(), content: "cargo build failed".into() },
404            Drawer { wing: "memorywhale".into(), room: "note".into(), content: "the fix was X".into() },
405        ];
406        let out = checkpoint("sh", &[script.to_string()], "mempalace_checkpoint", &items).unwrap();
407        // Fixture reports two added, one duplicate, no errors.
408        assert_eq!(out, CheckpointOutcome { added: 2, duplicates: 1, errors: 0 });
409    }
410
411    #[cfg(feature = "mempalace")]
412    #[test]
413    fn missing_server_is_a_clear_error() {
414        let eng = MemPalaceEngine::new("mw-definitely-not-installed", vec![]);
415        let err = eng
416            .try_retrieve(&Query::new("x", now()), 3)
417            .unwrap_err()
418            .to_string();
419        assert!(err.contains("failed to start MCP server"), "{err}");
420    }
421}
mw_memory/engine.rs

mw_memory/
engine.rs
