contrag_core/vector_store/
mod.rs

1pub mod stable_memory_store;
2
3use crate::error::Result;
4use crate::types::{Vector, SearchResult};
5
6/// Trait for vector storage backends
7#[async_trait::async_trait]
8pub trait VectorStore: Send + Sync {
9    /// Store a single vector
10    async fn store(&mut self, namespace: &str, vector: Vector) -> Result<()>;
11
12    /// Store multiple vectors
13    async fn store_batch(&mut self, namespace: &str, vectors: Vec<Vector>) -> Result<()> {
14        for vector in vectors {
15            self.store(namespace, vector).await?;
16        }
17        Ok(())
18    }
19
20    /// Search for similar vectors
21    async fn search(
22        &self,
23        namespace: &str,
24        query_embedding: Vec<f32>,
25        k: usize,
26    ) -> Result<Vec<SearchResult>>;
27
28    /// Delete a vector by ID
29    async fn delete(&mut self, namespace: &str, vector_id: &str) -> Result<()>;
30
31    /// Delete all vectors in a namespace
32    async fn delete_namespace(&mut self, namespace: &str) -> Result<()>;
33
34    /// Get vector count in namespace
35    async fn count(&self, namespace: &str) -> Result<usize>;
36
37    /// List all namespaces
38    async fn list_namespaces(&self) -> Result<Vec<String>>;
39}
40
41/// Cosine similarity calculation
42pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
43    if a.len() != b.len() {
44        return 0.0;
45    }
46
47    let dot_product: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
48    let magnitude_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
49    let magnitude_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
50
51    if magnitude_a == 0.0 || magnitude_b == 0.0 {
52        return 0.0;
53    }
54
55    dot_product / (magnitude_a * magnitude_b)
56}
57
58/// Euclidean distance calculation
59pub fn euclidean_distance(a: &[f32], b: &[f32]) -> f32 {
60    if a.len() != b.len() {
61        return f32::MAX;
62    }
63
64    a.iter()
65        .zip(b.iter())
66        .map(|(x, y)| (x - y).powi(2))
67        .sum::<f32>()
68        .sqrt()
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    #[test]
76    fn test_cosine_similarity() {
77        let a = vec![1.0, 0.0, 0.0];
78        let b = vec![1.0, 0.0, 0.0];
79        assert!((cosine_similarity(&a, &b) - 1.0).abs() < 0.001);
80
81        let c = vec![1.0, 0.0, 0.0];
82        let d = vec![0.0, 1.0, 0.0];
83        assert!((cosine_similarity(&c, &d) - 0.0).abs() < 0.001);
84    }
85
86    #[test]
87    fn test_euclidean_distance() {
88        let a = vec![0.0, 0.0, 0.0];
89        let b = vec![1.0, 0.0, 0.0];
90        assert!((euclidean_distance(&a, &b) - 1.0).abs() < 0.001);
91
92        let c = vec![0.0, 0.0, 0.0];
93        let d = vec![1.0, 1.0, 0.0];
94        assert!((euclidean_distance(&c, &d) - 1.414).abs() < 0.01);
95    }
96}
contrag_core/vector_store/mod.rs

contrag_core/vector_store/
mod.rs
