rskit_vectorstore/
registry.rs1use std::collections::BTreeMap;
4use std::sync::Arc;
5
6use rskit_errors::{AppError, AppResult, ErrorCode};
7
8use crate::{InMemoryVectorStore, VectorStore, VectorStoreConfig};
9
10pub trait VectorFactory: Send + Sync {
12 fn create(&self, config: &VectorStoreConfig) -> AppResult<Arc<dyn VectorStore>>;
14}
15
16#[derive(Default)]
18pub struct VectorStoreRegistry {
19 factories: BTreeMap<String, Arc<dyn VectorFactory>>,
20}
21
22impl VectorStoreRegistry {
23 #[must_use]
25 pub fn new() -> Self {
26 Self::default()
27 }
28
29 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 #[must_use]
54 pub fn contains(&self, name: &str) -> bool {
55 self.factories.contains_key(name)
56 }
57
58 #[must_use]
60 pub fn len(&self) -> usize {
61 self.factories.len()
62 }
63
64 #[must_use]
66 pub fn is_empty(&self) -> bool {
67 self.factories.is_empty()
68 }
69
70 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
102pub 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}