1use std::collections::HashMap;
11use std::sync::Arc;
12
13use rusqlite::Connection;
14
15use crate::embed::Embedder;
16use crate::scorer::score_with_lexical;
17use crate::{Memory, Query, ScoredMemory, Weights};
18
19fn bm25_similarities(memories: &[Memory], query: &str) -> HashMap<i64, f32> {
32 let mut out = HashMap::new();
33 let match_expr = fts_match_expr(query);
34 if match_expr.is_empty() {
35 return out;
36 }
37 let conn = match Connection::open_in_memory() {
38 Ok(c) => c,
39 Err(_) => return out,
40 };
41 if conn
42 .execute("CREATE VIRTUAL TABLE mem_fts USING fts5(text)", [])
43 .is_err()
44 {
45 return out;
46 }
47 {
48 let mut ins = match conn.prepare("INSERT INTO mem_fts(rowid, text) VALUES (?1, ?2)") {
49 Ok(s) => s,
50 Err(_) => return out,
51 };
52 for m in memories {
54 let _ = ins.execute(rusqlite::params![m.id, m.text]);
55 }
56 }
57 let mut stmt = match conn
58 .prepare("SELECT rowid, bm25(mem_fts) FROM mem_fts WHERE mem_fts MATCH ?1")
59 {
60 Ok(s) => s,
61 Err(_) => return out,
62 };
63 let rows = stmt.query_map([&match_expr], |r| {
64 Ok((r.get::<_, i64>(0)?, r.get::<_, f64>(1)?))
65 });
66 if let Ok(rows) = rows {
67 for (id, bm25) in rows.flatten() {
68 let x = (-bm25).max(0.0) as f32; out.insert(id, x / (1.0 + x));
70 }
71 }
72 out
73}
74
75fn fts_match_expr(query: &str) -> String {
80 let mut seen = std::collections::BTreeSet::new();
81 query
82 .split(|c: char| !c.is_alphanumeric())
83 .map(|w| w.to_lowercase())
84 .filter(|w| w.len() >= 2 && seen.insert(w.clone()))
85 .map(|t| format!("\"{t}\""))
86 .collect::<Vec<_>>()
87 .join(" OR ")
88}
89
90pub trait MemoryEngine {
91 fn name(&self) -> &str;
92 fn retrieve(&self, query: &Query, k: usize) -> Vec<ScoredMemory>;
94 fn explain(&self, id: i64, query: &Query) -> Option<ScoredMemory>;
96}
97
98pub struct BuiltinEngine {
104 pub memories: Vec<Memory>,
105 pub weights: Weights,
106 embedder: Option<Arc<dyn Embedder>>,
107}
108
109impl BuiltinEngine {
110 pub fn new(memories: Vec<Memory>) -> Self {
111 Self {
112 memories,
113 weights: Weights::default(),
114 embedder: None,
115 }
116 }
117
118 pub fn with_weights(mut self, weights: Weights) -> Self {
119 self.weights = weights;
120 self
121 }
122
123 pub fn with_embedder(mut self, embedder: Arc<dyn Embedder>) -> anyhow::Result<Self> {
126 for m in &mut self.memories {
127 if m.embedding.is_none() {
128 m.embedding = Some(embedder.embed(&m.text)?);
129 }
130 }
131 self.embedder = Some(embedder);
132 Ok(self)
133 }
134
135 fn query_embedding(&self, query: &Query) -> Option<Vec<f32>> {
137 self.embedder.as_ref().and_then(|e| e.embed(&query.text).ok())
138 }
139}
140
141impl MemoryEngine for BuiltinEngine {
142 fn name(&self) -> &str {
143 if self.embedder.is_some() {
144 "builtin+embeddings"
145 } else {
146 "builtin"
147 }
148 }
149
150 fn retrieve(&self, query: &Query, k: usize) -> Vec<ScoredMemory> {
151 let qe = self.query_embedding(query);
152 let sims = if qe.is_none() {
155 bm25_similarities(&self.memories, &query.text)
156 } else {
157 HashMap::new()
158 };
159 let mut scored: Vec<ScoredMemory> = self
160 .memories
161 .iter()
162 .map(|m| {
163 score_with_lexical(m, query, &self.weights, qe.as_deref(), sims.get(&m.id).copied())
164 })
165 .collect();
166 scored.sort_by(|a, b| {
167 b.score
168 .partial_cmp(&a.score)
169 .unwrap_or(std::cmp::Ordering::Equal)
170 });
171 scored.truncate(k);
172 scored
173 }
174
175 fn explain(&self, id: i64, query: &Query) -> Option<ScoredMemory> {
176 let qe = self.query_embedding(query);
177 let sims = if qe.is_none() {
178 bm25_similarities(&self.memories, &query.text)
179 } else {
180 HashMap::new()
181 };
182 self.memories
183 .iter()
184 .find(|m| m.id == id)
185 .map(|m| {
186 score_with_lexical(m, query, &self.weights, qe.as_deref(), sims.get(&m.id).copied())
187 })
188 }
189}
190
191#[cfg(feature = "mempalace")]
206pub struct MemPalaceEngine {
207 pub command: String,
209 pub args: Vec<String>,
211 pub tool: String,
213}
214
215#[cfg(feature = "mempalace")]
216impl Default for MemPalaceEngine {
217 fn default() -> Self {
218 Self {
219 command: "mempalace-mcp".into(),
220 args: Vec::new(),
221 tool: "search".into(),
222 }
223 }
224}
225
226#[cfg(feature = "mempalace")]
227impl MemPalaceEngine {
228 pub fn new(command: impl Into<String>, args: Vec<String>) -> Self {
229 Self {
230 command: command.into(),
231 args,
232 tool: "search".into(),
233 }
234 }
235
236 pub fn with_tool(mut self, tool: impl Into<String>) -> Self {
237 self.tool = tool.into();
238 self
239 }
240
241 pub fn try_retrieve(&self, query: &Query, k: usize) -> anyhow::Result<Vec<ScoredMemory>> {
245 let mut client = crate::mcp::McpClient::spawn(&self.command, &self.args)?;
246 let tools = client.list_tools()?;
247 if let Some(names) = tools.get("tools").and_then(serde_json::Value::as_array) {
248 if !names
249 .iter()
250 .any(|t| t.get("name").and_then(serde_json::Value::as_str) == Some(&self.tool))
251 {
252 anyhow::bail!(
253 "`{}` does not expose a `{}` tool",
254 self.command,
255 self.tool
256 );
257 }
258 }
259 let text = client.call_tool(
260 &self.tool,
261 serde_json::json!({"query": query.text, "limit": k}),
262 )?;
263 let mut out = map_hits(&text, query)?;
264 out.truncate(k);
265 Ok(out)
266 }
267}
268
269#[cfg(feature = "mempalace")]
272fn map_hits(payload: &str, query: &Query) -> anyhow::Result<Vec<ScoredMemory>> {
273 use anyhow::Context;
274 use serde_json::Value;
275
276 let parsed: Value = serde_json::from_str(payload)
277 .context("mempalace search result was not JSON")?;
278 let hits = match &parsed {
279 Value::Array(a) => a.clone(),
280 Value::Object(o) => o
281 .get("results")
282 .and_then(Value::as_array)
283 .cloned()
284 .ok_or_else(|| anyhow::anyhow!("mempalace search result has no `results` array"))?,
285 _ => anyhow::bail!("unexpected mempalace search result shape"),
286 };
287
288 let ts = |h: &Value, key: &str| {
289 h.get(key)
290 .and_then(Value::as_str)
291 .and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok())
292 .map(|d| d.with_timezone(&chrono::Utc))
293 .unwrap_or(query.now)
294 };
295
296 Ok(hits
297 .iter()
298 .enumerate()
299 .map(|(i, h)| {
300 let score = h.get("score").and_then(Value::as_f64).unwrap_or(0.0) as f32;
301 let score = score.clamp(0.0, 1.0);
302 let memory = Memory {
303 id: h.get("id").and_then(Value::as_i64).unwrap_or(i as i64),
304 text: h
305 .get("text")
306 .and_then(Value::as_str)
307 .unwrap_or_default()
308 .to_string(),
309 created_at: ts(h, "created_at"),
310 last_used: ts(h, "last_used"),
311 mentions: h.get("mentions").and_then(Value::as_u64).unwrap_or(0) as u32,
312 importance: h.get("importance").and_then(Value::as_f64).unwrap_or(0.0) as f32,
313 tags: h
314 .get("tags")
315 .and_then(Value::as_array)
316 .map(|t| {
317 t.iter()
318 .filter_map(Value::as_str)
319 .map(str::to_string)
320 .collect()
321 })
322 .unwrap_or_default(),
323 embedding: None,
324 };
325 ScoredMemory {
326 memory,
327 score,
328 signals: vec![crate::Signal {
329 name: "similarity".into(),
330 weight: 1.0,
331 score,
332 applicable: true,
333 detail: format!("mempalace semantic score {score:.2}"),
334 }],
335 }
336 })
337 .collect())
338}
339
340#[cfg(feature = "mempalace")]
341impl MemoryEngine for MemPalaceEngine {
342 fn name(&self) -> &str {
343 "mempalace"
344 }
345
346 fn retrieve(&self, query: &Query, k: usize) -> Vec<ScoredMemory> {
347 match self.try_retrieve(query, k) {
348 Ok(hits) => hits,
349 Err(e) => {
350 eprintln!("[mw-memory] mempalace retrieval failed: {e:#}");
351 Vec::new()
352 }
353 }
354 }
355
356 fn explain(&self, id: i64, query: &Query) -> Option<ScoredMemory> {
357 self.try_retrieve(query, 50)
359 .ok()?
360 .into_iter()
361 .find(|s| s.memory.id == id)
362 }
363}
364
365#[cfg(feature = "mempalace")]
368pub struct Drawer {
369 pub wing: String,
370 pub room: String,
371 pub content: String,
372}
373
374#[cfg(feature = "mempalace")]
376#[derive(Debug, Default, PartialEq, Eq)]
377pub struct CheckpointOutcome {
378 pub added: usize,
379 pub duplicates: usize,
380 pub errors: usize,
381}
382
383#[cfg(feature = "mempalace")]
388pub fn checkpoint(
389 command: &str,
390 args: &[String],
391 tool: &str,
392 items: &[Drawer],
393) -> anyhow::Result<CheckpointOutcome> {
394 use serde_json::{json, Value};
395 let mut client = crate::mcp::McpClient::spawn(command, args)?;
396 let payload = json!({
397 "items": items
398 .iter()
399 .map(|d| json!({"wing": d.wing, "room": d.room, "content": d.content}))
400 .collect::<Vec<_>>(),
401 "added_by": "memorywhale",
402 });
403 let text = client.call_tool(tool, payload)?;
404 let v: Value = serde_json::from_str(&text).unwrap_or(Value::Null);
407 let count = |k: &str| v.get(k).and_then(Value::as_array).map(Vec::len).unwrap_or(0);
408 Ok(CheckpointOutcome {
409 added: count("added"),
410 duplicates: count("duplicates"),
411 errors: count("errors"),
412 })
413}
414
415#[cfg(feature = "mempalace")]
418pub enum SyncOp {
419 Add {
420 wing: String,
421 room: String,
422 content: String,
423 added_by: String,
424 },
425 Delete {
426 drawer_id: String,
427 },
428}
429
430#[cfg(feature = "mempalace")]
432pub enum SyncResult {
433 Added { drawer_id: String },
435 Deleted,
437}
438
439#[cfg(feature = "mempalace")]
446pub fn sync_ops(
447 command: &str,
448 args: &[String],
449 add_tool: &str,
450 delete_tool: &str,
451 ops: &[SyncOp],
452) -> anyhow::Result<Vec<SyncResult>> {
453 use anyhow::Context;
454 use serde_json::{json, Value};
455 let mut client = crate::mcp::McpClient::spawn(command, args)?;
456 let mut out = Vec::with_capacity(ops.len());
457 for op in ops {
458 match op {
459 SyncOp::Add { wing, room, content, added_by } => {
460 let text = client.call_tool(
461 add_tool,
462 json!({"wing": wing, "room": room, "content": content, "added_by": added_by}),
463 )?;
464 let v: Value =
465 serde_json::from_str(&text).context("add_drawer result was not JSON")?;
466 let drawer_id = v
467 .get("drawer_id")
468 .and_then(Value::as_str)
469 .ok_or_else(|| anyhow::anyhow!("add_drawer returned no drawer_id"))?
470 .to_string();
471 out.push(SyncResult::Added { drawer_id });
472 }
473 SyncOp::Delete { drawer_id } => {
474 client.call_tool(delete_tool, json!({"drawer_id": drawer_id}))?;
475 out.push(SyncResult::Deleted);
476 }
477 }
478 }
479 Ok(out)
480}
481
482#[cfg(test)]
483mod tests {
484 use super::*;
485 use chrono::{Duration, TimeZone, Utc};
486
487 fn now() -> chrono::DateTime<Utc> {
488 Utc.with_ymd_and_hms(2026, 6, 27, 12, 0, 0).unwrap()
489 }
490
491 fn sample() -> Vec<Memory> {
492 let n = now();
493 vec![
494 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 },
495 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 },
496 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 },
497 ]
498 }
499
500 #[test]
501 fn retrieve_ranks_and_truncates() {
502 let eng = BuiltinEngine::new(sample());
503 let q = Query::new("which language for systems programming?", now());
504 let top = eng.retrieve(&q, 2);
505 assert_eq!(top.len(), 2);
506 assert_eq!(top[0].memory.id, 143); }
508
509 #[test]
510 fn bm25_similarity_signal_fires() {
511 let eng = BuiltinEngine::new(sample());
515 let q = Query::new("pizza", now());
516
517 let matching = eng.explain(7, &q).unwrap(); let sim_m = matching.signals.iter().find(|s| s.name == "similarity").unwrap();
519 assert!(sim_m.detail.contains("BM25"), "should use BM25: {}", sim_m.detail);
520 assert!(sim_m.score > 0.0, "matching memory sim should fire: {}", sim_m.score);
521
522 let non = eng.explain(143, &q).unwrap(); let sim_n = non.signals.iter().find(|s| s.name == "similarity").unwrap();
524 assert_eq!(sim_n.score, 0.0, "non-matching sim should be 0: {}", sim_n.score);
525 assert!(sim_m.score > sim_n.score);
526 }
527
528 #[test]
529 fn explain_returns_breakdown() {
530 let eng = BuiltinEngine::new(sample());
531 let q = Query::new("rust", now());
532 let e = eng.explain(143, &q).unwrap();
533 assert!(e.explain().contains("memory explain 143"));
534 assert!(eng.explain(9999, &q).is_none());
535 }
536
537 #[cfg(feature = "mempalace")]
538 #[test]
539 fn mempalace_maps_hits_to_signals() {
540 let payload = include_str!("../tests/fixtures/mempalace_search.json");
541 let hits = map_hits(payload, &Query::new("rust", now())).unwrap();
542 assert_eq!(hits.len(), 2);
543 assert_eq!(hits[0].memory.id, 143);
544 assert_eq!(hits[0].memory.tags, vec!["rust".to_string()]);
545 assert_eq!(hits[0].percent(), 87);
546 let sig = &hits[0].signals[0];
547 assert_eq!(sig.name, "similarity");
548 assert_eq!(sig.detail, "mempalace semantic score 0.87");
549 assert_eq!(hits[1].memory.created_at, now());
551 }
552
553 #[cfg(feature = "mempalace")]
554 #[test]
555 fn mempalace_rejects_garbage() {
556 assert!(map_hits("not json", &Query::new("x", now())).is_err());
557 assert!(map_hits("{\"oops\": 1}", &Query::new("x", now())).is_err());
558 }
559
560 #[cfg(all(feature = "mempalace", unix))]
563 #[test]
564 fn mempalace_talks_to_a_fake_server() {
565 let script = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/fake-mempalace-mcp.sh");
566 let eng = MemPalaceEngine::new("sh", vec![script.to_string()]);
567 let hits = eng.try_retrieve(&Query::new("rust", now()), 10).unwrap();
568 assert_eq!(hits.len(), 2);
569 assert_eq!(hits[0].memory.id, 143);
570 assert!(hits[0].reasons()[0].starts_with("mempalace semantic score"));
571 }
572
573 #[cfg(all(feature = "mempalace", unix))]
574 #[test]
575 fn checkpoint_pushes_and_parses_the_summary() {
576 let script =
577 concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/fake-mempalace-checkpoint.sh");
578 let items = vec![
579 Drawer { wing: "memorywhale".into(), room: "command".into(), content: "cargo build failed".into() },
580 Drawer { wing: "memorywhale".into(), room: "note".into(), content: "the fix was X".into() },
581 ];
582 let out = checkpoint("sh", &[script.to_string()], "mempalace_checkpoint", &items).unwrap();
583 assert_eq!(out, CheckpointOutcome { added: 2, duplicates: 1, errors: 0 });
585 }
586
587 #[cfg(all(feature = "mempalace", unix))]
590 #[test]
591 fn sync_ops_adds_and_deletes_over_one_session() {
592 let script =
593 concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/fake-mempalace-sync.sh");
594 let ops = vec![
595 SyncOp::Add {
596 wing: "memorywhale".into(),
597 room: "note".into(),
598 content: "first".into(),
599 added_by: "you".into(),
600 },
601 SyncOp::Delete { drawer_id: "old-1".into() },
602 SyncOp::Add {
603 wing: "memorywhale".into(),
604 room: "note".into(),
605 content: "second".into(),
606 added_by: "memorywhale".into(),
607 },
608 ];
609 let out = sync_ops(
610 "sh",
611 &[script.to_string()],
612 "mempalace_add_drawer",
613 "mempalace_delete_drawer",
614 &ops,
615 )
616 .unwrap();
617 assert_eq!(out.len(), 3);
618 let ids: Vec<&str> = out
619 .iter()
620 .filter_map(|r| match r {
621 SyncResult::Added { drawer_id } => Some(drawer_id.as_str()),
622 SyncResult::Deleted => None,
623 })
624 .collect();
625 assert_eq!(ids, vec!["drawer-1", "drawer-2"]);
626 assert!(matches!(out[1], SyncResult::Deleted));
627 }
628
629 #[cfg(feature = "mempalace")]
630 #[test]
631 fn missing_server_is_a_clear_error() {
632 let eng = MemPalaceEngine::new("mw-definitely-not-installed", vec![]);
633 let err = eng
634 .try_retrieve(&Query::new("x", now()), 3)
635 .unwrap_err()
636 .to_string();
637 assert!(err.contains("failed to start MCP server"), "{err}");
638 }
639}