elph_core/floppy/
builder.rs

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

elph_core/floppy/
builder.rs
