use alloc::string::String;
use alloc::vec::Vec;
use core::fmt::Write as _;
use plugmem_arena::TermId;
use crate::error::Error;
use crate::id::{EntityId, FactId};
use crate::index::bm25::Bm25Scratch;
use crate::index::hnsw::HnswScratch;
use crate::index::vecpool::{VecScratch, dot_i8};
use crate::index::{IntersectScratch, intersect};
use crate::model::{FactRecord, VALID_TO_OPEN};
use crate::tokenizer::Tokenizer;
use super::Memory;
pub mod source {
pub const BM25: u8 = 1;
pub const GRAPH: u8 = 1 << 1;
pub const TIME: u8 = 1 << 2;
pub const VEC: u8 = 1 << 3;
}
const SOURCE_CAP: usize = 128;
const GRAPH_ENTITY_CAP: usize = 64;
const GRAPH_FACT_CAP: usize = 256;
const GRAPH_EDGE_CAP: usize = 128;
const GRAPH_EXAMINE_CAP: usize = 2048;
const STOP_DF_DIVISOR: u64 = 8;
const STOP_DF_FLOOR: u64 = 1024;
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct RecallQuery<'a> {
pub now: u64,
pub text: Option<&'a str>,
pub vector: Option<&'a [f32]>,
pub tags: &'a [&'a str],
pub entities: &'a [&'a str],
pub as_of: Option<u64>,
pub range: Option<(u64, u64)>,
pub k: usize,
pub token_budget: Option<usize>,
pub include_closed: bool,
pub ef: Option<usize>,
}
impl<'a> RecallQuery<'a> {
pub fn text(now: u64, text: &'a str) -> Self {
Self {
now,
text: Some(text),
vector: None,
tags: &[],
entities: &[],
as_of: None,
range: None,
k: 0,
token_budget: None,
include_closed: false,
ef: None,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct RecalledFact {
pub id: FactId,
pub score: f32,
pub sources: u8,
pub entity: EntityId,
pub recorded_at: u64,
pub valid_from: u64,
pub valid_to: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct RecalledEdge {
pub src: EntityId,
pub rel: TermId,
pub dst: EntityId,
pub provenance: FactId,
}
#[derive(Clone, Debug, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct RecallResult {
pub facts: Vec<RecalledFact>,
pub edges: Vec<RecalledEdge>,
pub rendered: String,
pub truncated: bool,
}
#[derive(Debug, Default)]
pub struct RecallScratch {
tokenizer: Tokenizer,
name_scratch: String,
bm25: Bm25Scratch,
intersect: IntersectScratch,
allow: Vec<FactId>,
tag_terms: Vec<u32>,
query_terms: Vec<u32>,
bm25_out: Vec<(FactId, f32)>,
vec: VecScratch,
vec_out: Vec<(FactId, f32)>,
hnsw: HnswScratch,
hnsw_out: Vec<(u32, f32)>,
graph_out: Vec<(FactId, f32)>,
time_out: Vec<(FactId, f32)>,
visited: Vec<(EntityId, f32)>,
edges_tmp: Vec<(EntityId, TermId, bool, FactId)>,
fused: hashbrown::HashMap<u32, (f32, u8), xxhash_rust::xxh3::Xxh3Builder>,
ranked: Vec<(FactId, f32, u8)>,
tags_tmp: Vec<TermId>,
}
impl RecallScratch {
pub fn new() -> Self {
Self::default()
}
}
impl Memory<'_> {
pub fn recall(&self, q: RecallQuery<'_>) -> Result<RecallResult, Error> {
let mut scratch = RecallScratch::default();
let mut out = RecallResult::default();
self.recall_into(q, &mut scratch, &mut out)?;
Ok(out)
}
pub fn recall_into(
&self,
q: RecallQuery<'_>,
s: &mut RecallScratch,
out: &mut RecallResult,
) -> Result<(), Error> {
out.facts.clear();
out.edges.clear();
out.rendered.clear();
out.truncated = false;
let k = if q.k == 0 { 8 } else { q.k.min(64) };
let budget = q.token_budget.unwrap_or(512);
let as_of = q.as_of.unwrap_or(q.now);
s.allow.clear();
s.tag_terms.clear();
let mut dead_tag = false;
for tag in q.tags {
match self.terms.lookup(tag) {
Some(term) => s.tag_terms.push(term.0),
None => dead_tag = true,
}
}
if !dead_tag && !s.tag_terms.is_empty() {
intersect(&self.tags_idx, &s.tag_terms, &mut s.intersect, &mut s.allow);
}
let filtered = !q.tags.is_empty();
if filtered && (dead_tag || s.allow.is_empty()) {
return Ok(());
}
s.bm25_out.clear();
if let Some(text) = q.text {
s.query_terms.clear();
let terms = &self.terms;
let query_terms = &mut s.query_terms;
s.tokenizer.tokenize(text, &mut |token| {
if let Some(term) = terms.lookup(token) {
query_terms.push(term.0);
}
});
let docs = self.bm25.docs();
let is_stop = |df: u64| df > STOP_DF_FLOOR && df * STOP_DF_DIVISOR > docs;
if s.query_terms
.iter()
.any(|&t| !is_stop(u64::from(self.bm25.df(t))))
{
let bm25 = &self.bm25;
s.query_terms.retain(|&t| !is_stop(u64::from(bm25.df(t))));
} else if let Some(&least) = s.query_terms.iter().min_by_key(|&&t| self.bm25.df(t)) {
s.query_terms.clear();
s.query_terms.push(least);
}
let facts = &self.facts;
let allow = &s.allow;
self.bm25.search(
(self.cfg.bm25_k1, self.cfg.bm25_b),
&s.query_terms,
SOURCE_CAP,
&mut |id| admit(facts, allow, filtered, as_of, q.include_closed, id).is_some(),
&mut s.bm25,
&mut s.bm25_out,
);
}
s.vec_out.clear();
if let Some(v) = q.vector
&& self.cfg.dim > 0
{
let res = if self.hnsw.indexed() == 0 {
let facts = &self.facts;
let allow = &s.allow;
self.vecs.search(
v,
SOURCE_CAP,
&mut |id| admit(facts, allow, filtered, as_of, q.include_closed, id).is_some(),
&mut s.vec,
&mut s.vec_out,
)
} else {
self.vec_graph_source(v, &q, as_of, filtered, s)
};
res?;
}
s.visited.clear();
for name in q.entities {
super::normalize_name(&mut s.tokenizer, name, &mut s.name_scratch);
let found = self.lookup_entity_by_norm(&s.name_scratch);
if let Some(id) = found
&& !s.visited.iter().any(|&(e, _)| e == id)
{
s.visited.push((id, 1.0));
}
}
self.graph_source(&q, as_of, filtered, s, out);
self.time_source(&q, as_of, filtered, s);
s.fused.clear();
for (list, weight, bit) in [
(&s.bm25_out, self.cfg.w_bm25, source::BM25),
(&s.vec_out, self.cfg.w_vec, source::VEC),
(&s.graph_out, self.cfg.w_graph, source::GRAPH),
(&s.time_out, self.cfg.w_time, source::TIME),
] {
for (rank, &(fact, _)) in list.iter().enumerate() {
let contribution = weight / (self.cfg.rrf_k as f32 + rank as f32 + 1.0);
let entry = s.fused.entry(fact.0).or_insert((0.0, 0));
entry.0 += contribution;
entry.1 |= bit;
}
}
let half_life_ms = self.cfg.half_life_days as f32 * 86_400_000.0;
s.ranked.clear();
for (&id, &(score, bits)) in &s.fused {
let record = self.facts.get(&id.to_be_bytes()).expect("fused ids exist");
let age = q.now.saturating_sub(record.recorded_at) as f32;
let boost = 1.0 + self.cfg.w_recency * libm::exp2f(-age / half_life_ms);
s.ranked.push((FactId(id), score * boost, bits));
}
s.ranked
.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
let mut spent = 0usize;
for &(id, score, bits) in &s.ranked {
if out.facts.len() == k {
out.truncated = true;
break;
}
let record = self
.facts
.get(&id.0.to_be_bytes())
.expect("ranked ids exist");
let cost = self.texts.get(record.text).len() / 4 + 8;
if spent + cost > budget {
out.truncated = true;
break;
}
spent += cost;
out.facts.push(RecalledFact {
id,
score,
sources: bits,
entity: record.entity,
recorded_at: record.recorded_at,
valid_from: record.valid_from,
valid_to: record.valid_to,
});
}
self.render(out, &mut s.tags_tmp);
Ok(())
}
fn vec_graph_source(
&self,
v: &[f32],
q: &RecallQuery<'_>,
as_of: u64,
filtered: bool,
s: &mut RecallScratch,
) -> Result<(), Error> {
let RecallScratch {
vec,
hnsw,
hnsw_out,
vec_out,
allow,
..
} = s;
self.vecs.quantize_query(v, vec)?;
let (q_scale, q_q) = self.vecs.quantized(vec);
let ef = q.ef.unwrap_or(self.cfg.hnsw_ef_search).max(1);
self.hnsw
.search_quantized(&self.vecs, (q_scale, q_q), ef, hnsw, hnsw_out);
for slot in self.hnsw.indexed()..self.vecs.len() as u32 {
let (s_scale, s_q) = self.vecs.quant(slot as usize);
hnsw_out.push((slot, q_scale * s_scale * dot_i8(q_q, s_q) as f32));
}
vec_out.clear();
for &(slot, sim) in hnsw_out.iter() {
let fact = FactId(self.vecs.slot_fact(slot as usize));
if admit(&self.facts, allow, filtered, as_of, q.include_closed, fact).is_some() {
vec_out.push((fact, sim));
}
}
vec_out.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
vec_out.truncate(SOURCE_CAP);
Ok(())
}
fn graph_source(
&self,
q: &RecallQuery<'_>,
as_of: u64,
filtered: bool,
s: &mut RecallScratch,
out: &mut RecallResult,
) {
let RecallScratch {
allow,
graph_out,
visited,
edges_tmp,
..
} = s;
graph_out.clear();
if visited.is_empty() {
return;
}
let mut frontier = 0usize;
let mut weight = 1.0f32;
for _ in 0..self.cfg.graph_depth {
let depth_end = visited.len();
weight *= self.cfg.graph_decay;
for at in frontier..depth_end {
let (entity, _) = visited[at];
self.neighbors(entity, edges_tmp);
let batch = core::mem::take(edges_tmp);
for &(neighbor, rel, this_side_src, provenance) in &batch {
let (src, dst) = if this_side_src {
(entity, neighbor)
} else {
(neighbor, entity)
};
let edge = RecalledEdge {
src,
rel,
dst,
provenance,
};
if out.edges.len() < GRAPH_EDGE_CAP && !out.edges.contains(&edge) {
out.edges.push(edge);
}
if visited.len() < GRAPH_ENTITY_CAP
&& !visited.iter().any(|&(e, _)| e == neighbor)
{
visited.push((neighbor, weight));
}
}
*edges_tmp = batch;
}
frontier = depth_end;
}
let mut examined = 0usize;
'entities: for &(entity, weight) in visited.iter() {
for (fact, _) in self.entity_facts.entries(entity.0) {
examined += 1;
if graph_out.len() >= GRAPH_FACT_CAP || examined > GRAPH_EXAMINE_CAP {
break 'entities;
}
if admit(&self.facts, allow, filtered, as_of, q.include_closed, fact).is_some() {
graph_out.push((fact, weight));
}
}
}
for edge in out.edges.iter() {
if graph_out.len() >= GRAPH_FACT_CAP {
break;
}
if let Some(fact) = edge.provenance.some()
&& !graph_out.iter().any(|&(f, _)| f == fact)
&& admit(&self.facts, allow, filtered, as_of, q.include_closed, fact).is_some()
{
graph_out.push((fact, self.cfg.graph_decay));
}
}
graph_out.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
graph_out.truncate(SOURCE_CAP);
graph_out.dedup_by_key(|&mut (f, _)| f);
}
fn time_source(&self, q: &RecallQuery<'_>, as_of: u64, filtered: bool, s: &mut RecallScratch) {
let RecallScratch {
allow, time_out, ..
} = s;
time_out.clear();
let Some((from, to)) = q.range else { return };
let mut from_key = [0u8; 12];
plugmem_arena::key::write_pair(&mut from_key, from, 0);
let mut to_key = [0u8; 12];
plugmem_arena::key::write_pair(&mut to_key, to, 0);
for slot in self.temporal.range(&from_key, &to_key) {
if admit(
&self.facts,
allow,
filtered,
as_of,
q.include_closed,
slot.fact,
)
.is_some()
{
time_out.push((slot.fact, slot.recorded_at as f32));
if time_out.len() > SOURCE_CAP * 2 {
time_out.drain(..SOURCE_CAP);
}
}
}
time_out.reverse();
time_out.truncate(SOURCE_CAP);
}
fn neighbors(&self, entity: EntityId, out: &mut Vec<(EntityId, TermId, bool, FactId)>) {
out.clear();
let mut from = [0u8; 12];
plugmem_arena::key::write_u32(&mut from, entity.0);
let mut to = [0u8; 12];
plugmem_arena::key::write_u32(&mut to, entity.0 + 1);
out.extend(
self.edges_out
.range(&from, &to)
.map(|e| (e.b, e.rel, true, e.fact)),
);
out.extend(
self.edges_in
.range(&from, &to)
.map(|e| (e.b, e.rel, false, e.fact)),
);
}
fn render(&self, out: &mut RecallResult, tags_tmp: &mut Vec<TermId>) {
if out.facts.is_empty() && out.edges.is_empty() {
return; }
out.rendered.push_str("## memory\n");
for fact in &out.facts {
let record = self
.facts
.get(&fact.id.0.to_be_bytes())
.expect("selected ids exist");
let text = core::str::from_utf8(self.texts.get(record.text)).unwrap_or("");
let _ = write!(out.rendered, "- [f{}] ", fact.id.0);
if let Some(entity) = fact.entity.some()
&& let Some(name) = self.entity_name(entity)
{
let _ = write!(out.rendered, "{name}: ");
}
out.rendered.push_str(text);
out.rendered.push_str(" (");
render_ym(&mut out.rendered, fact.valid_from);
if fact.valid_to == VALID_TO_OPEN {
out.rendered.push_str("; active)");
} else {
out.rendered.push_str(" → ");
render_ym(&mut out.rendered, fact.valid_to);
out.rendered.push_str("; closed)");
}
tags_tmp.clear();
self.tags_of(fact.id, tags_tmp);
for &tag in tags_tmp.iter() {
let _ = write!(out.rendered, " #{}", self.terms.resolve(tag));
}
out.rendered.push('\n');
}
for edge in &out.edges {
let _ = writeln!(
out.rendered,
"- links: {} —{}→ {}",
self.entity_name(edge.src)
.expect("edges reference existing entities"),
self.terms.resolve(edge.rel),
self.entity_name(edge.dst)
.expect("edges reference existing entities"),
);
}
}
}
fn admit(
facts: &plugmem_arena::Arena<'_, FactRecord>,
allow: &[FactId],
filtered: bool,
as_of: u64,
include_closed: bool,
id: FactId,
) -> Option<FactRecord> {
let record = facts.get(&id.0.to_be_bytes())?;
if record.is_tombstone() || record.recorded_at > as_of || record.valid_from > as_of {
return None;
}
if !include_closed && as_of >= record.valid_to {
return None;
}
if filtered && allow.binary_search(&id).is_err() {
return None;
}
Some(record)
}
fn render_ym(out: &mut String, ms: u64) {
let days = (ms / 86_400_000) as i64;
let z = days + 719_468;
let era = z.div_euclid(146_097);
let doe = z.rem_euclid(146_097);
let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
let mp = (5 * doy + 2) / 153;
let month = if mp < 10 { mp + 3 } else { mp - 9 };
let year = yoe + era * 400 + i64::from(month <= 2);
let _ = write!(out, "{year:04}-{month:02}");
}