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, edge_end, edge_floor, edge_history_ceiling, edge_history_floor,
};
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 TEMPORAL_TAG_FIRST_MAX: usize = SOURCE_CAP * 64;
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>,
pub graph_depth: Option<u32>,
}
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,
graph_depth: None,
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>,
allow_bits: AllowFilter,
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)>,
time_tag: Vec<(FactId, u64)>,
visited: Vec<(EntityId, f32)>,
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(());
}
if filtered {
s.allow_bits.fill(&s.allow);
} else {
s.allow_bits.clear();
}
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;
let allow_bits = &s.allow_bits;
self.bm25.search(
(self.cfg.bm25_k1, self.cfg.bm25_b),
&s.query_terms,
SOURCE_CAP,
&mut |id| {
admit(
facts,
allow,
allow_bits,
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;
let allow_bits = &s.allow_bits;
self.vecs.search(
v,
SOURCE_CAP,
&mut |id| {
admit(
facts,
allow,
allow_bits,
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,
allow_bits,
..
} = 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,
allow_bits,
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,
allow_bits,
graph_out,
visited,
..
} = s;
graph_out.clear();
if visited.is_empty() {
return;
}
let full = |edges: &Vec<RecalledEdge>, visited: &Vec<(EntityId, f32)>| {
edges.len() >= GRAPH_EDGE_CAP && visited.len() >= GRAPH_ENTITY_CAP
};
let mut frontier = 0usize;
let mut weight = 1.0f32;
let depth = q.graph_depth.unwrap_or(self.cfg.graph_depth);
'expand: for _ in 0..depth {
let depth_end = visited.len();
if depth_end == frontier {
break;
}
weight *= self.cfg.graph_decay;
for at in frontier..depth_end {
if full(&out.edges, visited) {
break 'expand;
}
let (entity, _) = visited[at];
self.neighbors(
entity,
as_of,
q.as_of.is_some(),
&mut |neighbor, rel, this_side_src, provenance| {
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));
}
!full(&out.edges, visited)
},
);
}
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,
allow_bits,
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,
allow_bits,
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) {
s.time_out.clear();
let Some((from, to)) = q.range else { return };
if filtered && !s.allow.is_empty() && s.allow.len() <= TEMPORAL_TAG_FIRST_MAX {
self.time_source_from_tags(from, to, as_of, q.include_closed, s);
return;
}
let RecallScratch {
allow,
allow_bits,
time_out,
..
} = s;
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_rev(&from_key, &to_key) {
if admit(
&self.facts,
allow,
allow_bits,
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 {
break;
}
}
}
}
fn time_source_from_tags(
&self,
from: u64,
to: u64,
as_of: u64,
include_closed: bool,
s: &mut RecallScratch,
) {
let RecallScratch {
allow,
time_tag,
time_out,
..
} = s;
time_tag.clear();
for &fact in allow.iter() {
let Some(record) = admit(
&self.facts,
&[],
&AllowFilter::default(),
false,
as_of,
include_closed,
fact,
) else {
continue;
};
if record.recorded_at >= from && record.recorded_at < to {
time_tag.push((fact, record.recorded_at));
}
}
time_tag.sort_unstable_by(|a, b| b.1.cmp(&a.1).then_with(|| b.0.cmp(&a.0)));
time_out.extend(
time_tag
.iter()
.take(SOURCE_CAP)
.map(|&(fact, recorded_at)| (fact, recorded_at as f32)),
);
}
fn neighbors(
&self,
entity: EntityId,
as_of: u64,
historical: bool,
visit: &mut impl FnMut(EntityId, TermId, bool, FactId) -> bool,
) {
if historical {
let from = edge_history_floor(entity);
let to = edge_history_ceiling(entity, as_of);
for (arena, entity_is_src) in
[(&self.edges_hist_out, true), (&self.edges_hist_in, false)]
{
for e in arena.range_rev(&from, &to) {
if as_of < e.valid_to && !visit(e.b, e.rel, entity_is_src, e.fact) {
return;
}
}
}
} else {
let from = edge_floor(entity);
let to = edge_end(entity);
for (arena, entity_is_src) in [(&self.edges_out, true), (&self.edges_in, false)] {
for e in arena.range(&from, &to) {
if !visit(e.b, e.rel, entity_is_src, e.fact) {
return;
}
}
}
}
}
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 (Some(src), Some(dst)) = (self.entity_name(edge.src), self.entity_name(edge.dst))
else {
continue;
};
let _ = writeln!(
out.rendered,
"- links: {src} —{}→ {dst}",
self.terms.resolve(edge.rel),
);
}
}
}
#[derive(Debug, Default)]
struct AllowFilter {
bits: Vec<u64>,
shift: u32,
}
const ALLOW_BITS_PER_MEMBER: usize = 8;
const ALLOW_MIN_WORDS: usize = 8;
const ALLOW_MAX_WORDS: usize = 1 << 14;
impl AllowFilter {
fn fill(&mut self, allow: &[FactId]) {
let words = allow
.len()
.saturating_mul(ALLOW_BITS_PER_MEMBER)
.div_ceil(64)
.clamp(ALLOW_MIN_WORDS, ALLOW_MAX_WORDS)
.next_power_of_two();
self.bits.clear();
self.bits.resize(words, 0);
self.shift = 64 - (words * 64).trailing_zeros();
for &id in allow {
let at = self.index(id);
self.bits[at / 64] |= 1u64 << (at % 64);
}
}
fn clear(&mut self) {
self.bits.clear();
}
fn index(&self, id: FactId) -> usize {
(u64::from(id.0).wrapping_mul(0x9E37_79B9_7F4A_7C15) >> self.shift) as usize
}
fn maybe_contains(&self, id: FactId) -> bool {
let at = self.index(id);
self.bits[at / 64] & (1u64 << (at % 64)) != 0
}
}
fn admit(
facts: &plugmem_arena::Arena<'_, FactRecord>,
allow: &[FactId],
filter: &AllowFilter,
filtered: bool,
as_of: u64,
include_closed: bool,
id: FactId,
) -> Option<FactRecord> {
if filtered && (!filter.maybe_contains(id) || allow.binary_search(&id).is_err()) {
return None;
}
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;
}
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}");
}