Skip to main content

rskit_vectorstore/
registry.rs

1//! Explicit vector store backend registry.
2
3use std::collections::BTreeMap;
4use std::sync::Arc;
5
6use rskit_errors::{AppError, AppResult, ErrorCode};
7
8use crate::{InMemoryVectorStore, VectorStore, VectorStoreConfig};
9
10/// Factory for a named vector store backend.
11pub trait VectorFactory: Send + Sync {
12    /// Create a vector store backend instance.
13    fn create(&self, config: &VectorStoreConfig) -> AppResult<Arc<dyn VectorStore>>;
14}
15
16/// Explicit vector store backend registry.
17#[derive(Default)]
18pub struct VectorStoreRegistry {
19    factories: BTreeMap<String, Arc<dyn VectorFactory>>,
20}
21
22impl VectorStoreRegistry {
23    /// Create an empty registry. No backends are registered implicitly.
24    #[must_use]
25    pub fn new() -> Self {
26        Self::default()
27    }
28
29    /// Register a backend factory under `name`.
30    pub fn register(
31        &mut self,
32        name: impl Into<String>,
33        factory: Arc<dyn VectorFactory>,
34    ) -> AppResult<()> {
35        let name = name.into().trim().to_owned();
36        if name.is_empty() {
37            return Err(AppError::new(
38                ErrorCode::InvalidInput,
39                "vectorstore backend name is required",
40            ));
41        }
42        if self.factories.contains_key(&name) {
43            return Err(AppError::new(
44                ErrorCode::AlreadyExists,
45                format!("vectorstore backend '{name}' is already registered"),
46            ));
47        }
48        self.factories.insert(name, factory);
49        Ok(())
50    }
51
52    /// Return true when a backend is registered.
53    #[must_use]
54    pub fn contains(&self, name: &str) -> bool {
55        self.factories.contains_key(name)
56    }
57
58    /// Number of registered backends.
59    #[must_use]
60    pub fn len(&self) -> usize {
61        self.factories.len()
62    }
63
64    /// Return true when no backend factories are registered.
65    #[must_use]
66    pub fn is_empty(&self) -> bool {
67        self.factories.is_empty()
68    }
69
70    /// Build the backend selected by [`VectorStoreConfig::backend`].
71    pub fn build(&self, config: &VectorStoreConfig) -> AppResult<Arc<dyn VectorStore>> {
72        let backend = config.backend.trim();
73        if backend.is_empty() {
74            return Err(AppError::new(
75                ErrorCode::InvalidInput,
76                "vectorstore backend name is required",
77            ));
78        }
79        self.factories
80            .get(backend)
81            .ok_or_else(|| {
82                AppError::new(
83                    ErrorCode::NotFound,
84                    format!("vectorstore backend '{backend}' is not registered"),
85                )
86            })?
87            .create(config)
88    }
89}
90
91struct MemoryFactory;
92
93impl VectorFactory for MemoryFactory {
94    fn create(&self, config: &VectorStoreConfig) -> AppResult<Arc<dyn VectorStore>> {
95        Ok(Arc::new(InMemoryVectorStore::with_options(
96            config.memory.metric,
97            config.limits,
98        )))
99    }
100}
101
102/// Explicitly register the in-memory vector store backend.
103pub fn register_memory(registry: &mut VectorStoreRegistry) -> AppResult<()> {
104    registry.register("memory", Arc::new(MemoryFactory))
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn registry_starts_empty() {
113        let registry = VectorStoreRegistry::new();
114
115        assert!(registry.is_empty());
116        assert_eq!(registry.len(), 0);
117        assert!(!registry.contains("memory"));
118    }
119
120    #[test]
121    fn registry_rejects_duplicate_backend() {
122        let mut registry = VectorStoreRegistry::new();
123        register_memory(&mut registry).unwrap();
124
125        let err = register_memory(&mut registry).expect_err("duplicate backend must fail");
126
127        assert_eq!(err.code(), ErrorCode::AlreadyExists);
128    }
129
130    #[test]
131    fn registry_rejects_blank_registration_name() {
132        let mut registry = VectorStoreRegistry::new();
133
134        let err = registry
135            .register(" ", Arc::new(MemoryFactory))
136            .err()
137            .unwrap();
138
139        assert_eq!(err.code(), ErrorCode::InvalidInput);
140    }
141
142    #[test]
143    fn registry_builds_memory_backend_after_explicit_registration() {
144        let mut registry = VectorStoreRegistry::new();
145        register_memory(&mut registry).unwrap();
146
147        let store = registry.build(&VectorStoreConfig::default()).unwrap();
148
149        assert!(Arc::strong_count(&store) >= 1);
150    }
151
152    #[test]
153    fn registry_rejects_unregistered_backend() {
154        let registry = VectorStoreRegistry::new();
155
156        let err = registry.build(&VectorStoreConfig::default()).err().unwrap();
157
158        assert_eq!(err.code(), ErrorCode::NotFound);
159    }
160
161    #[test]
162    fn registry_rejects_blank_backend() {
163        let registry = VectorStoreRegistry::new();
164        let config = VectorStoreConfig {
165            backend: "  ".to_owned(),
166            ..VectorStoreConfig::default()
167        };
168
169        let err = registry.build(&config).err().unwrap();
170
171        assert_eq!(err.code(), ErrorCode::InvalidInput);
172    }
173
174    #[test]
175    fn registry_normalizes_backend_before_lookup() {
176        let mut registry = VectorStoreRegistry::new();
177        register_memory(&mut registry).unwrap();
178        let config = VectorStoreConfig {
179            backend: " memory ".to_owned(),
180            ..VectorStoreConfig::default()
181        };
182
183        let store = registry.build(&config).unwrap();
184
185        assert!(Arc::strong_count(&store) >= 1);
186    }
187}