1use alloc::string::String;
30use alloc::vec::Vec;
31use core::fmt::Write as _;
32
33use plugmem_arena::TermId;
34
35use crate::error::Error;
36use crate::id::{EntityId, FactId};
37use crate::index::bm25::Bm25Scratch;
38use crate::index::hnsw::HnswScratch;
39use crate::index::vecpool::{VecScratch, dot_i8};
40use crate::index::{IntersectScratch, intersect};
41use crate::model::{FactRecord, VALID_TO_OPEN};
42use crate::tokenizer::Tokenizer;
43
44use super::Memory;
45
46pub mod source {
48 pub const BM25: u8 = 1;
50 pub const GRAPH: u8 = 1 << 1;
52 pub const TIME: u8 = 1 << 2;
54 pub const VEC: u8 = 1 << 3;
56}
57
58const SOURCE_CAP: usize = 128;
60
61const GRAPH_ENTITY_CAP: usize = 64;
63const GRAPH_FACT_CAP: usize = 256;
64const GRAPH_EDGE_CAP: usize = 128;
65const GRAPH_EXAMINE_CAP: usize = 2048;
70
71const STOP_DF_DIVISOR: u64 = 8;
78const STOP_DF_FLOOR: u64 = 1024;
81
82#[derive(Clone, Copy, Debug)]
85#[cfg_attr(feature = "serde", derive(serde::Serialize))]
86pub struct RecallQuery<'a> {
87 pub now: u64,
89 pub text: Option<&'a str>,
91 pub vector: Option<&'a [f32]>,
93 pub tags: &'a [&'a str],
95 pub entities: &'a [&'a str],
97 pub as_of: Option<u64>,
99 pub range: Option<(u64, u64)>,
101 pub k: usize,
103 pub token_budget: Option<usize>,
105 pub include_closed: bool,
107 pub ef: Option<usize>,
111}
112
113impl<'a> RecallQuery<'a> {
114 pub fn text(now: u64, text: &'a str) -> Self {
116 Self {
117 now,
118 text: Some(text),
119 vector: None,
120 tags: &[],
121 entities: &[],
122 as_of: None,
123 range: None,
124 k: 0,
125 token_budget: None,
126 include_closed: false,
127 ef: None,
128 }
129 }
130}
131
132#[derive(Clone, Copy, Debug, PartialEq)]
134#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
135pub struct RecalledFact {
136 pub id: FactId,
138 pub score: f32,
140 pub sources: u8,
142 pub entity: EntityId,
144 pub recorded_at: u64,
146 pub valid_from: u64,
148 pub valid_to: u64,
150}
151
152#[derive(Clone, Copy, Debug, PartialEq, Eq)]
155#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
156pub struct RecalledEdge {
157 pub src: EntityId,
159 pub rel: TermId,
161 pub dst: EntityId,
163 pub provenance: FactId,
165}
166
167#[derive(Clone, Debug, Default)]
170#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
171pub struct RecallResult {
172 pub facts: Vec<RecalledFact>,
174 pub edges: Vec<RecalledEdge>,
176 pub rendered: String,
178 pub truncated: bool,
181}
182
183#[derive(Debug, Default)]
194pub struct RecallScratch {
195 tokenizer: Tokenizer,
198 name_scratch: String,
200 bm25: Bm25Scratch,
201 intersect: IntersectScratch,
202 allow: Vec<FactId>,
203 tag_terms: Vec<u32>,
204 query_terms: Vec<u32>,
205 bm25_out: Vec<(FactId, f32)>,
206 vec: VecScratch,
207 vec_out: Vec<(FactId, f32)>,
208 hnsw: HnswScratch,
209 hnsw_out: Vec<(u32, f32)>,
210 graph_out: Vec<(FactId, f32)>,
211 time_out: Vec<(FactId, f32)>,
212 visited: Vec<(EntityId, f32)>,
213 edges_tmp: Vec<(EntityId, TermId, bool, FactId)>,
214 fused: hashbrown::HashMap<u32, (f32, u8), xxhash_rust::xxh3::Xxh3Builder>,
215 ranked: Vec<(FactId, f32, u8)>,
216 tags_tmp: Vec<TermId>,
217}
218
219impl RecallScratch {
220 pub fn new() -> Self {
223 Self::default()
224 }
225}
226
227impl Memory<'_> {
228 pub fn recall(&self, q: RecallQuery<'_>) -> Result<RecallResult, Error> {
233 let mut scratch = RecallScratch::default();
234 let mut out = RecallResult::default();
235 self.recall_into(q, &mut scratch, &mut out)?;
236 Ok(out)
237 }
238
239 pub fn recall_into(
246 &self,
247 q: RecallQuery<'_>,
248 s: &mut RecallScratch,
249 out: &mut RecallResult,
250 ) -> Result<(), Error> {
251 out.facts.clear();
252 out.edges.clear();
253 out.rendered.clear();
254 out.truncated = false;
255
256 let k = if q.k == 0 { 8 } else { q.k.min(64) };
257 let budget = q.token_budget.unwrap_or(512);
258 let as_of = q.as_of.unwrap_or(q.now);
259
260 s.allow.clear();
262 s.tag_terms.clear();
263 let mut dead_tag = false;
264 for tag in q.tags {
265 match self.terms.lookup(tag) {
266 Some(term) => s.tag_terms.push(term.0),
267 None => dead_tag = true,
268 }
269 }
270 if !dead_tag && !s.tag_terms.is_empty() {
271 intersect(&self.tags_idx, &s.tag_terms, &mut s.intersect, &mut s.allow);
272 }
273 let filtered = !q.tags.is_empty();
274 if filtered && (dead_tag || s.allow.is_empty()) {
275 return Ok(());
276 }
277
278 s.bm25_out.clear();
280 if let Some(text) = q.text {
281 s.query_terms.clear();
282 let terms = &self.terms;
283 let query_terms = &mut s.query_terms;
286 s.tokenizer.tokenize(text, &mut |token| {
287 if let Some(term) = terms.lookup(token) {
288 query_terms.push(term.0);
289 }
290 });
291 let docs = self.bm25.docs();
293 let is_stop = |df: u64| df > STOP_DF_FLOOR && df * STOP_DF_DIVISOR > docs;
294 if s.query_terms
295 .iter()
296 .any(|&t| !is_stop(u64::from(self.bm25.df(t))))
297 {
298 let bm25 = &self.bm25;
299 s.query_terms.retain(|&t| !is_stop(u64::from(bm25.df(t))));
300 } else if let Some(&least) = s.query_terms.iter().min_by_key(|&&t| self.bm25.df(t)) {
301 s.query_terms.clear();
302 s.query_terms.push(least);
303 }
304 let facts = &self.facts;
305 let allow = &s.allow;
306 self.bm25.search(
307 (self.cfg.bm25_k1, self.cfg.bm25_b),
308 &s.query_terms,
309 SOURCE_CAP,
310 &mut |id| admit(facts, allow, filtered, as_of, q.include_closed, id).is_some(),
311 &mut s.bm25,
312 &mut s.bm25_out,
313 );
314 }
315
316 s.vec_out.clear();
319 if let Some(v) = q.vector
320 && self.cfg.dim > 0
321 {
322 let res = if self.hnsw.indexed() == 0 {
323 let facts = &self.facts;
324 let allow = &s.allow;
325 self.vecs.search(
326 v,
327 SOURCE_CAP,
328 &mut |id| admit(facts, allow, filtered, as_of, q.include_closed, id).is_some(),
329 &mut s.vec,
330 &mut s.vec_out,
331 )
332 } else {
333 self.vec_graph_source(v, &q, as_of, filtered, s)
334 };
335 res?;
336 }
337
338 s.visited.clear();
342 for name in q.entities {
343 super::normalize_name(&mut s.tokenizer, name, &mut s.name_scratch);
344 let found = self.lookup_entity_by_norm(&s.name_scratch);
345 if let Some(id) = found
346 && !s.visited.iter().any(|&(e, _)| e == id)
347 {
348 s.visited.push((id, 1.0));
349 }
350 }
351 self.graph_source(&q, as_of, filtered, s, out);
352 self.time_source(&q, as_of, filtered, s);
353
354 s.fused.clear();
356 for (list, weight, bit) in [
357 (&s.bm25_out, self.cfg.w_bm25, source::BM25),
358 (&s.vec_out, self.cfg.w_vec, source::VEC),
359 (&s.graph_out, self.cfg.w_graph, source::GRAPH),
360 (&s.time_out, self.cfg.w_time, source::TIME),
361 ] {
362 for (rank, &(fact, _)) in list.iter().enumerate() {
363 let contribution = weight / (self.cfg.rrf_k as f32 + rank as f32 + 1.0);
364 let entry = s.fused.entry(fact.0).or_insert((0.0, 0));
365 entry.0 += contribution;
366 entry.1 |= bit;
367 }
368 }
369
370 let half_life_ms = self.cfg.half_life_days as f32 * 86_400_000.0;
372 s.ranked.clear();
373 for (&id, &(score, bits)) in &s.fused {
374 let record = self.facts.get(&id.to_be_bytes()).expect("fused ids exist");
375 let age = q.now.saturating_sub(record.recorded_at) as f32;
376 let boost = 1.0 + self.cfg.w_recency * libm::exp2f(-age / half_life_ms);
377 s.ranked.push((FactId(id), score * boost, bits));
378 }
379 s.ranked
380 .sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
381
382 let mut spent = 0usize;
384 for &(id, score, bits) in &s.ranked {
385 if out.facts.len() == k {
386 out.truncated = true;
387 break;
388 }
389 let record = self
390 .facts
391 .get(&id.0.to_be_bytes())
392 .expect("ranked ids exist");
393 let cost = self.texts.get(record.text).len() / 4 + 8;
394 if spent + cost > budget {
395 out.truncated = true;
396 break;
397 }
398 spent += cost;
399 out.facts.push(RecalledFact {
400 id,
401 score,
402 sources: bits,
403 entity: record.entity,
404 recorded_at: record.recorded_at,
405 valid_from: record.valid_from,
406 valid_to: record.valid_to,
407 });
408 }
409
410 self.render(out, &mut s.tags_tmp);
412 Ok(())
413 }
414
415 fn vec_graph_source(
420 &self,
421 v: &[f32],
422 q: &RecallQuery<'_>,
423 as_of: u64,
424 filtered: bool,
425 s: &mut RecallScratch,
426 ) -> Result<(), Error> {
427 let RecallScratch {
428 vec,
429 hnsw,
430 hnsw_out,
431 vec_out,
432 allow,
433 ..
434 } = s;
435 self.vecs.quantize_query(v, vec)?;
436 let (q_scale, q_q) = self.vecs.quantized(vec);
437 let ef = q.ef.unwrap_or(self.cfg.hnsw_ef_search).max(1);
438 self.hnsw
439 .search_quantized(&self.vecs, (q_scale, q_q), ef, hnsw, hnsw_out);
440 for slot in self.hnsw.indexed()..self.vecs.len() as u32 {
441 let (s_scale, s_q) = self.vecs.quant(slot as usize);
442 hnsw_out.push((slot, q_scale * s_scale * dot_i8(q_q, s_q) as f32));
443 }
444 vec_out.clear();
445 for &(slot, sim) in hnsw_out.iter() {
446 let fact = FactId(self.vecs.slot_fact(slot as usize));
447 if admit(&self.facts, allow, filtered, as_of, q.include_closed, fact).is_some() {
448 vec_out.push((fact, sim));
449 }
450 }
451 vec_out.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
452 vec_out.truncate(SOURCE_CAP);
453 Ok(())
454 }
455
456 fn graph_source(
459 &self,
460 q: &RecallQuery<'_>,
461 as_of: u64,
462 filtered: bool,
463 s: &mut RecallScratch,
464 out: &mut RecallResult,
465 ) {
466 let RecallScratch {
467 allow,
468 graph_out,
469 visited,
470 edges_tmp,
471 ..
472 } = s;
473 graph_out.clear();
474 if visited.is_empty() {
475 return;
476 }
477
478 let mut frontier = 0usize;
480 let mut weight = 1.0f32;
481 for _ in 0..self.cfg.graph_depth {
482 let depth_end = visited.len();
483 weight *= self.cfg.graph_decay;
484 for at in frontier..depth_end {
485 let (entity, _) = visited[at];
486 self.neighbors(entity, edges_tmp);
487 let batch = core::mem::take(edges_tmp);
488 for &(neighbor, rel, this_side_src, provenance) in &batch {
489 let (src, dst) = if this_side_src {
490 (entity, neighbor)
491 } else {
492 (neighbor, entity)
493 };
494 let edge = RecalledEdge {
495 src,
496 rel,
497 dst,
498 provenance,
499 };
500 if out.edges.len() < GRAPH_EDGE_CAP && !out.edges.contains(&edge) {
501 out.edges.push(edge);
502 }
503 if visited.len() < GRAPH_ENTITY_CAP
504 && !visited.iter().any(|&(e, _)| e == neighbor)
505 {
506 visited.push((neighbor, weight));
507 }
508 }
509 *edges_tmp = batch;
510 }
511 frontier = depth_end;
512 }
513
514 let mut examined = 0usize;
518 'entities: for &(entity, weight) in visited.iter() {
519 for (fact, _) in self.entity_facts.entries(entity.0) {
520 examined += 1;
521 if graph_out.len() >= GRAPH_FACT_CAP || examined > GRAPH_EXAMINE_CAP {
522 break 'entities;
523 }
524 if admit(&self.facts, allow, filtered, as_of, q.include_closed, fact).is_some() {
525 graph_out.push((fact, weight));
526 }
527 }
528 }
529 for edge in out.edges.iter() {
530 if graph_out.len() >= GRAPH_FACT_CAP {
531 break;
532 }
533 if let Some(fact) = edge.provenance.some()
534 && !graph_out.iter().any(|&(f, _)| f == fact)
535 && admit(&self.facts, allow, filtered, as_of, q.include_closed, fact).is_some()
536 {
537 graph_out.push((fact, self.cfg.graph_decay));
538 }
539 }
540 graph_out.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
541 graph_out.truncate(SOURCE_CAP);
542 graph_out.dedup_by_key(|&mut (f, _)| f);
543 }
544
545 fn time_source(&self, q: &RecallQuery<'_>, as_of: u64, filtered: bool, s: &mut RecallScratch) {
548 let RecallScratch {
549 allow, time_out, ..
550 } = s;
551 time_out.clear();
552 let Some((from, to)) = q.range else { return };
553 let mut from_key = [0u8; 12];
554 plugmem_arena::key::write_pair(&mut from_key, from, 0);
555 let mut to_key = [0u8; 12];
556 plugmem_arena::key::write_pair(&mut to_key, to, 0);
557 for slot in self.temporal.range(&from_key, &to_key) {
558 if admit(
559 &self.facts,
560 allow,
561 filtered,
562 as_of,
563 q.include_closed,
564 slot.fact,
565 )
566 .is_some()
567 {
568 time_out.push((slot.fact, slot.recorded_at as f32));
569 if time_out.len() > SOURCE_CAP * 2 {
572 time_out.drain(..SOURCE_CAP);
573 }
574 }
575 }
576 time_out.reverse();
577 time_out.truncate(SOURCE_CAP);
578 }
579
580 fn neighbors(&self, entity: EntityId, out: &mut Vec<(EntityId, TermId, bool, FactId)>) {
583 out.clear();
584 let mut from = [0u8; 12];
585 plugmem_arena::key::write_u32(&mut from, entity.0);
586 let mut to = [0u8; 12];
587 plugmem_arena::key::write_u32(&mut to, entity.0 + 1);
588 out.extend(
589 self.edges_out
590 .range(&from, &to)
591 .map(|e| (e.b, e.rel, true, e.fact)),
592 );
593 out.extend(
594 self.edges_in
595 .range(&from, &to)
596 .map(|e| (e.b, e.rel, false, e.fact)),
597 );
598 }
599
600 fn render(&self, out: &mut RecallResult, tags_tmp: &mut Vec<TermId>) {
602 if out.facts.is_empty() && out.edges.is_empty() {
603 return; }
605 out.rendered.push_str("## memory\n");
606 for fact in &out.facts {
607 let record = self
608 .facts
609 .get(&fact.id.0.to_be_bytes())
610 .expect("selected ids exist");
611 let text = core::str::from_utf8(self.texts.get(record.text)).unwrap_or("");
614 let _ = write!(out.rendered, "- [f{}] ", fact.id.0);
615 if let Some(entity) = fact.entity.some()
618 && let Some(name) = self.entity_name(entity)
619 {
620 let _ = write!(out.rendered, "{name}: ");
621 }
622 out.rendered.push_str(text);
623 out.rendered.push_str(" (");
624 render_ym(&mut out.rendered, fact.valid_from);
625 if fact.valid_to == VALID_TO_OPEN {
626 out.rendered.push_str("; active)");
627 } else {
628 out.rendered.push_str(" → ");
629 render_ym(&mut out.rendered, fact.valid_to);
630 out.rendered.push_str("; closed)");
631 }
632 tags_tmp.clear();
633 self.tags_of(fact.id, tags_tmp);
634 for &tag in tags_tmp.iter() {
635 let _ = write!(out.rendered, " #{}", self.terms.resolve(tag));
636 }
637 out.rendered.push('\n');
638 }
639 for edge in &out.edges {
640 let _ = writeln!(
641 out.rendered,
642 "- links: {} —{}→ {}",
643 self.entity_name(edge.src)
644 .expect("edges reference existing entities"),
645 self.terms.resolve(edge.rel),
646 self.entity_name(edge.dst)
647 .expect("edges reference existing entities"),
648 );
649 }
650 }
651}
652
653fn admit(
656 facts: &plugmem_arena::Arena<'_, FactRecord>,
657 allow: &[FactId],
658 filtered: bool,
659 as_of: u64,
660 include_closed: bool,
661 id: FactId,
662) -> Option<FactRecord> {
663 let record = facts.get(&id.0.to_be_bytes())?;
664 if record.is_tombstone() || record.recorded_at > as_of || record.valid_from > as_of {
665 return None;
666 }
667 if !include_closed && as_of >= record.valid_to {
668 return None;
669 }
670 if filtered && allow.binary_search(&id).is_err() {
671 return None;
672 }
673 Some(record)
674}
675
676fn render_ym(out: &mut String, ms: u64) {
679 let days = (ms / 86_400_000) as i64;
680 let z = days + 719_468;
681 let era = z.div_euclid(146_097);
682 let doe = z.rem_euclid(146_097);
683 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
684 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
685 let mp = (5 * doy + 2) / 153;
686 let month = if mp < 10 { mp + 3 } else { mp - 9 };
687 let year = yoe + era * 400 + i64::from(month <= 2);
688 let _ = write!(out, "{year:04}-{month:02}");
689}