use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::AtomicBool;
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum VectorStoreError {
#[error("connection error: {0}")]
Connection(String),
#[error("collection error: {0}")]
Collection(String),
#[error("upsert error: {0}")]
Upsert(String),
#[error("search error: {0}")]
Search(String),
#[error("delete error: {0}")]
Delete(String),
#[error("scroll error: {0}")]
Scroll(String),
#[error("serialization error: {0}")]
Serialization(String),
#[error("operation unsupported: {0}")]
Unsupported(String),
}
#[derive(Debug, Clone)]
pub struct VectorPoint {
pub id: String,
pub vector: Vec<f32>,
pub payload: HashMap<String, serde_json::Value>,
}
#[derive(Debug, Clone, Default)]
pub struct VectorFilter {
pub must: Vec<FieldCondition>,
pub must_not: Vec<FieldCondition>,
}
#[derive(Debug, Clone)]
pub struct FieldCondition {
pub field: String,
pub value: FieldValue,
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum FieldValue {
Integer(i64),
Text(String),
}
#[derive(Debug, Clone)]
pub struct ScoredVectorPoint {
pub id: String,
pub score: f32,
pub payload: HashMap<String, serde_json::Value>,
}
pub(crate) type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
pub type ScrollResult = HashMap<String, HashMap<String, String>>;
pub type ScrollWithIdsResult = Vec<(String, HashMap<String, String>)>;
fn clamp_search_limit(site: &'static str, limit: u64, warned: &AtomicBool) -> u64 {
if let Ok(requested) = usize::try_from(limit) {
crate::warn_if_search_limit_clamped(site, requested, warned);
}
limit.clamp(1, crate::MAX_SEARCH_LIMIT as u64)
}
pub trait VectorStore: Send + Sync {
fn ensure_collection(
&self,
collection: &str,
vector_size: u64,
) -> BoxFuture<'_, Result<(), VectorStoreError>>;
fn collection_exists(&self, collection: &str) -> BoxFuture<'_, Result<bool, VectorStoreError>>;
fn delete_collection(&self, collection: &str) -> BoxFuture<'_, Result<(), VectorStoreError>>;
fn upsert(
&self,
collection: &str,
points: Vec<VectorPoint>,
) -> BoxFuture<'_, Result<(), VectorStoreError>>;
fn search(
&self,
collection: &str,
vector: Vec<f32>,
limit: u64,
filter: Option<VectorFilter>,
) -> BoxFuture<'_, Result<Vec<ScoredVectorPoint>, VectorStoreError>> {
static CLAMP_WARNED: AtomicBool = AtomicBool::new(false);
let limit = clamp_search_limit("VectorStore::search", limit, &CLAMP_WARNED);
self.search_clamped(collection, vector, limit, filter)
}
fn search_clamped(
&self,
collection: &str,
vector: Vec<f32>,
limit: u64,
filter: Option<VectorFilter>,
) -> BoxFuture<'_, Result<Vec<ScoredVectorPoint>, VectorStoreError>>;
fn delete_by_ids(
&self,
collection: &str,
ids: Vec<String>,
) -> BoxFuture<'_, Result<(), VectorStoreError>>;
fn scroll_all(
&self,
collection: &str,
key_field: &str,
) -> BoxFuture<'_, Result<ScrollResult, VectorStoreError>>;
fn scroll_all_with_point_ids(
&self,
collection: &str,
key_field: &str,
) -> BoxFuture<'_, Result<ScrollWithIdsResult, VectorStoreError>>;
fn health_check(&self) -> BoxFuture<'_, Result<bool, VectorStoreError>>;
fn create_keyword_indexes(
&self,
_collection: &str,
_fields: &[&str],
) -> BoxFuture<'_, Result<(), VectorStoreError>> {
Box::pin(async { Ok(()) })
}
fn get_points(
&self,
_collection: &str,
_ids: Vec<String>,
) -> BoxFuture<'_, Result<Vec<VectorPoint>, VectorStoreError>> {
Box::pin(async {
Err(VectorStoreError::Unsupported(
"get_points not implemented for this backend".into(),
))
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
struct RecordingStore {
last_limit: Arc<AtomicU64>,
}
impl VectorStore for RecordingStore {
fn ensure_collection(
&self,
_collection: &str,
_vector_size: u64,
) -> BoxFuture<'_, Result<(), VectorStoreError>> {
Box::pin(async { Ok(()) })
}
fn collection_exists(
&self,
_collection: &str,
) -> BoxFuture<'_, Result<bool, VectorStoreError>> {
Box::pin(async { Ok(true) })
}
fn delete_collection(
&self,
_collection: &str,
) -> BoxFuture<'_, Result<(), VectorStoreError>> {
Box::pin(async { Ok(()) })
}
fn upsert(
&self,
_collection: &str,
_points: Vec<VectorPoint>,
) -> BoxFuture<'_, Result<(), VectorStoreError>> {
Box::pin(async { Ok(()) })
}
fn search_clamped(
&self,
_collection: &str,
_vector: Vec<f32>,
limit: u64,
_filter: Option<VectorFilter>,
) -> BoxFuture<'_, Result<Vec<ScoredVectorPoint>, VectorStoreError>> {
self.last_limit.store(limit, Ordering::SeqCst);
Box::pin(async { Ok(vec![]) })
}
fn delete_by_ids(
&self,
_collection: &str,
_ids: Vec<String>,
) -> BoxFuture<'_, Result<(), VectorStoreError>> {
Box::pin(async { Ok(()) })
}
fn scroll_all(
&self,
_collection: &str,
_key_field: &str,
) -> BoxFuture<'_, Result<ScrollResult, VectorStoreError>> {
Box::pin(async { Ok(ScrollResult::new()) })
}
fn scroll_all_with_point_ids(
&self,
_collection: &str,
_key_field: &str,
) -> BoxFuture<'_, Result<ScrollWithIdsResult, VectorStoreError>> {
Box::pin(async { Ok(Vec::new()) })
}
fn health_check(&self) -> BoxFuture<'_, Result<bool, VectorStoreError>> {
Box::pin(async { Ok(true) })
}
}
#[tokio::test]
async fn search_clamps_oversized_limit_before_delegating() {
let last_limit = Arc::new(AtomicU64::new(0));
let store = RecordingStore {
last_limit: last_limit.clone(),
};
store
.search("collection", vec![0.0], u64::MAX, None)
.await
.unwrap();
assert_eq!(
last_limit.load(Ordering::SeqCst),
crate::MAX_SEARCH_LIMIT as u64
);
}
#[tokio::test]
async fn search_passes_small_limit_through_unclamped() {
let last_limit = Arc::new(AtomicU64::new(0));
let store = RecordingStore {
last_limit: last_limit.clone(),
};
store
.search("collection", vec![0.0], 5, None)
.await
.unwrap();
assert_eq!(last_limit.load(Ordering::SeqCst), 5);
}
}