use std::collections::{HashMap, HashSet};
#[cfg(feature = "persistence")]
use std::path::Path;
use serde_json::{Map, Value};
pub type Metadata = Map<String, Value>;
use crate::clock;
use crate::embedder::Embedder;
use crate::error::MemoryError;
use crate::extract::{ExtractedAttribute, ExtractedRelation, Extractor};
use crate::id;
use crate::model::{
ColumnFilter, EntityProfile, EntityRelation, Explanation, Link, MemoryEdge, MemoryNode,
Recollection, RememberedExtraction, UnrelateOutcome,
};
#[cfg(feature = "persistence")]
use crate::storage::NativeStore;
use crate::storage::{is_reserved_key, strip_reserved_keys, MemoryStore, AUTO_DATE_FIELD};
#[path = "fused_recall.rs"]
mod fused_recall;
#[cfg(feature = "persistence")]
#[path = "reinforce.rs"]
mod reinforce;
#[cfg(feature = "context")]
#[path = "context/memory_bridge.rs"]
mod memory_bridge;
const HUB_FIELD: &str = "_veles_hub";
const HUB_ID_SALT: &str = "\u{0}_veles_entity_hub\u{0}";
const MENTIONS_RELATION: &str = "mentions";
const ABOUT_RELATION: &str = "about";
#[cfg(feature = "persistence")]
pub struct MemoryService<E: Embedder, S: MemoryStore = NativeStore> {
store: S,
embedder: E,
autograph: Option<crate::extract::DynExtractor>,
}
#[cfg(not(feature = "persistence"))]
pub struct MemoryService<E: Embedder, S: MemoryStore> {
store: S,
embedder: E,
autograph: Option<crate::extract::DynExtractor>,
}
#[cfg(feature = "persistence")]
impl<E: Embedder> MemoryService<E, NativeStore> {
pub fn open<P: AsRef<Path>>(path: P, embedder: E) -> Result<Self, MemoryError> {
let store = NativeStore::open(path, embedder.dimension())?;
Ok(Self {
store,
embedder,
autograph: None,
})
}
}
impl<E: Embedder, S: MemoryStore> MemoryService<E, S> {
pub fn with_store(store: S, embedder: E) -> Self {
Self {
store,
embedder,
autograph: None,
}
}
#[must_use]
pub fn with_autograph(mut self, extractor: crate::extract::DynExtractor) -> Self {
self.autograph = Some(extractor);
self
}
pub fn remember(
&self,
fact: &str,
links: &[Link],
metadata: Option<&Metadata>,
) -> Result<u64, MemoryError> {
self.remember_with_ttl(fact, links, metadata, None)
}
pub fn remember_with_ttl(
&self,
fact: &str,
links: &[Link],
metadata: Option<&Metadata>,
ttl_seconds: Option<u64>,
) -> Result<u64, MemoryError> {
self.remember_inner(fact, links, metadata, ttl_seconds, true)
}
fn remember_inner(
&self,
fact: &str,
links: &[Link],
metadata: Option<&Metadata>,
ttl_seconds: Option<u64>,
run_autograph: bool,
) -> Result<u64, MemoryError> {
let fact = fact.trim();
self.validate_write(fact, links, metadata, ttl_seconds)?;
let fact_id = id::stable_id(fact);
reject_self_links(fact_id, links)?;
let existed_before = !links.is_empty() && self.store.get(fact_id)?.is_some();
self.write_fact(fact_id, fact, metadata, ttl_seconds)?;
self.link_or_rollback(fact_id, links, existed_before)?;
self.autograph_if(run_autograph, fact_id, fact);
Ok(fact_id)
}
fn validate_write(
&self,
fact: &str,
links: &[Link],
metadata: Option<&Metadata>,
ttl_seconds: Option<u64>,
) -> Result<(), MemoryError> {
validate_fact(fact)?;
reject_zero_ttl(ttl_seconds)?;
reject_reserved_keys(metadata)?;
reject_oversized_metadata(metadata)?;
self.validate_links(links)
}
fn write_fact(
&self,
fact_id: u64,
fact: &str,
metadata: Option<&Metadata>,
ttl_seconds: Option<u64>,
) -> Result<(), MemoryError> {
let embedding = self.embedder.embed(fact)?;
let stamped = stamp_with_today(metadata);
self.store_fact(fact_id, fact, &embedding, stamped.as_ref(), ttl_seconds)
}
fn validate_links(&self, links: &[Link]) -> Result<(), MemoryError> {
for link in links {
validate_relation(&link.relation)?;
}
self.ensure_link_targets_exist(links)
}
fn link_or_rollback(
&self,
fact_id: u64,
links: &[Link],
existed_before: bool,
) -> Result<(), MemoryError> {
let Err(cause) = self.relate_links(fact_id, links) else {
return Ok(());
};
if existed_before {
return Err(cause);
}
match self.store.delete(fact_id) {
Ok(()) => Err(cause),
Err(rollback) => Err(MemoryError::RollbackFailed {
cause: Box::new(cause),
rollback: Box::new(rollback),
}),
}
}
fn autograph_if(&self, run: bool, fact_id: u64, fact: &str) {
if run {
self.autograph(fact_id, fact);
}
}
fn autograph(&self, fact_id: u64, fact: &str) {
let Some(extractor) = self.autograph.as_ref() else {
return;
};
let Ok(mut extraction) = extractor.extract_graph(fact) else {
return;
};
crate::extract::orient_kinship(fact, &mut extraction.relations);
let mut entity_ids: HashMap<String, u64> = HashMap::new();
let mut edges: HashSet<(u64, u64)> = HashSet::new();
let mut seeded: HashSet<u64> = HashSet::new();
for extracted in &extraction.facts {
let _ = self.wire_entities(
fact_id,
&extracted.entities,
&mut entity_ids,
&mut edges,
&mut seeded,
);
}
let _ = self.wire_relations(
&extraction.relations,
&mut entity_ids,
&mut edges,
&mut seeded,
);
let _ = self.wire_attributes(&extraction.attributes, &mut entity_ids);
}
fn relate_links(&self, fact_id: u64, links: &[Link]) -> Result<(), MemoryError> {
for link in links {
self.store.relate(fact_id, link.target, &link.relation)?;
}
Ok(())
}
pub fn remember_extracted<X: Extractor>(
&self,
text: &str,
extractor: &X,
metadata: Option<&Metadata>,
) -> Result<RememberedExtraction, MemoryError> {
let text = text.trim();
if text.is_empty() {
return Err(MemoryError::EmptyFact);
}
let mut extraction = extractor.extract_graph(text)?;
crate::extract::orient_kinship(text, &mut extraction.relations);
let mut entity_ids: HashMap<String, u64> = HashMap::new();
let mut edges: HashSet<(u64, u64)> = HashSet::new();
let mut seeded: HashSet<u64> = HashSet::new();
let outcome = self.store_extracted_facts(
&extraction.facts,
metadata,
&mut entity_ids,
&mut edges,
&mut seeded,
)?;
self.wire_relations(
&extraction.relations,
&mut entity_ids,
&mut edges,
&mut seeded,
)?;
self.wire_attributes(&extraction.attributes, &mut entity_ids)?;
Ok(outcome)
}
pub fn entity_profile(&self, name: &str) -> Result<Option<EntityProfile>, MemoryError> {
let key = canonical_entity_name(name);
if key.is_empty() {
return Ok(None);
}
let id = id::stable_id(&format!("{HUB_ID_SALT}{key}"));
if self.store.get(id)?.is_none() {
return Ok(None);
}
Ok(Some(EntityProfile {
id,
name: key,
attributes: strip_reserved_keys(self.store.get_metadata(id)?).unwrap_or_default(),
relations: self.outgoing_entity_relations(id)?,
relations_in: self.incoming_entity_relations(id)?,
}))
}
fn outgoing_entity_relations(&self, id: u64) -> Result<Vec<EntityRelation>, MemoryError> {
self.resolve_entity_relations(self.store.relations(id)?, |edge| edge.to)
}
fn incoming_entity_relations(&self, id: u64) -> Result<Vec<EntityRelation>, MemoryError> {
self.resolve_entity_relations(self.store.incoming_relations(id)?, |edge| edge.from)
}
fn resolve_entity_relations(
&self,
edges: Vec<MemoryEdge>,
far_end: impl Fn(&MemoryEdge) -> u64,
) -> Result<Vec<EntityRelation>, MemoryError> {
let mut relations = Vec::new();
for edge in edges {
if edge.relation == MENTIONS_RELATION || edge.relation == ABOUT_RELATION {
continue;
}
let far = far_end(&edge);
let content = self.store.get(far)?.map(|(content, _)| content);
relations.push(EntityRelation {
predicate: edge.relation,
target_id: far,
target: content.unwrap_or_default(),
});
}
Ok(relations)
}
fn wire_relations(
&self,
relations: &[ExtractedRelation],
entity_ids: &mut HashMap<String, u64>,
edges: &mut HashSet<(u64, u64)>,
seeded: &mut HashSet<u64>,
) -> Result<(), MemoryError> {
for relation in relations {
if validate_relation(&relation.predicate).is_err() {
continue;
}
let subject_id = self.entity_hub(&relation.subject, entity_ids)?;
let object_id = self.entity_hub(&relation.object, entity_ids)?;
if subject_id == object_id {
continue;
}
self.seed_existing_edges(subject_id, edges, seeded)?;
self.add_edge(subject_id, object_id, &relation.predicate, edges)?;
}
Ok(())
}
fn wire_attributes(
&self,
attributes: &[ExtractedAttribute],
entity_ids: &mut HashMap<String, u64>,
) -> Result<(), MemoryError> {
let mut per_entity: HashMap<String, Metadata> = HashMap::new();
for attribute in attributes {
if is_reserved_key(&attribute.key) {
continue;
}
per_entity
.entry(attribute.entity.clone())
.or_default()
.insert(attribute.key.clone(), attribute.value.clone());
}
for (entity, meta) in per_entity {
if meta.is_empty() {
continue;
}
reject_oversized_metadata(Some(&meta))?;
let hub_id = self.entity_hub(&entity, entity_ids)?;
self.store.update_metadata(hub_id, &meta)?;
}
Ok(())
}
fn store_extracted_facts(
&self,
facts: &[crate::extract::ExtractedFact],
metadata: Option<&Metadata>,
entity_ids: &mut HashMap<String, u64>,
edges: &mut HashSet<(u64, u64)>,
seeded: &mut HashSet<u64>,
) -> Result<RememberedExtraction, MemoryError> {
let mut ids = Vec::with_capacity(facts.len());
let mut skipped_over_cap = 0;
for fact in facts {
let content = fact.text.trim();
if content.is_empty() {
continue;
}
let fact_id = match self.remember_inner(content, &[], metadata, None, false) {
Ok(id) => id,
Err(MemoryError::FactTooLarge { .. }) => {
skipped_over_cap += 1;
continue;
}
Err(error) => return Err(error),
};
ids.push(fact_id);
self.wire_entities(fact_id, &fact.entities, entity_ids, edges, seeded)?;
}
Ok(RememberedExtraction {
ids,
skipped_over_cap,
})
}
fn wire_entities(
&self,
fact_id: u64,
entities: &[String],
entity_ids: &mut HashMap<String, u64>,
edges: &mut HashSet<(u64, u64)>,
seeded: &mut HashSet<u64>,
) -> Result<(), MemoryError> {
for entity in entities {
if entity.chars().any(char::is_alphanumeric) {
self.wire_entity(fact_id, entity, entity_ids, edges, seeded)?;
}
}
Ok(())
}
fn wire_entity(
&self,
fact_id: u64,
entity: &str,
entity_ids: &mut HashMap<String, u64>,
edges: &mut HashSet<(u64, u64)>,
seeded: &mut HashSet<u64>,
) -> Result<(), MemoryError> {
let entity_id = self.entity_hub(entity, entity_ids)?;
if entity_id == fact_id {
return Ok(());
}
self.seed_existing_edges(fact_id, edges, seeded)?;
self.seed_existing_edges(entity_id, edges, seeded)?;
self.add_edge(fact_id, entity_id, ABOUT_RELATION, edges)?;
self.add_edge(entity_id, fact_id, MENTIONS_RELATION, edges)?;
Ok(())
}
fn add_edge(
&self,
from: u64,
to: u64,
label: &str,
edges: &mut HashSet<(u64, u64)>,
) -> Result<(), MemoryError> {
if edges.insert((from, to)) {
self.relate(from, to, label)?;
}
Ok(())
}
fn seed_existing_edges(
&self,
node: u64,
edges: &mut HashSet<(u64, u64)>,
seeded: &mut HashSet<u64>,
) -> Result<(), MemoryError> {
if !seeded.insert(node) {
return Ok(());
}
for edge in self.store.relations(node)? {
edges.insert((node, edge.to));
}
Ok(())
}
fn entity_hub(
&self,
entity: &str,
entity_ids: &mut HashMap<String, u64>,
) -> Result<u64, MemoryError> {
let key = entity.trim().to_lowercase();
if let Some(&id) = entity_ids.get(&key) {
return Ok(id);
}
let id = self.remember_hub(&key)?;
entity_ids.insert(key, id);
Ok(id)
}
fn remember_hub(&self, key: &str) -> Result<u64, MemoryError> {
let id = id::stable_id(&format!("{HUB_ID_SALT}{key}"));
if self.store.get(id)?.is_some() {
return Ok(id);
}
let content = format!("Entity: {key}");
let embedding = self.embedder.embed(&content)?;
let mut meta = Map::new();
meta.insert(HUB_FIELD.to_string(), Value::Bool(true));
self.store_fact(id, &content, &embedding, Some(&meta), None)?;
Ok(id)
}
fn ensure_exists(&self, id: u64) -> Result<(), MemoryError> {
if self.store.get(id)?.is_none() {
return Err(MemoryError::UnknownMemory(id));
}
Ok(())
}
fn ensure_link_targets_exist(&self, links: &[Link]) -> Result<(), MemoryError> {
for link in links {
self.ensure_exists(link.target)?;
}
Ok(())
}
fn store_fact(
&self,
id: u64,
fact: &str,
embedding: &[f32],
metadata: Option<&Metadata>,
ttl_seconds: Option<u64>,
) -> Result<(), MemoryError> {
match (metadata, ttl_seconds) {
(Some(meta), Some(ttl)) => {
self.store
.store_with_metadata_and_ttl(id, fact, embedding, meta, ttl)?;
}
(Some(meta), None) => self.store.store_with_metadata(id, fact, embedding, meta)?,
(None, Some(ttl)) => self.store.store_with_ttl(id, fact, embedding, ttl)?,
(None, None) => self.store.store(id, fact, embedding)?,
}
Ok(())
}
pub fn recall(
&self,
query: &str,
k: usize,
filter: Option<&Metadata>,
) -> Result<Vec<Recollection>, MemoryError> {
let query = query.trim();
if query.is_empty() {
return Ok(Vec::new());
}
reject_reserved_keys(filter)?;
let embedding = self.embedder.embed(query)?;
let hits = self.search(&embedding, k, filter)?;
let ids: Vec<u64> = hits.iter().map(|(id, _, _)| *id).collect();
let payloads = self.store.get_metadata_batch(&ids)?;
#[cfg(feature = "persistence")]
let (hits, payloads) = Self::rl_rerank(hits, payloads);
Ok(hits
.into_iter()
.zip(payloads)
.map(|((id, score, content), payload)| Recollection {
id,
score,
content,
metadata: strip_reserved_keys(payload),
})
.collect())
}
fn search(
&self,
embedding: &[f32],
k: usize,
filter: Option<&Metadata>,
) -> Result<Vec<(u64, f32, String)>, MemoryError> {
match filter {
Some(meta) if !meta.is_empty() => self.store.query_filtered(embedding, k, meta, 0),
_ => self
.store
.query_excluding(embedding, k, &hub_exclude_filter()),
}
}
pub fn recall_where(
&self,
query: &str,
k: usize,
filters: &[ColumnFilter],
) -> Result<Vec<Recollection>, MemoryError> {
let query = query.trim();
if query.is_empty() || k == 0 {
return Ok(Vec::new());
}
if filters.is_empty() {
return self.recall(query, k, None);
}
let embedding = self.embedder.embed(query)?;
self.store.query_columnar(&embedding, k, filters)
}
pub fn relate(&self, from: u64, to: u64, relation: &str) -> Result<u64, MemoryError> {
validate_relation(relation)?;
if from == to {
return Err(MemoryError::SelfRelation(from));
}
self.ensure_exists(from)?;
self.ensure_exists(to)?;
self.store.relate(from, to, relation)
}
pub fn unrelate(
&self,
from: u64,
to: u64,
relation: &str,
) -> Result<UnrelateOutcome, MemoryError> {
validate_relation(relation)?;
if from == to {
return Err(MemoryError::SelfRelation(from));
}
let removed = self.remove_matching_edges(from, to, relation)?;
Ok(UnrelateOutcome {
found: removed > 0,
removed,
})
}
fn remove_matching_edges(
&self,
from: u64,
to: u64,
relation: &str,
) -> Result<usize, MemoryError> {
let mut removed = 0usize;
for edge in self.store.relations(from)? {
if edge.to == to && edge.relation == relation && self.store.unrelate(edge.id)? {
removed += 1;
}
}
Ok(removed)
}
pub fn forget(&self, fact_id: u64) -> Result<bool, MemoryError> {
let found = self.store.get(fact_id)?.is_some();
let hubs = self.hubs_linked_from(fact_id)?;
self.store.delete(fact_id)?;
self.collect_orphan_hubs(&hubs)?;
Ok(found)
}
fn hubs_linked_from(&self, fact_id: u64) -> Result<Vec<u64>, MemoryError> {
let mut hubs = Vec::new();
for edge in self.store.relations(fact_id)? {
if self.is_hub(edge.to)? {
hubs.push(edge.to);
}
}
Ok(hubs)
}
fn collect_orphan_hubs(&self, hubs: &[u64]) -> Result<(), MemoryError> {
for &hub in hubs {
if !self.hub_still_mentioned(hub)? {
self.store.delete(hub)?;
}
}
Ok(())
}
fn hub_still_mentioned(&self, hub: u64) -> Result<bool, MemoryError> {
for edge in self.store.relations(hub)? {
if edge.relation == MENTIONS_RELATION && self.store.get(edge.to)?.is_some() {
return Ok(true);
}
}
self.hub_has_live_referent(hub)
}
fn hub_has_live_referent(&self, hub: u64) -> Result<bool, MemoryError> {
for edge in self.store.incoming_relations(hub)? {
if self.store.get(edge.from)?.is_some() && !self.is_hub(edge.from)? {
return Ok(true);
}
}
Ok(false)
}
fn is_hub(&self, id: u64) -> Result<bool, MemoryError> {
Ok(self
.store
.get_metadata(id)?
.is_some_and(|meta| meta.contains_key(HUB_FIELD)))
}
pub fn why(
&self,
decision: &str,
max_hops: usize,
filter: Option<&Metadata>,
) -> Result<Explanation, MemoryError> {
let decision = decision.trim();
if decision.is_empty() {
return Ok(Explanation::default());
}
reject_reserved_keys(filter)?;
let embedding = self.embedder.embed(decision)?;
let seeds = self.search(&embedding, 1, filter)?;
let Some((seed_id, _score, seed_content)) = seeds.into_iter().next() else {
return Ok(Explanation::default());
};
self.traverse(seed_id, seed_content, max_hops)
}
fn traverse(
&self,
seed_id: u64,
seed_content: String,
max_hops: usize,
) -> Result<Explanation, MemoryError> {
let mut explanation = Explanation {
nodes: vec![MemoryNode {
id: seed_id,
content: seed_content,
hop: 0,
}],
edges: Vec::new(),
};
let mut visited: HashSet<u64> = HashSet::from([seed_id]);
let mut frontier = vec![seed_id];
let mut next: Vec<u64> = Vec::new();
for hop in 1..=max_hops {
next.clear();
for node_id in frontier.drain(..) {
self.expand(node_id, hop, &mut explanation, &mut visited, &mut next)?;
}
if next.is_empty() {
break;
}
std::mem::swap(&mut frontier, &mut next);
}
Ok(explanation)
}
fn expand(
&self,
node_id: u64,
hop: usize,
explanation: &mut Explanation,
visited: &mut HashSet<u64>,
next: &mut Vec<u64>,
) -> Result<(), MemoryError> {
for edge in self.store.relations(node_id)? {
let target = edge.to;
if !visited.contains(&target) {
let Some((content, _embedding)) = self.store.get(target)? else {
continue; };
visited.insert(target);
explanation.nodes.push(MemoryNode {
id: target,
content,
hop,
});
next.push(target);
}
explanation.edges.push(edge);
}
Ok(())
}
}
fn hub_exclude_filter() -> Metadata {
let mut exclude = Map::new();
exclude.insert(HUB_FIELD.to_string(), Value::Bool(true));
exclude
}
fn reject_reserved_keys(metadata: Option<&Metadata>) -> Result<(), MemoryError> {
let Some(meta) = metadata else {
return Ok(());
};
for key in meta.keys() {
if is_reserved_key(key) {
return Err(MemoryError::ReservedKey(key.clone()));
}
}
Ok(())
}
fn reject_oversized_metadata(metadata: Option<&Metadata>) -> Result<(), MemoryError> {
let Some(meta) = metadata else {
return Ok(());
};
let bytes = crate::limits::metadata_bytes(meta);
if bytes > crate::limits::MAX_METADATA_BYTES {
return Err(MemoryError::MetadataTooLarge {
bytes,
max: crate::limits::MAX_METADATA_BYTES,
});
}
Ok(())
}
#[cfg(feature = "context")]
pub(crate) fn positive_ttl(ttl_seconds: Option<u64>) -> Option<u64> {
ttl_seconds.filter(|&seconds| seconds > 0)
}
#[must_use]
pub fn canonical_entity_name(name: &str) -> String {
name.trim().to_lowercase()
}
fn validate_fact(fact: &str) -> Result<(), MemoryError> {
if fact.is_empty() {
return Err(MemoryError::EmptyFact);
}
validate_embeddable(fact)
}
pub(crate) fn validate_embeddable(text: &str) -> Result<(), MemoryError> {
if text.len() > crate::limits::MAX_EMBEDDABLE_TEXT_BYTES {
return Err(MemoryError::FactTooLarge {
bytes: text.len(),
max: crate::limits::MAX_EMBEDDABLE_TEXT_BYTES,
});
}
Ok(())
}
#[cfg(feature = "context")]
pub(crate) fn embeddable_prefix(text: &str) -> &str {
let cap = crate::limits::MAX_EMBEDDABLE_TEXT_BYTES;
if text.len() <= cap {
return text;
}
let mut end = cap;
while end > 0 && !text.is_char_boundary(end) {
end -= 1;
}
&text[..end]
}
fn reject_zero_ttl(ttl_seconds: Option<u64>) -> Result<(), MemoryError> {
if ttl_seconds == Some(0) {
return Err(MemoryError::ZeroTtl);
}
Ok(())
}
fn reject_self_links(fact_id: u64, links: &[Link]) -> Result<(), MemoryError> {
if links.iter().any(|link| link.target == fact_id) {
return Err(MemoryError::SelfRelation(fact_id));
}
Ok(())
}
fn stamp_with_today(metadata: Option<&Metadata>) -> Option<Metadata> {
if metadata.is_some_and(|meta| meta.contains_key(AUTO_DATE_FIELD)) {
return metadata.cloned();
}
let Some(today) = clock::today_ymd() else {
return metadata.cloned();
};
let mut stamped = metadata.cloned().unwrap_or_default();
stamped.insert(AUTO_DATE_FIELD.to_owned(), Value::from(today));
Some(stamped)
}
const MAX_RELATION_BYTES: usize = 512;
fn validate_relation(label: &str) -> Result<(), MemoryError> {
if label.is_empty() {
return Err(MemoryError::InvalidRelation(
"relation label must not be empty".to_owned(),
));
}
if label.len() > MAX_RELATION_BYTES {
return Err(MemoryError::InvalidRelation(format!(
"relation label exceeds maximum of {MAX_RELATION_BYTES} bytes ({} given)",
label.len()
)));
}
if label.chars().any(|c| c.is_ascii_control()) {
return Err(MemoryError::InvalidRelation(
"relation label must not contain ASCII control characters".to_owned(),
));
}
Ok(())
}