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;
60const TEMPORAL_TAG_FIRST_MAX: usize = SOURCE_CAP * 64;
64
65const GRAPH_ENTITY_CAP: usize = 64;
67const GRAPH_FACT_CAP: usize = 256;
68const GRAPH_EDGE_CAP: usize = 128;
69const GRAPH_EXAMINE_CAP: usize = 2048;
74
75const STOP_DF_DIVISOR: u64 = 8;
82const STOP_DF_FLOOR: u64 = 1024;
85
86#[derive(Clone, Copy, Debug)]
89#[cfg_attr(feature = "serde", derive(serde::Serialize))]
90pub struct RecallQuery<'a> {
91 pub now: u64,
93 pub text: Option<&'a str>,
95 pub vector: Option<&'a [f32]>,
97 pub tags: &'a [&'a str],
99 pub entities: &'a [&'a str],
101 pub as_of: Option<u64>,
103 pub range: Option<(u64, u64)>,
105 pub k: usize,
107 pub token_budget: Option<usize>,
109 pub include_closed: bool,
111 pub ef: Option<usize>,
115}
116
117impl<'a> RecallQuery<'a> {
118 pub fn text(now: u64, text: &'a str) -> Self {
120 Self {
121 now,
122 text: Some(text),
123 vector: None,
124 tags: &[],
125 entities: &[],
126 as_of: None,
127 range: None,
128 k: 0,
129 token_budget: None,
130 include_closed: false,
131 ef: None,
132 }
133 }
134}
135
136#[derive(Clone, Copy, Debug, PartialEq)]
138#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
139pub struct RecalledFact {
140 pub id: FactId,
142 pub score: f32,
144 pub sources: u8,
146 pub entity: EntityId,
148 pub recorded_at: u64,
150 pub valid_from: u64,
152 pub valid_to: u64,
154}
155
156#[derive(Clone, Copy, Debug, PartialEq, Eq)]
159#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
160pub struct RecalledEdge {
161 pub src: EntityId,
163 pub rel: TermId,
165 pub dst: EntityId,
167 pub provenance: FactId,
169}
170
171#[derive(Clone, Debug, Default)]
174#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
175pub struct RecallResult {
176 pub facts: Vec<RecalledFact>,
178 pub edges: Vec<RecalledEdge>,
180 pub rendered: String,
182 pub truncated: bool,
185}
186
187#[derive(Debug, Default)]
198pub struct RecallScratch {
199 tokenizer: Tokenizer,
202 name_scratch: String,
204 bm25: Bm25Scratch,
205 intersect: IntersectScratch,
206 allow: Vec<FactId>,
207 tag_terms: Vec<u32>,
208 query_terms: Vec<u32>,
209 bm25_out: Vec<(FactId, f32)>,
210 vec: VecScratch,
211 vec_out: Vec<(FactId, f32)>,
212 hnsw: HnswScratch,
213 hnsw_out: Vec<(u32, f32)>,
214 graph_out: Vec<(FactId, f32)>,
215 time_out: Vec<(FactId, f32)>,
216 time_tag: Vec<(FactId, u64)>,
217 visited: Vec<(EntityId, f32)>,
218 edges_tmp: Vec<(EntityId, TermId, bool, FactId)>,
219 fused: hashbrown::HashMap<u32, (f32, u8), xxhash_rust::xxh3::Xxh3Builder>,
220 ranked: Vec<(FactId, f32, u8)>,
221 tags_tmp: Vec<TermId>,
222}
223
224impl RecallScratch {
225 pub fn new() -> Self {
228 Self::default()
229 }
230}
231
232impl Memory<'_> {
233 pub fn recall(&self, q: RecallQuery<'_>) -> Result<RecallResult, Error> {
238 let mut scratch = RecallScratch::default();
239 let mut out = RecallResult::default();
240 self.recall_into(q, &mut scratch, &mut out)?;
241 Ok(out)
242 }
243
244 pub fn recall_into(
251 &self,
252 q: RecallQuery<'_>,
253 s: &mut RecallScratch,
254 out: &mut RecallResult,
255 ) -> Result<(), Error> {
256 out.facts.clear();
257 out.edges.clear();
258 out.rendered.clear();
259 out.truncated = false;
260
261 let k = if q.k == 0 { 8 } else { q.k.min(64) };
262 let budget = q.token_budget.unwrap_or(512);
263 let as_of = q.as_of.unwrap_or(q.now);
264
265 s.allow.clear();
267 s.tag_terms.clear();
268 let mut dead_tag = false;
269 for tag in q.tags {
270 match self.terms.lookup(tag) {
271 Some(term) => s.tag_terms.push(term.0),
272 None => dead_tag = true,
273 }
274 }
275 if !dead_tag && !s.tag_terms.is_empty() {
276 intersect(&self.tags_idx, &s.tag_terms, &mut s.intersect, &mut s.allow);
277 }
278 let filtered = !q.tags.is_empty();
279 if filtered && (dead_tag || s.allow.is_empty()) {
280 return Ok(());
281 }
282
283 s.bm25_out.clear();
285 if let Some(text) = q.text {
286 s.query_terms.clear();
287 let terms = &self.terms;
288 let query_terms = &mut s.query_terms;
291 s.tokenizer.tokenize(text, &mut |token| {
292 if let Some(term) = terms.lookup(token) {
293 query_terms.push(term.0);
294 }
295 });
296 let docs = self.bm25.docs();
298 let is_stop = |df: u64| df > STOP_DF_FLOOR && df * STOP_DF_DIVISOR > docs;
299 if s.query_terms
300 .iter()
301 .any(|&t| !is_stop(u64::from(self.bm25.df(t))))
302 {
303 let bm25 = &self.bm25;
304 s.query_terms.retain(|&t| !is_stop(u64::from(bm25.df(t))));
305 } else if let Some(&least) = s.query_terms.iter().min_by_key(|&&t| self.bm25.df(t)) {
306 s.query_terms.clear();
307 s.query_terms.push(least);
308 }
309 let facts = &self.facts;
310 let allow = &s.allow;
311 self.bm25.search(
312 (self.cfg.bm25_k1, self.cfg.bm25_b),
313 &s.query_terms,
314 SOURCE_CAP,
315 &mut |id| admit(facts, allow, filtered, as_of, q.include_closed, id).is_some(),
316 &mut s.bm25,
317 &mut s.bm25_out,
318 );
319 }
320
321 s.vec_out.clear();
324 if let Some(v) = q.vector
325 && self.cfg.dim > 0
326 {
327 let res = if self.hnsw.indexed() == 0 {
328 let facts = &self.facts;
329 let allow = &s.allow;
330 self.vecs.search(
331 v,
332 SOURCE_CAP,
333 &mut |id| admit(facts, allow, filtered, as_of, q.include_closed, id).is_some(),
334 &mut s.vec,
335 &mut s.vec_out,
336 )
337 } else {
338 self.vec_graph_source(v, &q, as_of, filtered, s)
339 };
340 res?;
341 }
342
343 s.visited.clear();
347 for name in q.entities {
348 super::normalize_name(&mut s.tokenizer, name, &mut s.name_scratch);
349 let found = self.lookup_entity_by_norm(&s.name_scratch);
350 if let Some(id) = found
351 && !s.visited.iter().any(|&(e, _)| e == id)
352 {
353 s.visited.push((id, 1.0));
354 }
355 }
356 self.graph_source(&q, as_of, filtered, s, out);
357 self.time_source(&q, as_of, filtered, s);
358
359 s.fused.clear();
361 for (list, weight, bit) in [
362 (&s.bm25_out, self.cfg.w_bm25, source::BM25),
363 (&s.vec_out, self.cfg.w_vec, source::VEC),
364 (&s.graph_out, self.cfg.w_graph, source::GRAPH),
365 (&s.time_out, self.cfg.w_time, source::TIME),
366 ] {
367 for (rank, &(fact, _)) in list.iter().enumerate() {
368 let contribution = weight / (self.cfg.rrf_k as f32 + rank as f32 + 1.0);
369 let entry = s.fused.entry(fact.0).or_insert((0.0, 0));
370 entry.0 += contribution;
371 entry.1 |= bit;
372 }
373 }
374
375 let half_life_ms = self.cfg.half_life_days as f32 * 86_400_000.0;
377 s.ranked.clear();
378 for (&id, &(score, bits)) in &s.fused {
379 let record = self.facts.get(&id.to_be_bytes()).expect("fused ids exist");
380 let age = q.now.saturating_sub(record.recorded_at) as f32;
381 let boost = 1.0 + self.cfg.w_recency * libm::exp2f(-age / half_life_ms);
382 s.ranked.push((FactId(id), score * boost, bits));
383 }
384 s.ranked
385 .sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
386
387 let mut spent = 0usize;
389 for &(id, score, bits) in &s.ranked {
390 if out.facts.len() == k {
391 out.truncated = true;
392 break;
393 }
394 let record = self
395 .facts
396 .get(&id.0.to_be_bytes())
397 .expect("ranked ids exist");
398 let cost = self.texts.get(record.text).len() / 4 + 8;
399 if spent + cost > budget {
400 out.truncated = true;
401 break;
402 }
403 spent += cost;
404 out.facts.push(RecalledFact {
405 id,
406 score,
407 sources: bits,
408 entity: record.entity,
409 recorded_at: record.recorded_at,
410 valid_from: record.valid_from,
411 valid_to: record.valid_to,
412 });
413 }
414
415 self.render(out, &mut s.tags_tmp);
417 Ok(())
418 }
419
420 fn vec_graph_source(
425 &self,
426 v: &[f32],
427 q: &RecallQuery<'_>,
428 as_of: u64,
429 filtered: bool,
430 s: &mut RecallScratch,
431 ) -> Result<(), Error> {
432 let RecallScratch {
433 vec,
434 hnsw,
435 hnsw_out,
436 vec_out,
437 allow,
438 ..
439 } = s;
440 self.vecs.quantize_query(v, vec)?;
441 let (q_scale, q_q) = self.vecs.quantized(vec);
442 let ef = q.ef.unwrap_or(self.cfg.hnsw_ef_search).max(1);
443 self.hnsw
444 .search_quantized(&self.vecs, (q_scale, q_q), ef, hnsw, hnsw_out);
445 for slot in self.hnsw.indexed()..self.vecs.len() as u32 {
446 let (s_scale, s_q) = self.vecs.quant(slot as usize);
447 hnsw_out.push((slot, q_scale * s_scale * dot_i8(q_q, s_q) as f32));
448 }
449 vec_out.clear();
450 for &(slot, sim) in hnsw_out.iter() {
451 let fact = FactId(self.vecs.slot_fact(slot as usize));
452 if admit(&self.facts, allow, filtered, as_of, q.include_closed, fact).is_some() {
453 vec_out.push((fact, sim));
454 }
455 }
456 vec_out.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
457 vec_out.truncate(SOURCE_CAP);
458 Ok(())
459 }
460
461 fn graph_source(
464 &self,
465 q: &RecallQuery<'_>,
466 as_of: u64,
467 filtered: bool,
468 s: &mut RecallScratch,
469 out: &mut RecallResult,
470 ) {
471 let RecallScratch {
472 allow,
473 graph_out,
474 visited,
475 edges_tmp,
476 ..
477 } = s;
478 graph_out.clear();
479 if visited.is_empty() {
480 return;
481 }
482
483 let mut frontier = 0usize;
485 let mut weight = 1.0f32;
486 for _ in 0..self.cfg.graph_depth {
487 let depth_end = visited.len();
488 weight *= self.cfg.graph_decay;
489 for at in frontier..depth_end {
490 let (entity, _) = visited[at];
491 self.neighbors(entity, edges_tmp);
492 let batch = core::mem::take(edges_tmp);
493 for &(neighbor, rel, this_side_src, provenance) in &batch {
494 let (src, dst) = if this_side_src {
495 (entity, neighbor)
496 } else {
497 (neighbor, entity)
498 };
499 let edge = RecalledEdge {
500 src,
501 rel,
502 dst,
503 provenance,
504 };
505 if out.edges.len() < GRAPH_EDGE_CAP && !out.edges.contains(&edge) {
506 out.edges.push(edge);
507 }
508 if visited.len() < GRAPH_ENTITY_CAP
509 && !visited.iter().any(|&(e, _)| e == neighbor)
510 {
511 visited.push((neighbor, weight));
512 }
513 }
514 *edges_tmp = batch;
515 }
516 frontier = depth_end;
517 }
518
519 let mut examined = 0usize;
523 'entities: for &(entity, weight) in visited.iter() {
524 for (fact, _) in self.entity_facts.entries(entity.0) {
525 examined += 1;
526 if graph_out.len() >= GRAPH_FACT_CAP || examined > GRAPH_EXAMINE_CAP {
527 break 'entities;
528 }
529 if admit(&self.facts, allow, filtered, as_of, q.include_closed, fact).is_some() {
530 graph_out.push((fact, weight));
531 }
532 }
533 }
534 for edge in out.edges.iter() {
535 if graph_out.len() >= GRAPH_FACT_CAP {
536 break;
537 }
538 if let Some(fact) = edge.provenance.some()
539 && !graph_out.iter().any(|&(f, _)| f == fact)
540 && admit(&self.facts, allow, filtered, as_of, q.include_closed, fact).is_some()
541 {
542 graph_out.push((fact, self.cfg.graph_decay));
543 }
544 }
545 graph_out.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
546 graph_out.truncate(SOURCE_CAP);
547 graph_out.dedup_by_key(|&mut (f, _)| f);
548 }
549
550 fn time_source(&self, q: &RecallQuery<'_>, as_of: u64, filtered: bool, s: &mut RecallScratch) {
553 s.time_out.clear();
554 let Some((from, to)) = q.range else { return };
555 if filtered && !s.allow.is_empty() && s.allow.len() <= TEMPORAL_TAG_FIRST_MAX {
556 self.time_source_from_tags(from, to, as_of, q.include_closed, s);
557 return;
558 }
559 let RecallScratch {
560 allow, time_out, ..
561 } = s;
562 let mut from_key = [0u8; 12];
563 plugmem_arena::key::write_pair(&mut from_key, from, 0);
564 let mut to_key = [0u8; 12];
565 plugmem_arena::key::write_pair(&mut to_key, to, 0);
566 for slot in self.temporal.range_rev(&from_key, &to_key) {
567 if admit(
568 &self.facts,
569 allow,
570 filtered,
571 as_of,
572 q.include_closed,
573 slot.fact,
574 )
575 .is_some()
576 {
577 time_out.push((slot.fact, slot.recorded_at as f32));
578 if time_out.len() == SOURCE_CAP {
582 break;
583 }
584 }
585 }
586 }
587
588 fn time_source_from_tags(
592 &self,
593 from: u64,
594 to: u64,
595 as_of: u64,
596 include_closed: bool,
597 s: &mut RecallScratch,
598 ) {
599 let RecallScratch {
600 allow,
601 time_tag,
602 time_out,
603 ..
604 } = s;
605 time_tag.clear();
606 for &fact in allow.iter() {
607 let Some(record) = admit(&self.facts, &[], false, as_of, include_closed, fact) else {
608 continue;
609 };
610 if record.recorded_at >= from && record.recorded_at < to {
611 time_tag.push((fact, record.recorded_at));
612 }
613 }
614 time_tag.sort_unstable_by(|a, b| b.1.cmp(&a.1).then_with(|| b.0.cmp(&a.0)));
615 time_out.extend(
616 time_tag
617 .iter()
618 .take(SOURCE_CAP)
619 .map(|&(fact, recorded_at)| (fact, recorded_at as f32)),
620 );
621 }
622
623 fn neighbors(&self, entity: EntityId, out: &mut Vec<(EntityId, TermId, bool, FactId)>) {
626 out.clear();
627 let mut from = [0u8; 12];
628 plugmem_arena::key::write_u32(&mut from, entity.0);
629 let mut to = [0u8; 12];
630 plugmem_arena::key::write_u32(&mut to, entity.0 + 1);
631 out.extend(
632 self.edges_out
633 .range(&from, &to)
634 .map(|e| (e.b, e.rel, true, e.fact)),
635 );
636 out.extend(
637 self.edges_in
638 .range(&from, &to)
639 .map(|e| (e.b, e.rel, false, e.fact)),
640 );
641 }
642
643 fn render(&self, out: &mut RecallResult, tags_tmp: &mut Vec<TermId>) {
645 if out.facts.is_empty() && out.edges.is_empty() {
646 return; }
648 out.rendered.push_str("## memory\n");
649 for fact in &out.facts {
650 let record = self
651 .facts
652 .get(&fact.id.0.to_be_bytes())
653 .expect("selected ids exist");
654 let text = core::str::from_utf8(self.texts.get(record.text)).unwrap_or("");
657 let _ = write!(out.rendered, "- [f{}] ", fact.id.0);
658 if let Some(entity) = fact.entity.some()
661 && let Some(name) = self.entity_name(entity)
662 {
663 let _ = write!(out.rendered, "{name}: ");
664 }
665 out.rendered.push_str(text);
666 out.rendered.push_str(" (");
667 render_ym(&mut out.rendered, fact.valid_from);
668 if fact.valid_to == VALID_TO_OPEN {
669 out.rendered.push_str("; active)");
670 } else {
671 out.rendered.push_str(" → ");
672 render_ym(&mut out.rendered, fact.valid_to);
673 out.rendered.push_str("; closed)");
674 }
675 tags_tmp.clear();
676 self.tags_of(fact.id, tags_tmp);
677 for &tag in tags_tmp.iter() {
678 let _ = write!(out.rendered, " #{}", self.terms.resolve(tag));
679 }
680 out.rendered.push('\n');
681 }
682 for edge in &out.edges {
683 let _ = writeln!(
684 out.rendered,
685 "- links: {} —{}→ {}",
686 self.entity_name(edge.src)
687 .expect("edges reference existing entities"),
688 self.terms.resolve(edge.rel),
689 self.entity_name(edge.dst)
690 .expect("edges reference existing entities"),
691 );
692 }
693 }
694}
695
696fn admit(
699 facts: &plugmem_arena::Arena<'_, FactRecord>,
700 allow: &[FactId],
701 filtered: bool,
702 as_of: u64,
703 include_closed: bool,
704 id: FactId,
705) -> Option<FactRecord> {
706 let record = facts.get(&id.0.to_be_bytes())?;
707 if record.is_tombstone() || record.recorded_at > as_of || record.valid_from > as_of {
708 return None;
709 }
710 if !include_closed && as_of >= record.valid_to {
711 return None;
712 }
713 if filtered && allow.binary_search(&id).is_err() {
714 return None;
715 }
716 Some(record)
717}
718
719fn render_ym(out: &mut String, ms: u64) {
722 let days = (ms / 86_400_000) as i64;
723 let z = days + 719_468;
724 let era = z.div_euclid(146_097);
725 let doe = z.rem_euclid(146_097);
726 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
727 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
728 let mp = (5 * doy + 2) / 153;
729 let month = if mp < 10 { mp + 3 } else { mp - 9 };
730 let year = yoe + era * 400 + i64::from(month <= 2);
731 let _ = write!(out, "{year:04}-{month:02}");
732}