#[cfg(feature = "persistence")]
use std::collections::HashMap;
#[cfg(feature = "persistence")]
use std::path::Path;
#[cfg(feature = "persistence")]
use std::sync::Arc;
#[cfg(feature = "persistence")]
use serde_json::json;
use serde_json::Value;
#[cfg(feature = "persistence")]
use velesdb_core::agent::AgentMemory;
#[cfg(feature = "persistence")]
use velesdb_core::{Database, SearchResult};
use crate::error::MemoryError;
use crate::model::{BoundedMemoryEdges, ColumnFilter, MemoryEdge, Recollection};
#[cfg(feature = "persistence")]
use crate::mutation::{DirtyKey, MutationCapture, MutationObserver};
use crate::service::Metadata;
#[cfg(feature = "persistence")]
mod migration;
pub trait FactStore {
fn store(&self, id: u64, content: &str, embedding: &[f32]) -> Result<(), MemoryError>;
fn store_with_metadata(
&self,
id: u64,
content: &str,
embedding: &[f32],
metadata: &Metadata,
) -> Result<(), MemoryError>;
fn store_with_ttl(
&self,
id: u64,
content: &str,
embedding: &[f32],
ttl_seconds: u64,
) -> Result<(), MemoryError>;
fn store_with_metadata_and_ttl(
&self,
id: u64,
content: &str,
embedding: &[f32],
metadata: &Metadata,
ttl_seconds: u64,
) -> Result<(), MemoryError> {
self.store_with_ttl(id, content, embedding, ttl_seconds)?;
self.update_metadata(id, metadata)
}
fn update_metadata(&self, id: u64, metadata: &Metadata) -> Result<(), MemoryError>;
fn get(&self, id: u64) -> Result<Option<(String, Vec<f32>)>, MemoryError>;
fn get_metadata(&self, id: u64) -> Result<Option<Metadata>, MemoryError>;
fn get_metadata_batch(&self, ids: &[u64]) -> Result<Vec<Option<Metadata>>, MemoryError>;
fn delete(&self, id: u64) -> Result<(), MemoryError>;
fn count(&self) -> usize;
fn list(
&self,
cursor: Option<u64>,
limit: usize,
) -> Result<(Vec<RawListedFact>, Option<u64>), MemoryError> {
let _ = (cursor, limit);
Err(MemoryError::Unsupported(
"this storage backend does not support listing",
))
}
}
pub trait RecallStore: FactStore {
fn query_filtered(
&self,
embedding: &[f32],
k: usize,
filter: &Metadata,
offset: usize,
) -> Result<Vec<(u64, f32, String)>, MemoryError>;
fn query_excluding(
&self,
embedding: &[f32],
k: usize,
exclude: &Metadata,
) -> Result<Vec<(u64, f32, String)>, MemoryError>;
}
pub trait ColumnStore {
fn query_columnar(
&self,
embedding: &[f32],
k: usize,
filters: &[ColumnFilter],
) -> Result<Vec<Recollection>, MemoryError>;
}
pub trait GraphStore {
fn relate(&self, from: u64, to: u64, relation: &str) -> Result<u64, MemoryError>;
fn relations(&self, id: u64) -> Result<Vec<MemoryEdge>, MemoryError>;
fn incoming_relations(&self, id: u64) -> Result<Vec<MemoryEdge>, MemoryError>;
fn relations_bounded(&self, id: u64, cap: usize) -> Result<BoundedMemoryEdges, MemoryError>;
fn incoming_relations_bounded(
&self,
id: u64,
cap: usize,
) -> Result<BoundedMemoryEdges, MemoryError>;
fn unrelate(&self, edge_id: u64) -> Result<bool, MemoryError>;
fn unrelate_from(&self, from: u64, edge_id: u64) -> Result<bool, MemoryError> {
let _ = from;
self.unrelate(edge_id)
}
fn edge_count(&self) -> Option<usize> {
None
}
}
pub trait MemoryStore: RecallStore + GraphStore + ColumnStore {}
impl<T: RecallStore + GraphStore + ColumnStore> MemoryStore for T {}
#[derive(Debug, Clone)]
pub struct RawListedFact {
pub id: u64,
pub content: String,
pub payload: Metadata,
}
#[cfg(feature = "persistence")]
impl RawListedFact {
pub(crate) fn from_raw(fact: &crate::migration::RawFact) -> Self {
let mut payload: Metadata = serde_json::from_str(&fact.payload).unwrap_or_default();
let content = match payload.remove("content") {
Some(Value::String(text)) => text,
_ => String::new(),
};
Self {
id: fact.id,
content,
payload,
}
}
}
#[cfg(feature = "persistence")]
pub struct NativeStore {
memory: AgentMemory,
db: Arc<Database>,
capture: MutationCapture,
}
#[cfg(feature = "persistence")]
impl NativeStore {
pub fn open<P: AsRef<Path>>(path: P, dimension: usize) -> Result<Self, MemoryError> {
let db = Arc::new(Database::open(path)?);
let memory = AgentMemory::with_dimension(Arc::clone(&db), dimension)?;
Ok(Self {
memory,
db,
capture: MutationCapture::default(),
})
}
pub(crate) fn set_mutation_observer(
&self,
observer: Option<Arc<dyn MutationObserver>>,
) -> Result<(), MemoryError> {
self.capture.replace(observer)
}
pub(crate) fn mutation_capture_active(&self) -> bool {
self.capture.is_active()
}
fn unrelate_unobserved(&self, edge_id: u64) -> Result<bool, MemoryError> {
self.memory
.semantic()
.unrelate(edge_id)
.map_err(MemoryError::from)
}
}
#[cfg(feature = "persistence")]
impl FactStore for NativeStore {
fn store(&self, id: u64, content: &str, embedding: &[f32]) -> Result<(), MemoryError> {
self.capture.observe(DirtyKey::Fact(id))?;
self.memory
.semantic()
.store(id, content, embedding)
.map_err(MemoryError::from)
}
fn store_with_metadata(
&self,
id: u64,
content: &str,
embedding: &[f32],
metadata: &Metadata,
) -> Result<(), MemoryError> {
self.capture.observe(DirtyKey::Fact(id))?;
self.memory
.semantic()
.store_with_metadata(id, content, embedding, metadata)
.map_err(MemoryError::from)
}
fn store_with_ttl(
&self,
id: u64,
content: &str,
embedding: &[f32],
ttl_seconds: u64,
) -> Result<(), MemoryError> {
self.capture.observe(DirtyKey::Fact(id))?;
self.memory
.semantic()
.store_with_ttl(id, content, embedding, ttl_seconds)
.map_err(MemoryError::from)
}
fn update_metadata(&self, id: u64, metadata: &Metadata) -> Result<(), MemoryError> {
self.capture.observe(DirtyKey::Fact(id))?;
self.memory
.semantic()
.update_metadata(id, metadata)
.map_err(MemoryError::from)
}
fn store_with_metadata_and_ttl(
&self,
id: u64,
content: &str,
embedding: &[f32],
metadata: &Metadata,
ttl_seconds: u64,
) -> Result<(), MemoryError> {
self.capture.observe(DirtyKey::Fact(id))?;
self.memory
.semantic()
.store_with_metadata(id, content, embedding, metadata)
.map_err(MemoryError::from)?;
self.memory
.semantic()
.set_ttl_durable(id, ttl_seconds)
.map_err(MemoryError::from)
}
fn get(&self, id: u64) -> Result<Option<(String, Vec<f32>)>, MemoryError> {
self.memory.semantic().get(id).map_err(MemoryError::from)
}
fn get_metadata(&self, id: u64) -> Result<Option<Metadata>, MemoryError> {
self.memory
.semantic()
.get_metadata(id)
.map_err(MemoryError::from)
}
fn get_metadata_batch(&self, ids: &[u64]) -> Result<Vec<Option<Metadata>>, MemoryError> {
self.memory
.semantic()
.get_metadata_batch(ids)
.map_err(MemoryError::from)
}
fn delete(&self, id: u64) -> Result<(), MemoryError> {
self.capture.observe(DirtyKey::Fact(id))?;
self.memory.semantic().delete(id).map_err(MemoryError::from)
}
fn count(&self) -> usize {
self.memory.semantic().count()
}
fn list(
&self,
cursor: Option<u64>,
limit: usize,
) -> Result<(Vec<RawListedFact>, Option<u64>), MemoryError> {
let (facts, next) = crate::migration::scroll_page(
&self.db,
self.memory.semantic().collection_name(),
cursor,
limit,
)?;
let listed = facts.iter().map(RawListedFact::from_raw).collect();
Ok((listed, next))
}
}
#[cfg(feature = "persistence")]
impl RecallStore for NativeStore {
fn query_filtered(
&self,
embedding: &[f32],
k: usize,
filter: &Metadata,
offset: usize,
) -> Result<Vec<(u64, f32, String)>, MemoryError> {
self.memory
.semantic()
.query_filtered(embedding, k, filter, offset)
.map_err(MemoryError::from)
}
fn query_excluding(
&self,
embedding: &[f32],
k: usize,
exclude: &Metadata,
) -> Result<Vec<(u64, f32, String)>, MemoryError> {
self.memory
.semantic()
.query_excluding(embedding, k, exclude)
.map_err(MemoryError::from)
}
}
#[cfg(feature = "persistence")]
impl ColumnStore for NativeStore {
fn query_columnar(
&self,
embedding: &[f32],
k: usize,
filters: &[ColumnFilter],
) -> Result<Vec<Recollection>, MemoryError> {
let (sql, params) = self.build_fused_query(embedding, k, filters)?;
for field in filters
.iter()
.map(|filter| filter.field.as_str())
.chain(INTERNAL_MARKER_FIELDS.iter().copied())
{
self.memory
.semantic()
.ensure_index(field)
.map_err(MemoryError::from)?;
}
let results = self
.memory
.query_semantic(&sql, ¶ms)
.map_err(MemoryError::from)?;
Ok(results.iter().map(to_recollection).collect())
}
}
#[cfg(feature = "persistence")]
impl GraphStore for NativeStore {
fn relate(&self, from: u64, to: u64, relation: &str) -> Result<u64, MemoryError> {
self.capture.observe(DirtyKey::OutgoingEdges(from))?;
self.memory
.semantic()
.relate(from, to, relation, None)
.map_err(MemoryError::from)
}
fn relations(&self, id: u64) -> Result<Vec<MemoryEdge>, MemoryError> {
Ok(to_memory_edges(self.memory.semantic().relations(id)?))
}
fn incoming_relations(&self, id: u64) -> Result<Vec<MemoryEdge>, MemoryError> {
Ok(to_memory_edges(
self.memory.semantic().incoming_relations(id)?,
))
}
fn relations_bounded(&self, id: u64, cap: usize) -> Result<BoundedMemoryEdges, MemoryError> {
let bounded = self.memory.semantic().relations_bounded(id, cap)?;
Ok(BoundedMemoryEdges {
edges: to_memory_edges(bounded.edges),
truncated: bounded.truncated,
})
}
fn incoming_relations_bounded(
&self,
id: u64,
cap: usize,
) -> Result<BoundedMemoryEdges, MemoryError> {
let bounded = self.memory.semantic().incoming_relations_bounded(id, cap)?;
Ok(BoundedMemoryEdges {
edges: to_memory_edges(bounded.edges),
truncated: bounded.truncated,
})
}
fn unrelate(&self, edge_id: u64) -> Result<bool, MemoryError> {
if self.capture.is_active() {
return Err(MemoryError::MigrationCapture(format!(
"cannot remove edge {edge_id} without its source id"
)));
}
self.unrelate_unobserved(edge_id)
}
fn unrelate_from(&self, from: u64, edge_id: u64) -> Result<bool, MemoryError> {
self.capture.observe(DirtyKey::OutgoingEdges(from))?;
self.unrelate_unobserved(edge_id)
}
fn edge_count(&self) -> Option<usize> {
self.memory.semantic().edge_count().ok()
}
}
#[cfg(feature = "persistence")]
fn to_memory_edges(edges: Vec<velesdb_core::collection::graph::GraphEdge>) -> Vec<MemoryEdge> {
edges
.into_iter()
.map(|edge| MemoryEdge {
id: edge.id(),
from: edge.source(),
to: edge.target(),
relation: edge.label().to_owned(),
})
.collect()
}
#[cfg(feature = "persistence")]
impl NativeStore {
fn build_fused_query(
&self,
embedding: &[f32],
k: usize,
filters: &[ColumnFilter],
) -> Result<(String, HashMap<String, Value>), MemoryError> {
use std::fmt::Write as _;
let mut params: HashMap<String, Value> = HashMap::new();
params.insert("q".to_string(), json!(embedding));
let mut predicate = String::from("vector NEAR $q");
for (index, filter) in filters.iter().enumerate() {
validate_column_filter(filter)?;
let key = format!("p{index}");
if matches!(filter.op, crate::model::ColumnOp::Ne) {
let _ = write!(predicate, " AND {} IS NOT NULL", filter.field);
}
let _ = write!(
predicate,
" AND {} {} ${key}",
filter.field,
filter.op.as_sql()
);
params.insert(key, filter.value.clone());
}
for (index, marker) in INTERNAL_MARKER_FIELDS.iter().enumerate() {
let key = format!("m{index}");
let _ = write!(predicate, " AND {marker} != ${key}");
params.insert(key, json!(true));
}
let sql = format!(
"SELECT * FROM {} WHERE {predicate} LIMIT {k}",
self.memory.semantic().collection_name()
);
Ok((sql, params))
}
}
pub const AUTO_DATE_FIELD: &str = "_veles_date";
pub(crate) fn is_reserved_key(key: &str) -> bool {
key != AUTO_DATE_FIELD && (key == "content" || key.starts_with("_veles_"))
}
pub const HUB_FIELD: &str = "_veles_hub";
pub const CTX_EVENT_FIELD: &str = "_veles_ctx_event";
pub const CTX_SOURCE_FIELD: &str = "_veles_ctx_source";
pub const CTX_WORKING_FIELD: &str = "_veles_ctx_working";
pub const CTX_WORKING_INDEX_FIELD: &str = "_veles_ctx_working_index";
pub const INTERNAL_MARKER_FIELDS: &[&str] = &[
HUB_FIELD,
CTX_EVENT_FIELD,
CTX_SOURCE_FIELD,
CTX_WORKING_FIELD,
CTX_WORKING_INDEX_FIELD,
];
#[must_use]
pub fn is_internal_scaffolding(payload: &Metadata) -> bool {
INTERNAL_MARKER_FIELDS
.iter()
.any(|marker| payload.contains_key(*marker))
}
#[must_use]
pub fn strip_reserved_keys(payload: Option<Metadata>) -> Option<Metadata> {
payload.and_then(|payload| {
let metadata: Metadata = payload
.into_iter()
.filter(|(key, _)| !is_reserved_key(key))
.collect();
(!metadata.is_empty()).then_some(metadata)
})
}
#[must_use]
pub fn strip_reserved_keys_ref(payload: Option<&Metadata>) -> Option<Metadata> {
payload.and_then(|payload| {
let metadata: Metadata = payload
.iter()
.filter(|(key, _)| !is_reserved_key(key))
.map(|(key, value)| (key.clone(), value.clone()))
.collect();
(!metadata.is_empty()).then_some(metadata)
})
}
#[cfg(feature = "persistence")]
fn to_recollection(result: &SearchResult) -> Recollection {
let payload = result.point.payload.as_ref().and_then(Value::as_object);
let content = payload
.and_then(|payload| payload.get("content"))
.and_then(Value::as_str)
.unwrap_or_default()
.to_owned();
Recollection {
id: result.point.id,
score: result.score,
content,
metadata: strip_reserved_keys_ref(payload),
}
}
pub fn validate_column_filter(filter: &ColumnFilter) -> Result<(), MemoryError> {
let field = &filter.field;
let plain = !field.is_empty() && field.chars().all(|c| c.is_ascii_alphanumeric() || c == '_');
if !plain || is_reserved_key(field) {
return Err(MemoryError::InvalidFilter(field.clone()));
}
match &filter.value {
Value::String(_) | Value::Number(_) | Value::Bool(_) => Ok(()),
value => Err(MemoryError::InvalidFilter(format!(
"value must be a string, number, or boolean, got {value}"
))),
}
}
#[cfg(all(test, feature = "persistence"))]
#[path = "storage_tests.rs"]
mod tests;