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::Extractor;
use crate::id;
use crate::model::{ColumnFilter, Explanation, Link, MemoryNode, Recollection};
#[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";
#[cfg(feature = "persistence")]
pub struct MemoryService<E: Embedder, S: MemoryStore = NativeStore> {
store: S,
embedder: E,
}
#[cfg(not(feature = "persistence"))]
pub struct MemoryService<E: Embedder, S: MemoryStore> {
store: S,
embedder: E,
}
#[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 })
}
}
impl<E: Embedder, S: MemoryStore> MemoryService<E, S> {
pub fn with_store(store: S, embedder: E) -> Self {
Self { store, embedder }
}
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> {
let fact = fact.trim();
if fact.is_empty() {
return Err(MemoryError::EmptyFact);
}
reject_reserved_keys(metadata)?;
reject_oversized_metadata(metadata)?;
for link in links {
validate_relation(&link.relation)?;
}
self.ensure_link_targets_exist(links)?;
let fact_id = id::stable_id(fact);
let embedding = self.embedder.embed(fact)?;
let existed_before = !links.is_empty() && self.store.get(fact_id)?.is_some();
let stamped = stamp_with_today(metadata);
self.store_fact(
fact_id,
fact,
&embedding,
stamped.as_ref(),
positive_ttl(ttl_seconds),
)?;
if let Err(e) = self.relate_links(fact_id, links) {
if !existed_before {
if let Err(rollback) = self.store.delete(fact_id) {
return Err(MemoryError::RollbackFailed {
cause: Box::new(e),
rollback: Box::new(rollback),
});
}
}
return Err(e);
}
Ok(fact_id)
}
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<Vec<u64>, MemoryError> {
let text = text.trim();
if text.is_empty() {
return Err(MemoryError::EmptyFact);
}
let facts = extractor.extract(text)?;
let mut fact_ids = Vec::with_capacity(facts.len());
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 fact in &facts {
let content = fact.text.trim();
if content.is_empty() {
continue;
}
let fact_id = self.remember(content, &[], metadata)?;
fact_ids.push(fact_id);
self.wire_entities(
fact_id,
&fact.entities,
&mut entity_ids,
&mut edges,
&mut seeded,
)?;
}
Ok(fact_ids)
}
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", 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}"));
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_ttl(id, fact, embedding, ttl)?;
self.store.update_metadata(id, meta)?;
}
(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)?;
self.ensure_exists(from)?;
self.ensure_exists(to)?;
self.store.relate(from, to, relation)
}
pub fn forget(&self, fact_id: u64) -> Result<bool, MemoryError> {
let found = self.store.get(fact_id)?.is_some();
self.store.delete(fact_id)?;
Ok(found)
}
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(())
}
fn positive_ttl(ttl_seconds: Option<u64>) -> Option<u64> {
ttl_seconds.filter(|&seconds| seconds > 0)
}
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(())
}