daimon_core/
vector_store.rs1use std::future::Future;
8use std::pin::Pin;
9use std::sync::Arc;
10
11use crate::document::{Document, ScoredDocument};
12use crate::error::Result;
13
14pub trait VectorStore: Send + Sync {
20 fn upsert(
22 &self,
23 id: &str,
24 embedding: Vec<f32>,
25 document: Document,
26 ) -> impl Future<Output = Result<()>> + Send;
27
28 fn query(
31 &self,
32 embedding: Vec<f32>,
33 top_k: usize,
34 ) -> impl Future<Output = Result<Vec<ScoredDocument>>> + Send;
35
36 fn delete(&self, id: &str) -> impl Future<Output = Result<bool>> + Send;
39
40 fn count(&self) -> impl Future<Output = Result<usize>> + Send;
42}
43
44pub trait ErasedVectorStore: Send + Sync {
46 fn upsert_erased<'a>(
47 &'a self,
48 id: &'a str,
49 embedding: Vec<f32>,
50 document: Document,
51 ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
52
53 fn query_erased<'a>(
54 &'a self,
55 embedding: Vec<f32>,
56 top_k: usize,
57 ) -> Pin<Box<dyn Future<Output = Result<Vec<ScoredDocument>>> + Send + 'a>>;
58
59 fn delete_erased<'a>(
60 &'a self,
61 id: &'a str,
62 ) -> Pin<Box<dyn Future<Output = Result<bool>> + Send + 'a>>;
63
64 fn count_erased(&self) -> Pin<Box<dyn Future<Output = Result<usize>> + Send + '_>>;
65}
66
67impl<T: VectorStore> ErasedVectorStore for T {
68 fn upsert_erased<'a>(
69 &'a self,
70 id: &'a str,
71 embedding: Vec<f32>,
72 document: Document,
73 ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
74 Box::pin(self.upsert(id, embedding, document))
75 }
76
77 fn query_erased<'a>(
78 &'a self,
79 embedding: Vec<f32>,
80 top_k: usize,
81 ) -> Pin<Box<dyn Future<Output = Result<Vec<ScoredDocument>>> + Send + 'a>> {
82 Box::pin(self.query(embedding, top_k))
83 }
84
85 fn delete_erased<'a>(
86 &'a self,
87 id: &'a str,
88 ) -> Pin<Box<dyn Future<Output = Result<bool>> + Send + 'a>> {
89 Box::pin(self.delete(id))
90 }
91
92 fn count_erased(&self) -> Pin<Box<dyn Future<Output = Result<usize>> + Send + '_>> {
93 Box::pin(self.count())
94 }
95}
96
97pub type SharedVectorStore = Arc<dyn ErasedVectorStore>;