memory_mcp/index/
mod.rs

1//! Vector index module — defines the `VectorStore` trait and provides
2//! concrete implementations backed by usearch (HNSW) and an in-memory
3//! HashMap (for tests).
4
5use std::path::Path;
6
7use crate::{
8    error::MemoryError,
9    types::{Scope, ScopeFilter},
10};
11
12/// HashMap-based `VectorStore` implementation for tests.
13pub mod in_memory;
14/// Usearch HNSW-backed `VectorStore` implementation.
15pub mod usearch;
16
17pub use in_memory::InMemoryStore;
18pub use usearch::UsearchStore;
19
20// ---------------------------------------------------------------------------
21// Sealed trait — prevents external implementations of VectorStore
22// ---------------------------------------------------------------------------
23
24pub(crate) mod sealed {
25    pub trait Sealed {}
26}
27
28// ---------------------------------------------------------------------------
29// VectorStore trait
30// ---------------------------------------------------------------------------
31
32/// A pluggable vector similarity store.
33///
34/// Implementations must be `Send + Sync` so they can be shared across async
35/// tasks and placed behind an `Arc` or `Box`.
36///
37/// # Object safety
38/// The trait is object-safe: `load` (which would return `Self`) is intentionally
39/// absent. Each implementation provides its own constructor.
40pub trait VectorStore: Send + Sync + sealed::Sealed {
41    /// Insert or upsert `vector` for `qualified_name` in the given `scope`.
42    ///
43    /// Returns the key assigned to the entry in the global "all" index.
44    fn add(
45        &self,
46        scope: &Scope,
47        vector: &[f32],
48        qualified_name: String,
49    ) -> Result<u64, MemoryError>;
50
51    /// Remove the entry for `qualified_name` from `scope` (and from the
52    /// all-index). Best-effort — does not fail if the entry is absent.
53    fn remove(&self, scope: &Scope, qualified_name: &str) -> Result<(), MemoryError>;
54
55    /// Search for the `limit` nearest neighbours of `query`, filtered by
56    /// `filter`.
57    ///
58    /// Returns `(key, qualified_name, distance)` triples sorted by ascending
59    /// distance (lower = more similar).
60    fn search(
61        &self,
62        filter: &ScopeFilter,
63        query: &[f32],
64        limit: usize,
65    ) -> Result<Vec<(u64, String, f32)>, MemoryError>;
66
67    /// Look up the vector key for a qualified name in the all-index.
68    ///
69    /// Returns `None` if the name is not indexed.
70    fn find_by_name(&self, qualified_name: &str) -> Option<u64>;
71
72    /// Persist all indexes to subdirectories under `dir`.
73    fn save(&self, dir: &Path) -> Result<(), MemoryError>;
74
75    /// Returns `true` when the store is ready to accept queries.
76    ///
77    /// For `UsearchStore` this is always `true` after construction. For
78    /// `InMemoryStore` it returns the configured value (useful to simulate
79    /// a not-yet-ready backend in tests).
80    fn is_ready(&self) -> bool;
81
82    /// The embedding dimensionality this store was initialised with.
83    fn dimensions(&self) -> usize;
84
85    /// The commit SHA last written to or read from the index metadata, if any.
86    fn commit_sha(&self) -> Option<String>;
87
88    /// Overwrite the stored commit SHA.
89    fn set_commit_sha(&self, sha: Option<&str>);
90}
91
92// ---------------------------------------------------------------------------
93// Trait-level tests — run against both implementations
94// ---------------------------------------------------------------------------
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use crate::index::{InMemoryStore, UsearchStore};
100
101    fn vec_a() -> Vec<f32> {
102        vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
103    }
104
105    fn vec_b() -> Vec<f32> {
106        vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
107    }
108
109    /// Run all trait-level contract tests against a given `VectorStore`.
110    fn check_contract(store: &dyn VectorStore) {
111        let scope = Scope::Global;
112        let name = "global/contract-test".to_string();
113
114        // TC-02a: add + find_by_name returns Some
115        store
116            .add(&scope, &vec_a(), name.clone())
117            .expect("add should succeed");
118        assert!(
119            store.find_by_name(&name).is_some(),
120            "TC-02a: find_by_name should return Some after add"
121        );
122
123        // TC-02a: search returns the added entry
124        let results = store
125            .search(&ScopeFilter::GlobalOnly, &vec_a(), 5)
126            .expect("search should succeed");
127        assert!(
128            results.iter().any(|(_, n, _)| n == &name),
129            "TC-02a: search should return added entry"
130        );
131
132        // TC-02c: upsert replaces old entry — only one result with that name
133        store
134            .add(&scope, &vec_b(), name.clone())
135            .expect("upsert should succeed");
136        let results = store
137            .search(&ScopeFilter::All, &vec_b(), 10)
138            .expect("search after upsert should succeed");
139        assert_eq!(
140            results.iter().filter(|(_, n, _)| n == &name).count(),
141            1,
142            "TC-02c: upsert should leave exactly one entry"
143        );
144
145        // TC-02b: remove makes entry unreachable
146        store.remove(&scope, &name).expect("remove should succeed");
147        assert!(
148            store.find_by_name(&name).is_none(),
149            "TC-02b: find_by_name should return None after remove"
150        );
151        let results_after = store
152            .search(&ScopeFilter::GlobalOnly, &vec_a(), 5)
153            .expect("search after remove should succeed");
154        assert!(
155            !results_after.iter().any(|(_, n, _)| n == &name),
156            "TC-02b: search should not return removed entry"
157        );
158
159        // TC-02d: search with ProjectAndGlobal returns correct entries.
160        let proj_scope = Scope::Project("testproj".to_string());
161        store
162            .add(
163                &Scope::Global,
164                &vec_a(),
165                "global/contract-global".to_string(),
166            )
167            .expect("re-add global entry for TC-02d");
168        store
169            .add(
170                &proj_scope,
171                &vec_b(),
172                "projects/testproj/contract-proj".to_string(),
173            )
174            .expect("add project entry should succeed");
175        let pag_results = store
176            .search(
177                &ScopeFilter::ProjectAndGlobal("testproj".to_string()),
178                &vec_a(),
179                10,
180            )
181            .expect("ProjectAndGlobal search should succeed");
182        let pag_names: Vec<&str> = pag_results.iter().map(|(_, n, _)| n.as_str()).collect();
183        assert!(
184            pag_names
185                .iter()
186                .any(|n| *n == "projects/testproj/contract-proj"),
187            "TC-02d: ProjectAndGlobal should include project entries"
188        );
189        assert!(
190            pag_names.iter().any(|n| *n == "global/contract-global"),
191            "TC-02d: ProjectAndGlobal should include global entries"
192        );
193        // Clean up
194        store
195            .remove(&proj_scope, "projects/testproj/contract-proj")
196            .expect("remove project entry");
197        store
198            .remove(&Scope::Global, "global/contract-global")
199            .expect("remove global entry");
200
201        // TC-06: is_ready() returns true for a freshly created store.
202        assert!(
203            store.is_ready(),
204            "TC-06: is_ready() should return true for a functioning store"
205        );
206
207        // dimensions() returns the value the store was created with.
208        assert_eq!(
209            store.dimensions(),
210            8,
211            "dimensions() should return 8 (the value passed to new)"
212        );
213
214        // commit_sha / set_commit_sha round-trip.
215        assert!(
216            store.commit_sha().is_none(),
217            "commit_sha() should be None on a fresh store"
218        );
219        store.set_commit_sha(Some("deadbeef"));
220        assert_eq!(
221            store.commit_sha(),
222            Some("deadbeef".to_string()),
223            "commit_sha() should reflect set_commit_sha(Some(...))"
224        );
225        store.set_commit_sha(None);
226        assert!(
227            store.commit_sha().is_none(),
228            "commit_sha() should be None after set_commit_sha(None)"
229        );
230    }
231
232    #[test]
233    fn trait_contract_usearch_store() {
234        let store = UsearchStore::new(8).expect("create UsearchStore");
235        check_contract(&store);
236    }
237
238    #[test]
239    fn trait_contract_in_memory_store() {
240        let store = InMemoryStore::new(8);
241        check_contract(&store);
242    }
243}
memory_mcp/index/mod.rs

memory_mcp/index/
mod.rs
