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}