oneai_vector/
in_memory.rs

1//! In-memory brute-force cosine `VectorBackend`.
2//!
3//! No external dependencies and always available (not behind a feature flag).
4//! Intended for tests and tiny document sets — O(n·d) per query is fine up to
5//! ~10⁴ vectors; beyond that switch to [`SqliteVecBackend`](crate::SqliteVecBackend)
6//! (exact KNN) or [`UsearchBackend`](crate::UsearchBackend) (HNSW ANN).
7
8use std::collections::HashMap;
9
10use async_trait::async_trait;
11use oneai_core::traits::{Filter, Metadata, VectorBackend, VectorHit};
12use tokio::sync::Mutex;
13
14use crate::cosine;
15
16/// A brute-force cosine vector backend backed by a `Vec` guarded by a mutex.
17pub struct InMemoryVectorBackend {
18    dim: usize,
19    rows: Mutex<Vec<(String, Vec<f32>, Metadata)>>,
20    by_id: Mutex<HashMap<String, usize>>,
21}
22
23impl InMemoryVectorBackend {
24    /// Create a backend that only accepts vectors of `dim` dimensions.
25    pub fn new(dim: usize) -> Self {
26        Self {
27            dim,
28            rows: Mutex::new(Vec::new()),
29            by_id: Mutex::new(HashMap::new()),
30        }
31    }
32}
33
34#[async_trait]
35impl VectorBackend for InMemoryVectorBackend {
36    async fn upsert(
37        &self,
38        id: &str,
39        embedding: &[f32],
40        metadata: Metadata,
41    ) -> oneai_core::Result<()> {
42        if embedding.len() != self.dim {
43            return Err(oneai_core::OneAIError::Rag(format!(
44                "InMemoryVectorBackend: dim mismatch (got {}, expected {})",
45                embedding.len(),
46                self.dim
47            )));
48        }
49        let mut rows = self.rows.lock().await;
50        let mut by_id = self.by_id.lock().await;
51        if let Some(&idx) = by_id.get(id) {
52            rows[idx].1 = embedding.to_vec();
53            rows[idx].2 = metadata;
54        } else {
55            by_id.insert(id.to_string(), rows.len());
56            rows.push((id.to_string(), embedding.to_vec(), metadata));
57        }
58        Ok(())
59    }
60
61    async fn search(
62        &self,
63        query: &[f32],
64        top_k: usize,
65        filter: Option<&Filter>,
66    ) -> oneai_core::Result<Vec<VectorHit>> {
67        if query.len() != self.dim {
68            return Err(oneai_core::OneAIError::Rag(format!(
69                "InMemoryVectorBackend: query dim mismatch (got {}, expected {})",
70                query.len(),
71                self.dim
72            )));
73        }
74        let rows = self.rows.lock().await;
75        let mut scored: Vec<VectorHit> = rows
76            .iter()
77            .filter(|(_, _, meta)| filter.map(|f| f.matches(meta)).unwrap_or(true))
78            .map(|(id, emb, meta)| VectorHit {
79                id: id.clone(),
80                score: cosine(query, emb),
81                metadata: meta.clone(),
82            })
83            .collect();
84        // Descending by score.
85        scored.sort_by(|a, b| {
86            b.score
87                .partial_cmp(&a.score)
88                .unwrap_or(std::cmp::Ordering::Equal)
89        });
90        scored.truncate(top_k);
91        Ok(scored)
92    }
93
94    async fn delete(&self, id: &str) -> oneai_core::Result<()> {
95        let mut by_id = self.by_id.lock().await;
96        if let Some(idx) = by_id.remove(id) {
97            let mut rows = self.rows.lock().await;
98            // Swap-remove; rebuild id→index map for the moved element.
99            let is_last = idx == rows.len() - 1;
100            rows.swap_remove(idx);
101            if !is_last {
102                // The element that moved into `idx` needs its index fixed.
103                let moved_id = rows[idx].0.clone();
104                by_id.insert(moved_id, idx);
105            }
106        }
107        Ok(())
108    }
109
110    fn dimension(&self) -> usize {
111        self.dim
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    #[tokio::test]
120    async fn upsert_search_filter_delete() {
121        let be = InMemoryVectorBackend::new(4);
122        be.upsert(
123            "a",
124            &[0.1, 0.2, 0.3, 0.4],
125            Metadata::from([("k".into(), "x".into())]),
126        )
127        .await
128        .unwrap();
129        be.upsert(
130            "b",
131            &[0.11, 0.21, 0.31, 0.41],
132            Metadata::from([("k".into(), "y".into())]),
133        )
134        .await
135        .unwrap();
136        be.upsert(
137            "c",
138            &[0.9, 0.8, 0.7, 0.6],
139            Metadata::from([("k".into(), "x".into())]),
140        )
141        .await
142        .unwrap();
143
144        // Nearest to a is b, then c.
145        let hits = be.search(&[0.1, 0.2, 0.3, 0.4], 3, None).await.unwrap();
146        assert_eq!(hits[0].id, "a");
147        assert_eq!(hits[1].id, "b");
148        assert!(hits[0].score > 0.99);
149
150        // Filter k==x → only a and c.
151        let f = Filter::new().with_eq("k", "x");
152        let hits = be.search(&[0.1, 0.2, 0.3, 0.4], 5, Some(&f)).await.unwrap();
153        assert_eq!(hits.len(), 2);
154        assert!(hits.iter().all(|h| h.metadata["k"] == "x"));
155
156        // Upsert replaces.
157        be.upsert("a", &[0.9, 0.8, 0.7, 0.6], Metadata::new())
158            .await
159            .unwrap();
160        let hits = be.search(&[0.9, 0.8, 0.7, 0.6], 2, None).await.unwrap();
161        assert!(hits.iter().any(|h| h.id == "a"));
162
163        be.delete("b").await.unwrap();
164        let hits = be
165            .search(&[0.11, 0.21, 0.31, 0.41], 10, None)
166            .await
167            .unwrap();
168        assert!(!hits.iter().any(|h| h.id == "b"));
169    }
170
171    #[tokio::test]
172    async fn dim_mismatch_errors() {
173        let be = InMemoryVectorBackend::new(4);
174        let err = be.upsert("a", &[0.1, 0.2, 0.3], Metadata::new()).await;
175        assert!(err.is_err());
176        assert!(be.search(&[0.1, 0.2, 0.3], 1, None).await.is_err());
177    }
178}
oneai_vector/in_memory.rs

oneai_vector/
in_memory.rs
