elph_core/memz/
builder.rs

1use anyhow::{Result, bail};
2
3use super::embed::{DEFAULT_EMBED_MODEL, FastEmbedOptions, create_fastembed};
4#[cfg(feature = "fastembed")]
5use super::embed::{embedding_dims, resolve_embedding_model};
6use super::paths::MemzPaths;
7use super::store::{EmbedFn, MemoryStore, noop_embedder};
8use super::types::{MemzConfig, VectorType};
9use super::util::DEFAULT_EMBEDDING_DIMS;
10
11/// Builder for a [`MemoryStore`] with explicit configuration (no environment variables).
12pub struct MemzBuilder {
13    config: MemzConfig,
14    custom_embed: Option<EmbedFn>,
15    fastembed_opts: Option<FastEmbedOptions>,
16}
17
18impl MemzBuilder {
19    pub fn new(db_path: impl Into<String>, session_id: impl Into<String>) -> Self {
20        Self {
21            config: MemzConfig::new(db_path, session_id),
22            custom_embed: None,
23            fastembed_opts: None,
24        }
25    }
26
27    pub fn from_paths(paths: &MemzPaths, session_id: impl Into<String>) -> Self {
28        Self::new(paths.db_path_string(), session_id)
29    }
30
31    pub fn vector_type(mut self, vector_type: VectorType) -> Self {
32        self.config = self.config.vector_type(vector_type);
33        self
34    }
35
36    pub fn dimensions(mut self, dimensions: u32) -> Self {
37        self.config = self.config.dimensions(dimensions);
38        self
39    }
40
41    pub fn top_k(mut self, top_k: u32) -> Self {
42        self.config = self.config.top_k(top_k);
43        self
44    }
45
46    pub fn learning_rate(mut self, learning_rate: f64) -> Self {
47        self.config = self.config.learning_rate(learning_rate);
48        self
49    }
50
51    pub fn decay_rate(mut self, decay_rate: f64) -> Self {
52        self.config = self.config.decay_rate(decay_rate);
53        self
54    }
55
56    /// Skip memz migrations in [`MemoryStore::init`] when the host already applied them.
57    pub fn apply_migrations(mut self, apply: bool) -> Self {
58        self.config = self.config.apply_migrations(apply);
59        self
60    }
61
62    /// Custom embedder; mutually exclusive with [`Self::fastembed`].
63    pub fn embed(mut self, embed: EmbedFn) -> Self {
64        self.custom_embed = Some(embed);
65        self.fastembed_opts = None;
66        self
67    }
68
69    /// Zero-vector embedder for read-only inspection without a local model.
70    pub fn noop_embed(mut self) -> Self {
71        let dims = self.config.dimensions.unwrap_or(DEFAULT_EMBEDDING_DIMS);
72        self.custom_embed = Some(noop_embedder(dims));
73        self.fastembed_opts = None;
74        self
75    }
76
77    /// Local fastembed backend. Sets [`MemzConfig::dimensions`] from the resolved model.
78    #[cfg(feature = "fastembed")]
79    pub fn fastembed(mut self, options: FastEmbedOptions) -> Result<Self> {
80        let model_name = options.model.as_deref().unwrap_or(DEFAULT_EMBED_MODEL);
81        let dims = resolve_embedding_model(model_name, options.quantized)
82            .map(|m| embedding_dims(&m))
83            .map_err(|e| anyhow::anyhow!("{e}"))?;
84        self.config = self.config.dimensions(dims);
85        self.fastembed_opts = Some(options);
86        self.custom_embed = None;
87        Ok(self)
88    }
89
90    #[cfg(not(feature = "fastembed"))]
91    pub fn fastembed(mut self, _options: FastEmbedOptions) -> Result<Self> {
92        bail!("fastembed backend requires the `fastembed` feature on this crate");
93    }
94
95    pub fn build(self) -> Result<MemoryStore> {
96        let embed = match (self.custom_embed, self.fastembed_opts) {
97            (Some(e), None) => e,
98            (None, Some(opts)) => create_fastembed(opts)?,
99            (None, None) => {
100                bail!("embedder required: call .embed(), .noop_embed(), or .fastembed()");
101            }
102            (Some(_), Some(_)) => bail!("cannot set both a custom embedder and fastembed options"),
103        };
104        Ok(MemoryStore::new(self.config, embed))
105    }
106}
107
108impl MemzPaths {
109    /// Start a [`MemzBuilder`] rooted at this data directory.
110    pub fn builder(&self, session_id: impl Into<String>) -> MemzBuilder {
111        MemzBuilder::from_paths(self, session_id)
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use std::sync::Arc;
119
120    fn mock_embed() -> EmbedFn {
121        Arc::new(|_| Box::pin(async { Ok(vec![1.0, 0.0, 0.0, 0.0]) }))
122    }
123
124    #[test]
125    fn builder_requires_embedder() {
126        let dir = tempfile::tempdir().expect("tempdir");
127        let db = dir.path().join("t.db").to_string_lossy().into_owned();
128        match MemzBuilder::new(db, "s").build() {
129            Err(e) => assert!(e.to_string().contains("embedder required")),
130            Ok(_) => panic!("expected embedder required error"),
131        }
132    }
133
134    #[test]
135    fn builder_with_custom_embed() {
136        let dir = tempfile::tempdir().expect("tempdir");
137        let db = dir.path().join("t.db").to_string_lossy().into_owned();
138        let store = MemzBuilder::new(db, "s")
139            .dimensions(4)
140            .embed(mock_embed())
141            .build()
142            .expect("build");
143        assert_eq!(store.dimensions(), 4);
144    }
145
146    #[test]
147    fn from_paths_sets_db_location() {
148        let paths = MemzPaths::project_local();
149        let store = paths.builder("sess").dimensions(4).noop_embed().build().expect("build");
150        assert_eq!(store.dimensions(), 4);
151    }
152}
elph_core/memz/builder.rs

elph_core/memz/
builder.rs
