Skip to main content

khive_storage/
vectors.rs

1//! Vector embedding storage and similarity search capability.
2
3use std::collections::HashSet;
4use std::sync::OnceLock;
5
6use async_trait::async_trait;
7use uuid::Uuid;
8
9use khive_types::SubstrateKind;
10
11use crate::capability::StorageCapability;
12use crate::error::StorageError;
13use crate::types::{
14    BatchWriteSummary, IndexRebuildScope, OrphanSweepConfig, OrphanSweepResult, StorageResult,
15    VectorMetadataFilter, VectorRecord, VectorSearchHit, VectorSearchRequest,
16    VectorStoreCapabilities, VectorStoreInfo,
17};
18
19/// Storage capability for dense vector embeddings and similarity search.
20#[async_trait]
21pub trait VectorStore: Send + Sync + 'static {
22    // --- Required methods ---
23
24    /// Store one or more dense vectors for a subject, identified by field name.
25    async fn insert(
26        &self,
27        subject_id: Uuid,
28        kind: SubstrateKind,
29        namespace: &str,
30        field: &str,
31        vectors: Vec<Vec<f32>>,
32    ) -> StorageResult<()>;
33    /// Insert a batch of pre-assembled vector records in one call.
34    async fn insert_batch(&self, records: Vec<VectorRecord>) -> StorageResult<BatchWriteSummary>;
35    /// Delete all vectors associated with the given subject ID.
36    async fn delete(&self, subject_id: Uuid) -> StorageResult<bool>;
37    /// Return the total number of vector entries in this store.
38    async fn count(&self) -> StorageResult<u64>;
39    /// Run approximate nearest-neighbor search and return ranked hits.
40    async fn search(&self, request: VectorSearchRequest) -> StorageResult<Vec<VectorSearchHit>>;
41    /// Return index metadata and health statistics for this backend.
42    async fn info(&self) -> StorageResult<VectorStoreInfo>;
43    /// Rebuild the ANN index, optionally scoped to a subset of entries.
44    async fn rebuild(&self, scope: IndexRebuildScope) -> StorageResult<VectorStoreInfo>;
45
46    // --- New methods (default impls; backends opt in by overriding) ---
47
48    /// Declare what this backend supports (called at runtime policy construction).
49    ///
50    /// Default returns a conservative, backend-neutral baseline with all optional
51    /// features disabled and no advertised dimension ceiling or index kind
52    /// (STORAGE-AUD-001, ADR-044, ADR-071 Phase 1). Backends that support filter
53    /// pushdown, batch search, quantization, in-place update, or that have a known
54    /// dimension ceiling or index kind should override this and return their own
55    /// `&'static VectorStoreCapabilities`.
56    fn capabilities(&self) -> &'static VectorStoreCapabilities {
57        static BASELINE: OnceLock<VectorStoreCapabilities> = OnceLock::new();
58        BASELINE.get_or_init(|| VectorStoreCapabilities {
59            supports_filter: false,
60            supports_batch_search: false,
61            supports_quantization: false,
62            supports_update: false,
63            supports_orphan_sweep: false,
64            supports_multi_field: false,
65            // Backend-neutral baseline: unknown dimension ceiling and no
66            // advertised index kind. Backends with a concrete limit (e.g.
67            // SqliteVecStore) must override capabilities().
68            max_dimensions: None,
69            index_kinds: vec![],
70        })
71    }
72
73    /// Search with metadata pre-filter.
74    ///
75    /// Default: delegates to [`Self::search`] when the filter carries no predicates;
76    /// returns [`StorageError::Unsupported`] otherwise. Backends with native filter
77    /// pushdown should override this method and set `supports_filter = true` in their
78    /// [`VectorStoreCapabilities`].
79    ///
80    /// Callers must check `capabilities().supports_filter` before calling; the
81    /// runtime layer is responsible for post-filtering when native pushdown is absent.
82    ///
83    /// A backend that claims `supports_filter = true` but does not override this
84    /// method will trigger a `debug_assert` at runtime.
85    async fn search_with_filter(
86        &self,
87        request: &VectorSearchRequest,
88        filter: &VectorMetadataFilter,
89    ) -> StorageResult<Vec<VectorSearchHit>> {
90        if filter.is_empty() {
91            return self.search(request.clone()).await;
92        }
93        debug_assert!(
94            !self.capabilities().supports_filter,
95            "backend claims supports_filter=true but did not override search_with_filter"
96        );
97        Err(StorageError::Unsupported {
98            capability: StorageCapability::Vectors,
99            operation: "search_with_filter".into(),
100            message: "filter pushdown not supported; set supports_filter=true only when overriding this method".into(),
101        })
102    }
103
104    /// Search with N query vectors in one round-trip (HyDE fan-out, multi-query).
105    ///
106    /// Default: sequential calls to [`Self::search`], isolating per-query errors so one
107    /// bad request does not abort the batch. Backends that support native batch
108    /// search should override this and set `supports_batch_search = true`.
109    async fn search_batch(
110        &self,
111        requests: &[VectorSearchRequest],
112    ) -> StorageResult<Vec<StorageResult<Vec<VectorSearchHit>>>> {
113        let mut out = Vec::with_capacity(requests.len());
114        for req in requests {
115            out.push(self.search(req.clone()).await);
116        }
117        Ok(out)
118    }
119
120    /// Re-embed an existing entry in place.
121    ///
122    /// Default: delete then insert. Backends that support atomic in-place update
123    /// should override this and set `supports_update = true` in their
124    /// [`VectorStoreCapabilities`].
125    async fn update(
126        &self,
127        subject_id: Uuid,
128        kind: SubstrateKind,
129        namespace: &str,
130        field: &str,
131        vectors: Vec<Vec<f32>>,
132    ) -> StorageResult<()> {
133        self.delete(subject_id).await?;
134        self.insert(subject_id, kind, namespace, field, vectors)
135            .await
136    }
137
138    /// Remove vectors with no live subject (orphan sweep).
139    ///
140    /// Default returns [`StorageError::Unsupported`]. Backends that implement
141    /// deletion must set `supports_orphan_sweep = true` and override this method.
142    async fn orphan_sweep(&self, config: &OrphanSweepConfig) -> StorageResult<OrphanSweepResult> {
143        let _ = config;
144        Err(StorageError::Unsupported {
145            capability: StorageCapability::Vectors,
146            operation: "orphan_sweep".into(),
147            message: "this backend does not support orphan sweep".into(),
148        })
149    }
150
151    /// Check which of the given subject IDs already have embeddings in this store
152    /// for the specified namespace.
153    ///
154    /// Returns a [`HashSet`] of IDs that are present. IDs not in the returned set
155    /// have no embedding. Default returns [`StorageError::Unsupported`]; backends
156    /// that support fast bulk existence checks should override this method.
157    async fn batch_exists(&self, ids: &[Uuid], namespace: &str) -> StorageResult<HashSet<Uuid>> {
158        let _ = (ids, namespace);
159        Err(StorageError::Unsupported {
160            capability: StorageCapability::Vectors,
161            operation: "batch_exists".into(),
162            message: "this backend does not support batch existence checks".into(),
163        })
164    }
165
166    /// Delete all rows for the given subject IDs, regardless of their stored namespace.
167    ///
168    /// This is a namespace-agnostic sweep — it removes every vector row whose
169    /// `subject_id` matches, no matter which namespace the row was written under.
170    /// Required when the vec table's PRIMARY KEY is `subject_id` alone (not
171    /// `(subject_id, namespace)`): a row from a prior namespace would collide on
172    /// re-insert after a relabel, so the pre-insert drop must target by subject
173    /// only. Returns the number of rows deleted across all chunks.
174    ///
175    /// Default returns [`StorageError::Unsupported`]; backends that store vectors
176    /// in a per-subject keyed table should override this method.
177    async fn delete_subjects(&self, ids: &[Uuid]) -> StorageResult<u64> {
178        let _ = ids;
179        Err(StorageError::Unsupported {
180            capability: StorageCapability::Vectors,
181            operation: "delete_subjects".into(),
182            message: "this backend does not support namespace-agnostic subject deletion".into(),
183        })
184    }
185}