do_memory_core/embeddings/
index.rs

1//! Vector index abstractions and implementations.
2
3use crate::embeddings::similarity::cosine_similarity;
4use crate::error::Result;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use std::path::Path;
8
9/// A hit from a vector search.
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct VectorHit {
12    /// The ID of the vector.
13    pub id: String,
14    /// The similarity score.
15    pub score: f32,
16}
17
18/// Trait for vector indexing and similarity search.
19pub trait VectorIndex: Send + Sync {
20    /// Add or update a vector in the index.
21    fn upsert(&mut self, id: &str, embedding: &[f32]) -> Result<()>;
22
23    /// Remove a vector from the index.
24    fn remove(&mut self, id: &str) -> Result<()>;
25
26    /// Search for the top-k most similar vectors.
27    fn search(&self, query: &[f32], top_k: usize) -> Result<Vec<VectorHit>>;
28
29    /// Save the index to a file.
30    fn save(&self, path: &Path) -> Result<()>;
31
32    /// Get the number of vectors in the index.
33    fn len(&self) -> usize;
34
35    /// Check if the index is empty.
36    fn is_empty(&self) -> bool {
37        self.len() == 0
38    }
39}
40
41/// A simple brute-force vector index.
42#[derive(Debug, Default, Serialize, Deserialize, Clone)]
43pub struct SimpleVectorIndex {
44    vectors: HashMap<String, Vec<f32>>,
45}
46
47impl SimpleVectorIndex {
48    /// Create a new empty SimpleVectorIndex.
49    pub fn new() -> Self {
50        Self::default()
51    }
52
53    /// Load a SimpleVectorIndex from a file.
54    pub fn load(path: &Path) -> Result<Self> {
55        let file = std::fs::File::open(path)?;
56        let index: Self = serde_json::from_reader(file)?;
57        Ok(index)
58    }
59}
60
61impl VectorIndex for SimpleVectorIndex {
62    fn upsert(&mut self, id: &str, embedding: &[f32]) -> Result<()> {
63        self.vectors.insert(id.to_string(), embedding.to_vec());
64        Ok(())
65    }
66
67    fn remove(&mut self, id: &str) -> Result<()> {
68        self.vectors.remove(id);
69        Ok(())
70    }
71
72    fn search(&self, query: &[f32], top_k: usize) -> Result<Vec<VectorHit>> {
73        let mut hits: Vec<VectorHit> = self
74            .vectors
75            .iter()
76            .map(|(id, vec)| {
77                let score = cosine_similarity(query, vec);
78                VectorHit {
79                    id: id.clone(),
80                    score,
81                }
82            })
83            .collect();
84
85        // Sort by score descending
86        hits.sort_by(|a, b| {
87            b.score
88                .partial_cmp(&a.score)
89                .unwrap_or(std::cmp::Ordering::Equal)
90        });
91
92        // Take top-k
93        hits.truncate(top_k);
94
95        Ok(hits)
96    }
97
98    fn save(&self, path: &Path) -> Result<()> {
99        if let Some(parent) = path.parent() {
100            std::fs::create_dir_all(parent)?;
101        }
102        let file = std::fs::File::create(path)?;
103        serde_json::to_writer(file, self)?;
104        Ok(())
105    }
106
107    fn len(&self) -> usize {
108        self.vectors.len()
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use tempfile::tempdir;
116
117    #[test]
118    fn test_simple_vector_index_search() {
119        let mut index = SimpleVectorIndex::new();
120        index.upsert("1", &[1.0, 0.0, 0.0]).unwrap();
121        index.upsert("2", &[0.0, 1.0, 0.0]).unwrap();
122        index.upsert("3", &[0.5, 0.5, 0.0]).unwrap();
123
124        let query = [1.0, 0.1, 0.0];
125        let hits = index.search(&query, 2).unwrap();
126
127        assert_eq!(hits.len(), 2);
128        assert_eq!(hits[0].id, "1");
129        assert_eq!(hits[1].id, "3");
130    }
131
132    #[test]
133    fn test_simple_vector_index_persistence() {
134        let dir = tempdir().unwrap();
135        let path = dir.path().join("index.json");
136
137        let mut index = SimpleVectorIndex::new();
138        index.upsert("1", &[1.0, 0.0]).unwrap();
139        index.save(&path).unwrap();
140
141        let loaded = SimpleVectorIndex::load(&path).unwrap();
142        assert_eq!(loaded.len(), 1);
143
144        let hits = loaded.search(&[1.0, 0.0], 1).unwrap();
145        assert_eq!(hits[0].id, "1");
146    }
147
148    #[test]
149    fn test_simple_vector_index_remove() {
150        let mut index = SimpleVectorIndex::new();
151        index.upsert("1", &[1.0, 0.0]).unwrap();
152        assert_eq!(index.len(), 1);
153        index.remove("1").unwrap();
154        assert_eq!(index.len(), 0);
155    }
156}
do_memory_core/embeddings/index.rs

do_memory_core/embeddings/
index.rs
