use parking_lot::RwLock;
use super::{DistanceMetric, VelesCollection, VelesDatabase, VelesError, VelesPoint};
#[derive(Debug, Clone, uniffi::Record)]
pub struct SemanticResult {
pub id: u64,
pub score: f32,
pub content: String,
}
#[derive(uniffi::Object)]
pub struct VelesSemanticMemory {
collection: std::sync::Arc<VelesCollection>,
contents: RwLock<std::collections::HashMap<u64, String>>,
}
#[uniffi::export]
impl VelesSemanticMemory {
#[uniffi::constructor]
#[allow(deprecated)]
pub fn new(db: &VelesDatabase, dimension: u32) -> Result<Self, VelesError> {
let collection_name = "_semantic_memory";
let collection = match db.get_collection(collection_name.to_string())? {
Some(coll) => coll,
None => {
db.create_collection(
collection_name.to_string(),
dimension,
DistanceMetric::Cosine,
)?;
db.get_collection(collection_name.to_string())?
.ok_or(VelesError::Database {
message: "Failed to retrieve collection after creation".to_string(),
})?
}
};
Ok(Self {
collection,
contents: RwLock::new(std::collections::HashMap::new()),
})
}
pub fn store(&self, id: u64, content: String, embedding: Vec<f32>) -> Result<(), VelesError> {
let point = VelesPoint {
id,
vector: embedding,
payload: None,
};
self.collection.upsert(point)?;
self.contents.write().insert(id, content);
Ok(())
}
pub fn query(
&self,
embedding: Vec<f32>,
top_k: u32,
) -> Result<Vec<SemanticResult>, VelesError> {
let results = self.collection.search(embedding, top_k)?;
let contents = self.contents.read();
Ok(results
.into_iter()
.map(|r| SemanticResult {
id: r.id,
score: r.score,
content: contents.get(&r.id).cloned().unwrap_or_default(),
})
.collect())
}
pub fn len(&self) -> Result<u64, VelesError> {
let contents = self.contents.read();
Ok(contents.len() as u64)
}
pub fn is_empty(&self) -> Result<bool, VelesError> {
Ok(self.len()? == 0)
}
pub fn remove(&self, id: u64) -> Result<bool, VelesError> {
self.collection.delete(id)?;
let mut contents = self.contents.write();
Ok(contents.remove(&id).is_some())
}
pub fn clear(&self) -> Result<(), VelesError> {
let ids: Vec<u64> = {
let contents = self.contents.read();
contents.keys().copied().collect()
};
{
let mut contents = self.contents.write();
contents.clear();
}
for id in ids {
let _ = self.collection.delete(id);
}
Ok(())
}
pub fn dimension(&self) -> u32 {
self.collection.dimension()
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn create_test_db() -> (TempDir, std::sync::Arc<VelesDatabase>) {
let dir = TempDir::new().unwrap();
let db = VelesDatabase::open(dir.path().to_string_lossy().to_string()).unwrap();
(dir, db)
}
#[test]
fn test_semantic_memory_new() {
let (_dir, db) = create_test_db();
let memory = VelesSemanticMemory::new(&db, 4).unwrap();
assert_eq!(memory.dimension(), 4);
assert!(memory.is_empty().unwrap());
}
#[test]
fn test_semantic_memory_store_and_query() {
let (_dir, db) = create_test_db();
let memory = VelesSemanticMemory::new(&db, 4).unwrap();
memory
.store(1, "Test content".to_string(), vec![0.1, 0.2, 0.3, 0.4])
.unwrap();
assert_eq!(memory.len().unwrap(), 1);
assert!(!memory.is_empty().unwrap());
let results = memory.query(vec![0.1, 0.2, 0.3, 0.4], 5).unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].id, 1);
assert_eq!(results[0].content, "Test content");
}
#[test]
fn test_semantic_memory_remove() {
let (_dir, db) = create_test_db();
let memory = VelesSemanticMemory::new(&db, 4).unwrap();
memory
.store(1, "Content".to_string(), vec![0.1, 0.2, 0.3, 0.4])
.unwrap();
assert_eq!(memory.len().unwrap(), 1);
let removed = memory.remove(1).unwrap();
assert!(removed);
assert!(memory.is_empty().unwrap());
}
}