#[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::{ColumnFilter, MemoryEdge, Recollection};
use crate::service::Metadata;
pub trait MemoryStore {
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 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>;
fn query_columnar(
&self,
embedding: &[f32],
k: usize,
filters: &[ColumnFilter],
) -> Result<Vec<Recollection>, MemoryError>;
fn relate(&self, from: u64, to: u64, relation: &str) -> Result<u64, MemoryError>;
fn relations(&self, id: u64) -> Result<Vec<MemoryEdge>, MemoryError>;
fn count(&self) -> usize;
}
#[cfg(feature = "persistence")]
pub struct NativeStore {
memory: AgentMemory,
}
#[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(db, dimension)?;
Ok(Self { memory })
}
}
#[cfg(feature = "persistence")]
impl MemoryStore for NativeStore {
fn store(&self, id: u64, content: &str, embedding: &[f32]) -> Result<(), MemoryError> {
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.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.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.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.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.memory.semantic().delete(id).map_err(MemoryError::from)
}
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)
}
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 filter in filters {
self.memory
.semantic()
.ensure_index(&filter.field)
.map_err(MemoryError::from)?;
}
let results = self
.memory
.query_semantic(&sql, ¶ms)
.map_err(MemoryError::from)?;
Ok(results.iter().map(to_recollection).collect())
}
fn relate(&self, from: u64, to: u64, relation: &str) -> Result<u64, MemoryError> {
self.memory
.semantic()
.relate(from, to, relation, None)
.map_err(MemoryError::from)
}
fn relations(&self, id: u64) -> Result<Vec<MemoryEdge>, MemoryError> {
Ok(self
.memory
.semantic()
.relations(id)?
.into_iter()
.map(|edge| MemoryEdge {
from: edge.source(),
to: edge.target(),
relation: edge.label().to_owned(),
})
.collect())
}
fn count(&self) -> usize {
self.memory.semantic().count()
}
}
#[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}");
let _ = write!(
predicate,
" AND {} {} ${key}",
filter.field,
filter.op.as_sql()
);
params.insert(key, filter.value.clone());
}
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_"))
}
#[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;