umi_memory/umi/
builder.rs1use super::Memory;
9use crate::embedding::EmbeddingProvider;
10use crate::llm::LLMProvider;
11use crate::storage::{StorageBackend, VectorBackend};
12
13pub struct MemoryBuilder<L, E, V, S> {
39 llm: Option<L>,
40 embedder: Option<E>,
41 vector: Option<V>,
42 storage: Option<S>,
43}
44
45impl<L, E, V, S> MemoryBuilder<L, E, V, S>
46where
47 L: LLMProvider + Clone,
48 E: EmbeddingProvider + Clone,
49 V: VectorBackend + Clone,
50 S: StorageBackend + Clone,
51{
52 #[must_use]
54 pub fn new() -> Self {
55 Self {
56 llm: None,
57 embedder: None,
58 vector: None,
59 storage: None,
60 }
61 }
62
63 #[must_use]
68 pub fn with_llm(mut self, llm: L) -> Self {
69 self.llm = Some(llm);
70 self
71 }
72
73 #[must_use]
78 pub fn with_embedder(mut self, embedder: E) -> Self {
79 self.embedder = Some(embedder);
80 self
81 }
82
83 #[must_use]
88 pub fn with_vector(mut self, vector: V) -> Self {
89 self.vector = Some(vector);
90 self
91 }
92
93 #[must_use]
98 pub fn with_storage(mut self, storage: S) -> Self {
99 self.storage = Some(storage);
100 self
101 }
102
103 #[must_use]
111 pub fn build(self) -> Memory<L, E, S, V> {
112 let llm = self.llm.expect("LLM provider is required");
113 let embedder = self.embedder.expect("Embedder is required");
114 let vector = self.vector.expect("Vector backend is required");
115 let storage = self.storage.expect("Storage backend is required");
116
117 Memory::new(llm, embedder, vector, storage)
118 }
119}
120
121impl<L, E, V, S> Default for MemoryBuilder<L, E, V, S>
122where
123 L: LLMProvider + Clone,
124 E: EmbeddingProvider + Clone,
125 V: VectorBackend + Clone,
126 S: StorageBackend + Clone,
127{
128 fn default() -> Self {
129 Self::new()
130 }
131}
132
133#[cfg(test)]
138mod tests {
139 use super::*;
140 use crate::dst::SimConfig;
141 use crate::embedding::SimEmbeddingProvider;
142 use crate::llm::SimLLMProvider;
143 use crate::storage::{SimStorageBackend, SimVectorBackend};
144
145 #[test]
146 fn test_builder_construction() {
147 let builder = MemoryBuilder::<
148 SimLLMProvider,
149 SimEmbeddingProvider,
150 SimVectorBackend,
151 SimStorageBackend,
152 >::new();
153
154 assert!(builder.llm.is_none());
156 assert!(builder.embedder.is_none());
157 assert!(builder.vector.is_none());
158 assert!(builder.storage.is_none());
159 }
160
161 #[test]
162 fn test_builder_with_methods() {
163 let llm = SimLLMProvider::with_seed(42);
164 let embedder = SimEmbeddingProvider::with_seed(42);
165 let vector = SimVectorBackend::new(42);
166 let storage = SimStorageBackend::new(SimConfig::with_seed(42));
167
168 let builder = MemoryBuilder::new()
169 .with_llm(llm)
170 .with_embedder(embedder)
171 .with_vector(vector)
172 .with_storage(storage);
173
174 assert!(builder.llm.is_some());
176 assert!(builder.embedder.is_some());
177 assert!(builder.vector.is_some());
178 assert!(builder.storage.is_some());
179 }
180
181 #[test]
182 #[should_panic(expected = "LLM provider is required")]
183 fn test_builder_missing_llm() {
184 use crate::umi::Memory;
185 let _memory: Memory<SimLLMProvider, SimEmbeddingProvider, SimStorageBackend, SimVectorBackend> =
186 MemoryBuilder::new()
187 .with_embedder(SimEmbeddingProvider::with_seed(42))
188 .with_vector(SimVectorBackend::new(42))
189 .with_storage(SimStorageBackend::new(SimConfig::with_seed(42)))
190 .build();
191 }
192
193 #[test]
194 #[should_panic(expected = "Embedder is required")]
195 fn test_builder_missing_embedder() {
196 use crate::umi::Memory;
197 let _memory: Memory<SimLLMProvider, SimEmbeddingProvider, SimStorageBackend, SimVectorBackend> =
198 MemoryBuilder::new()
199 .with_llm(SimLLMProvider::with_seed(42))
200 .with_vector(SimVectorBackend::new(42))
201 .with_storage(SimStorageBackend::new(SimConfig::with_seed(42)))
202 .build();
203 }
204}