Skip to main content

zeph_memory/
vector_store.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4
5#[derive(Debug, thiserror::Error)]
6pub enum VectorStoreError {
7    #[error("connection error: {0}")]
8    Connection(String),
9    #[error("collection error: {0}")]
10    Collection(String),
11    #[error("upsert error: {0}")]
12    Upsert(String),
13    #[error("search error: {0}")]
14    Search(String),
15    #[error("delete error: {0}")]
16    Delete(String),
17    #[error("scroll error: {0}")]
18    Scroll(String),
19    #[error("serialization error: {0}")]
20    Serialization(String),
21}
22
23#[derive(Debug, Clone)]
24pub struct VectorPoint {
25    pub id: String,
26    pub vector: Vec<f32>,
27    pub payload: HashMap<String, serde_json::Value>,
28}
29
30#[derive(Debug, Clone, Default)]
31pub struct VectorFilter {
32    pub must: Vec<FieldCondition>,
33    pub must_not: Vec<FieldCondition>,
34}
35
36#[derive(Debug, Clone)]
37pub struct FieldCondition {
38    pub field: String,
39    pub value: FieldValue,
40}
41
42#[derive(Debug, Clone)]
43pub enum FieldValue {
44    Integer(i64),
45    Text(String),
46}
47
48#[derive(Debug, Clone)]
49pub struct ScoredVectorPoint {
50    pub id: String,
51    pub score: f32,
52    pub payload: HashMap<String, serde_json::Value>,
53}
54
55type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
56
57pub type ScrollResult = HashMap<String, HashMap<String, String>>;
58
59pub trait VectorStore: Send + Sync {
60    fn ensure_collection(
61        &self,
62        collection: &str,
63        vector_size: u64,
64    ) -> BoxFuture<'_, Result<(), VectorStoreError>>;
65
66    fn collection_exists(&self, collection: &str) -> BoxFuture<'_, Result<bool, VectorStoreError>>;
67
68    fn delete_collection(&self, collection: &str) -> BoxFuture<'_, Result<(), VectorStoreError>>;
69
70    fn upsert(
71        &self,
72        collection: &str,
73        points: Vec<VectorPoint>,
74    ) -> BoxFuture<'_, Result<(), VectorStoreError>>;
75
76    fn search(
77        &self,
78        collection: &str,
79        vector: Vec<f32>,
80        limit: u64,
81        filter: Option<VectorFilter>,
82    ) -> BoxFuture<'_, Result<Vec<ScoredVectorPoint>, VectorStoreError>>;
83
84    fn delete_by_ids(
85        &self,
86        collection: &str,
87        ids: Vec<String>,
88    ) -> BoxFuture<'_, Result<(), VectorStoreError>>;
89
90    fn scroll_all(
91        &self,
92        collection: &str,
93        key_field: &str,
94    ) -> BoxFuture<'_, Result<ScrollResult, VectorStoreError>>;
95}