use anyhow::Result;
use async_trait::async_trait;
use uuid::Uuid;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use serde::{Deserialize, Serialize};
use rocksdb::{DB, Options, WriteBatch};
use super::traits::{VectorStorage, VectorPayload, SearchFilter, SearchHit};
const CF_VECTORS: &str = "vectors";
const CF_PAYLOADS: &str = "payloads";
#[derive(Debug, Clone, Serialize, Deserialize)]
struct StoredVector {
embedding: Vec<f32>,
payload: VectorPayload,
}
pub struct RocksDBVectorStore {
db: Arc<DB>,
dimension: usize,
cache: Arc<Mutex<HashMap<Uuid, StoredVector>>>,
dirty: Arc<Mutex<bool>>,
}
impl RocksDBVectorStore {
pub fn open(path: &std::path::Path, dimension: usize) -> Result<Self> {
let mut opts = Options::default();
opts.create_if_missing(true);
opts.create_missing_column_families(true);
let cf_vectors = rocksdb::ColumnFamilyDescriptor::new(CF_VECTORS, Options::default());
let cf_payloads = rocksdb::ColumnFamilyDescriptor::new(CF_PAYLOADS, Options::default());
let db = DB::open_cf_descriptors(&opts, path, vec![cf_vectors, cf_payloads])?;
let cache = Arc::new(Mutex::new(HashMap::new()));
let db = Arc::new(db);
{
let cf = db.cf_handle(CF_VECTORS).expect("CF_VECTORS not found");
let mut cache_lock = cache.lock().unwrap();
for (key, value) in db.iterator_cf(&cf, rocksdb::IteratorMode::Start).flatten() {
if let Ok(id_str) = std::str::from_utf8(&key) {
if let Ok(id) = Uuid::parse_str(id_str) {
if let Ok(stored) = serde_json::from_slice::<StoredVector>(&value) {
cache_lock.insert(id, stored);
}
}
}
}
}
Ok(Self {
db,
dimension,
cache,
dirty: Arc::new(Mutex::new(false)),
})
}
pub fn save(&self) -> Result<()> {
let dirty = *self.dirty.lock().unwrap();
if !dirty {
return Ok(());
}
let cache = self.cache.lock().unwrap();
let cf_vectors = self.db.cf_handle(CF_VECTORS).expect("CF_VECTORS not found");
let mut batch = WriteBatch::default();
for (id, stored) in cache.iter() {
let key = id.to_string();
let value = serde_json::to_vec(stored)?;
batch.put_cf(&cf_vectors, key.as_bytes(), &value);
}
self.db.write(batch)?;
*self.dirty.lock().unwrap() = false;
Ok(())
}
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
let dot = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum::<f32>();
let norm_a = (a.iter().map(|x| x * x).sum::<f32>()).sqrt();
let norm_b = (b.iter().map(|x| x * x).sum::<f32>()).sqrt();
if norm_a == 0.0 || norm_b == 0.0 {
0.0
} else {
dot / (norm_a * norm_b)
}
}
pub fn count_cached(&self) -> usize {
self.cache.lock().unwrap().len()
}
}
#[async_trait]
impl VectorStorage for RocksDBVectorStore {
async fn add(&self, id: &Uuid, embedding: &[f32], payload: VectorPayload) -> Result<()> {
if embedding.len() != self.dimension {
return Err(anyhow::anyhow!(
"Embedding dimension mismatch: expected {}, got {}",
self.dimension,
embedding.len()
));
}
let stored = StoredVector {
embedding: embedding.to_vec(),
payload,
};
self.cache.lock().unwrap().insert(*id, stored);
*self.dirty.lock().unwrap() = true;
Ok(())
}
async fn remove(&self, id: &Uuid) -> Result<bool> {
let existed = self.cache.lock().unwrap().remove(id).is_some();
if existed {
*self.dirty.lock().unwrap() = true;
let cf_vectors = self.db.cf_handle(CF_VECTORS).expect("CF_VECTORS not found");
self.db.delete_cf(&cf_vectors, id.to_string().as_bytes())?;
}
Ok(existed)
}
async fn search(&self, query: &[f32], filter: SearchFilter, top_k: usize) -> Result<Vec<SearchHit>> {
if query.len() != self.dimension {
return Err(anyhow::anyhow!(
"Query dimension mismatch: expected {}, got {}",
self.dimension,
query.len()
));
}
let cache = self.cache.lock().unwrap();
let mut similarities: Vec<(Uuid, f32, &StoredVector)> = cache.iter()
.filter_map(|(id, stored)| {
if filter.project_id.is_some() && stored.payload.project_id != filter.project_id
&& (!filter.include_global || stored.payload.project_id != Some(Uuid::nil())) {
return None;
}
if filter.memory_type.is_some() && Some(stored.payload.memory_type.as_str()) != filter.memory_type.as_deref() {
return None;
}
let sim = Self::cosine_similarity(query, &stored.embedding);
if sim < filter.min_score {
return None;
}
Some((*id, sim, stored))
})
.collect();
similarities.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
similarities.truncate(top_k);
let hits = similarities.into_iter()
.map(|(id, score, stored)| SearchHit {
memory_id: id,
score,
payload: stored.payload.clone(),
})
.collect();
Ok(hits)
}
async fn get(&self, id: &Uuid) -> Result<Option<Vec<f32>>> {
Ok(self.cache.lock().unwrap().get(id).map(|s| s.embedding.clone()))
}
async fn count(&self) -> Result<usize> {
Ok(self.cache.lock().unwrap().len())
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[tokio::test]
async fn test_rocksdb_vector_store() {
let dir = tempdir().unwrap();
let id = Uuid::new_v4();
let embedding = vec![1.0, 2.0, 3.0];
let payload = VectorPayload {
content_preview: "test".to_string(),
..Default::default()
};
{
let store = RocksDBVectorStore::open(dir.path(), 3).unwrap();
store.add(&id, &embedding, payload).await.unwrap();
let retrieved = store.get(&id).await.unwrap();
assert!(retrieved.is_some());
assert_eq!(retrieved.unwrap(), embedding);
store.save().unwrap();
}
let store2 = RocksDBVectorStore::open(dir.path(), 3).unwrap();
let retrieved2 = store2.get(&id).await.unwrap();
assert!(retrieved2.is_some());
assert_eq!(retrieved2.unwrap(), embedding);
}
}