use async_trait::async_trait;
use crate::errors::VectorStoreError;
use crate::models::{Filters, MemoryRecord, Payload, ScoredMemory};
#[derive(Debug, Clone)]
pub struct VectorSearchResult {
pub id: String,
pub score: f32,
pub payload: Payload,
}
#[async_trait]
pub trait VectorStore: Send + Sync {
async fn insert(
&self,
id: &str,
embedding: Vec<f32>,
payload: Payload,
) -> Result<(), VectorStoreError>;
async fn search(
&self,
embedding: &[f32],
limit: usize,
filters: Option<&Filters>,
) -> Result<Vec<VectorSearchResult>, VectorStoreError>;
async fn get(&self, id: &str) -> Result<Option<VectorSearchResult>, VectorStoreError>;
async fn delete(&self, id: &str) -> Result<(), VectorStoreError>;
async fn update(
&self,
id: &str,
embedding: Option<Vec<f32>>,
payload: Payload,
) -> Result<(), VectorStoreError>;
async fn list(
&self,
filters: Option<&Filters>,
limit: usize,
) -> Result<Vec<VectorSearchResult>, VectorStoreError>;
async fn delete_all(&self, filters: Option<&Filters>) -> Result<usize, VectorStoreError>;
async fn collection_exists(&self) -> Result<bool, VectorStoreError>;
async fn create_collection(&self) -> Result<(), VectorStoreError>;
}
impl VectorSearchResult {
pub fn to_memory_record(&self) -> MemoryRecord {
MemoryRecord {
id: uuid::Uuid::parse_str(&self.id).unwrap_or_else(|_| uuid::Uuid::new_v4()),
content: self.payload.data.clone(),
metadata: self.payload.metadata.clone(),
user_id: self.payload.user_id.clone(),
agent_id: self.payload.agent_id.clone(),
run_id: self.payload.run_id.clone(),
hash: self.payload.hash.clone(),
created_at: self.payload.created_at,
updated_at: self.payload.created_at, }
}
pub fn to_scored_memory(&self) -> ScoredMemory {
ScoredMemory {
record: self.to_memory_record(),
score: self.score,
}
}
}