wesichain_core/
vector_store.rs

1use async_trait::async_trait;
2
3use crate::{Document, MetadataFilter, StoreError};
4
5#[derive(Clone, Debug)]
6pub struct SearchResult {
7    pub document: Document,
8    pub score: f32,
9}
10
11#[async_trait]
12pub trait VectorStore: Send + Sync {
13    async fn add(&self, docs: Vec<Document>) -> Result<(), StoreError>;
14    async fn search(
15        &self,
16        query_embedding: &[f32],
17        top_k: usize,
18        filter: Option<&MetadataFilter>,
19    ) -> Result<Vec<SearchResult>, StoreError>;
20    async fn delete(&self, ids: &[String]) -> Result<(), StoreError>;
21
22    async fn delete_strs(&self, ids: &[&str]) -> Result<(), StoreError>
23    where
24        Self: Sized,
25    {
26        let owned: Vec<String> = ids.iter().map(|id| (*id).to_string()).collect();
27        self.delete(&owned).await
28    }
29
30    async fn delete_ref<T>(&self, ids: &[T]) -> Result<(), StoreError>
31    where
32        T: AsRef<str> + Sync,
33        Self: Sized,
34    {
35        let owned: Vec<String> = ids.iter().map(|id| id.as_ref().to_string()).collect();
36        self.delete(&owned).await
37    }
38}
39
40#[async_trait]
41impl<T: VectorStore + ?Sized> VectorStore for std::sync::Arc<T> {
42    async fn add(&self, docs: Vec<Document>) -> Result<(), StoreError> {
43        self.as_ref().add(docs).await
44    }
45
46    async fn search(
47        &self,
48        query_embedding: &[f32],
49        top_k: usize,
50        filter: Option<&MetadataFilter>,
51    ) -> Result<Vec<SearchResult>, StoreError> {
52        self.as_ref().search(query_embedding, top_k, filter).await
53    }
54
55    async fn delete(&self, ids: &[String]) -> Result<(), StoreError> {
56        self.as_ref().delete(ids).await
57    }
58}
59
60pub async fn delete_strs_dyn(store: &dyn VectorStore, ids: &[&str]) -> Result<(), StoreError> {
61    let owned: Vec<String> = ids.iter().map(|id| (*id).to_string()).collect();
62    store.delete(&owned).await
63}
64
65pub async fn delete_ref_dyn<T: AsRef<str> + Sync>(
66    store: &dyn VectorStore,
67    ids: &[T],
68) -> Result<(), StoreError> {
69    let owned: Vec<String> = ids.iter().map(|id| id.as_ref().to_string()).collect();
70    store.delete(&owned).await
71}
wesichain_core/vector_store.rs

wesichain_core/
vector_store.rs
