Skip to main content

daimon_core/
vector_store.rs

1//! Trait-based plugin system for vector stores.
2//!
3//! Implement [`VectorStore`] for your storage backend (pgvector, Qdrant,
4//! Pinecone, Chroma, Weaviate, Milvus, etc.), then compose with a knowledge
5//! base for a complete embedding-backed retrieval pipeline.
6
7use std::future::Future;
8use std::pin::Pin;
9use std::sync::Arc;
10
11use crate::document::{Document, ScoredDocument};
12use crate::error::Result;
13
14/// Low-level vector storage interface.
15///
16/// Implement this trait for your vector database. Operations deal in
17/// pre-computed embeddings — combine with an `EmbeddingModel` and knowledge
18/// base for automatic embedding computation.
19pub trait VectorStore: Send + Sync {
20    /// Inserts or updates a document with a pre-computed embedding vector.
21    fn upsert(
22        &self,
23        id: &str,
24        embedding: Vec<f32>,
25        document: Document,
26    ) -> impl Future<Output = Result<()>> + Send;
27
28    /// Queries the store for the `top_k` most similar documents to the
29    /// given embedding vector. Returns results sorted by descending similarity.
30    fn query(
31        &self,
32        embedding: Vec<f32>,
33        top_k: usize,
34    ) -> impl Future<Output = Result<Vec<ScoredDocument>>> + Send;
35
36    /// Deletes a document by ID. Returns `true` if the document was found
37    /// and deleted.
38    fn delete(&self, id: &str) -> impl Future<Output = Result<bool>> + Send;
39
40    /// Returns the number of documents in the store.
41    fn count(&self) -> impl Future<Output = Result<usize>> + Send;
42}
43
44/// Object-safe wrapper for [`VectorStore`].
45pub 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
97/// Shared ownership of a vector store.
98pub type SharedVectorStore = Arc<dyn ErasedVectorStore>;