lean_ctx/core/embeddings/
model_registry.rs

1//! Embedding model registry — model configs, selection, and metadata.
2//!
3//! Supports multiple ONNX embedding models with different dimensions,
4//! tokenizers, and download sources. Models are selected via the
5//! `LEAN_CTX_EMBEDDING_MODEL` env var or the `[embedding].model` key in `config.toml`
6//! (env var wins) — see [`resolve_model`].
7//!
8//! Besides the built-ins, any compatible HuggingFace repo can be loaded with
9//! `model = "hf:org/repo[@revision]"` (GL #397, upstream #328): the repo must
10//! ship an ONNX export (`onnx/model.onnx`) and a `tokenizer.json`. This custom
11//! path probes the ONNX graph for its real input/output signature, so it suits
12//! code-specialized models (e.g. `hf:jinaai/jina-embeddings-v2-base-code`) that
13//! need no hand-maintained config. See `docs/guides/custom-embeddings.md`.
14
15use std::fmt;
16
17/// Supported embedding models.
18#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
19#[serde(rename_all = "kebab-case")]
20pub enum EmbeddingModel {
21    /// all-MiniLM-L6-v2 — generic sentence embeddings (384d, ~91MB).
22    /// Default model for backward compatibility.
23    AllMiniLmL6V2,
24    /// nomic-embed-text-v1.5 — top MTEB general-purpose (768d, ~547MB).
25    /// Matryoshka representation learning, supports dimension truncation.
26    NomicEmbedV1_5,
27    /// Any HuggingFace repo with an ONNX export + tokenizer.json
28    /// (`hf:org/repo[@revision]`, GL #397).
29    Custom(CustomModelSpec),
30}
31
32/// A user-supplied HuggingFace embedding model (`hf:org/repo[@revision]`).
33#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
34pub struct CustomModelSpec {
35    /// HuggingFace repo id, e.g. `jinaai/jina-embeddings-v2-base-code`.
36    pub repo: String,
37    /// Optional revision pin (tag/branch/commit). `None` resolves `main` —
38    /// supply-chain-wise a pin is strongly recommended and the resolver warns
39    /// without one.
40    pub revision: Option<String>,
41    /// Embedding dimensions (`[embedding].dimensions`). When unset, the real
42    /// value is detected from a probe inference at load time; this is only the
43    /// declared fallback.
44    pub dimensions: Option<usize>,
45}
46
47impl CustomModelSpec {
48    /// Parse `org/repo[@revision]` (the part after the `hf:` scheme).
49    /// Returns `None` when the repo id is not plausibly a HuggingFace repo.
50    fn parse(s: &str) -> Option<Self> {
51        let (repo, revision) = match s.split_once('@') {
52            Some((r, rev)) => (
53                r.trim(),
54                Some(rev.trim().to_string()).filter(|v| !v.is_empty()),
55            ),
56            None => (s.trim(), None),
57        };
58        // A HF repo id is exactly `owner/name` with no whitespace.
59        let mut parts = repo.split('/');
60        let (owner, name) = (parts.next()?, parts.next()?);
61        if parts.next().is_some()
62            || owner.is_empty()
63            || name.is_empty()
64            || repo.chars().any(char::is_whitespace)
65        {
66            return None;
67        }
68        Some(Self {
69            repo: repo.to_string(),
70            revision,
71            dimensions: None,
72        })
73    }
74
75    /// Filesystem-safe storage slug, unique per repo+revision.
76    fn storage_slug(&self) -> String {
77        let mut slug = String::from("hf-");
78        for c in self.repo.chars() {
79            slug.push(match c {
80                'a'..='z' | '0'..='9' | '-' => c,
81                'A'..='Z' => c.to_ascii_lowercase(),
82                _ => '-',
83            });
84        }
85        if let Some(rev) = &self.revision {
86            slug.push('-');
87            for c in rev.chars().take(16) {
88                slug.push(match c {
89                    'a'..='z' | '0'..='9' | '-' => c,
90                    'A'..='Z' => c.to_ascii_lowercase(),
91                    _ => '-',
92                });
93            }
94        }
95        slug
96    }
97}
98
99impl EmbeddingModel {
100    pub const DEFAULT: Self = Self::AllMiniLmL6V2;
101
102    pub fn config(&self) -> ModelConfig {
103        match self {
104            Self::AllMiniLmL6V2 => ModelConfig {
105                model: self.clone(),
106                name: "all-MiniLM-L6-v2".into(),
107                hf_repo: "sentence-transformers/all-MiniLM-L6-v2".into(),
108                revision: None,
109                onnx_path: "onnx/model.onnx".into(),
110                vocab_file: VocabSource::VocabTxt("vocab.txt".into()),
111                dimensions: 384,
112                max_seq_len: 256,
113                model_min_bytes: 1_000_000,
114                vocab_min_bytes: 100_000,
115                query_prefix: None,
116                document_prefix: None,
117                needs_token_type_ids: true,
118            },
119            Self::NomicEmbedV1_5 => ModelConfig {
120                model: self.clone(),
121                name: "nomic-embed-text-v1.5".into(),
122                hf_repo: "nomic-ai/nomic-embed-text-v1.5".into(),
123                revision: None,
124                onnx_path: "onnx/model.onnx".into(),
125                vocab_file: VocabSource::VocabTxt("vocab.txt".into()),
126                dimensions: 768,
127                max_seq_len: 512,
128                model_min_bytes: 100_000_000,
129                vocab_min_bytes: 100_000,
130                query_prefix: Some("search_query: ".into()),
131                document_prefix: Some("search_document: ".into()),
132                needs_token_type_ids: false,
133            },
134            Self::Custom(spec) => ModelConfig {
135                model: self.clone(),
136                // The canonical name doubles as the index `model_id`, so a
137                // repo or revision change triggers the one-shot re-index.
138                name: match &spec.revision {
139                    Some(rev) => format!("hf:{}@{rev}", spec.repo),
140                    None => format!("hf:{}", spec.repo),
141                },
142                hf_repo: spec.repo.clone(),
143                revision: spec.revision.clone(),
144                onnx_path: "onnx/model.onnx".into(),
145                // Custom repos must ship a HuggingFace tokenizer.json — the
146                // universal format (WordPiece/BPE/Unigram all serialize to it).
147                vocab_file: VocabSource::TokenizerJson("tokenizer.json".into()),
148                // Declared fallback; the probe inference at load time detects
149                // the real width (`detect_dimensions`) and wins.
150                dimensions: spec.dimensions.unwrap_or(768),
151                max_seq_len: 512,
152                model_min_bytes: 1_000_000,
153                vocab_min_bytes: 1_000,
154                query_prefix: None,
155                document_prefix: None,
156                // Probed from the ONNX graph at load time; BERT-style models
157                // with a third input still get token_type_ids wired up.
158                needs_token_type_ids: false,
159            },
160        }
161    }
162
163    /// Parse model name from string (env var / config file).
164    ///
165    /// Accepts the built-in aliases plus the `hf:org/repo[@revision]` scheme
166    /// for custom HuggingFace models (GL #397).
167    pub fn from_str_name(s: &str) -> Option<Self> {
168        let trimmed = s.trim();
169        if let Some(rest) = trimmed.strip_prefix("hf:") {
170            return CustomModelSpec::parse(rest).map(Self::Custom);
171        }
172        match trimmed.to_lowercase().replace('_', "-").as_str() {
173            "all-minilm-l6-v2" | "minilm" | "default" => Some(Self::AllMiniLmL6V2),
174            "nomic-embed-v1.5" | "nomic-embed-text-v1.5" | "nomic" | "nomic-embed" => {
175                Some(Self::NomicEmbedV1_5)
176            }
177            _ => None,
178        }
179    }
180
181    /// All built-in model variants (custom models are user-defined).
182    pub const ALL: &'static [Self] = &[Self::AllMiniLmL6V2, Self::NomicEmbedV1_5];
183
184    /// Unique subdirectory name for model storage isolation.
185    pub fn storage_dir_name(&self) -> String {
186        match self {
187            Self::AllMiniLmL6V2 => "all-minilm-l6-v2".to_string(),
188            Self::NomicEmbedV1_5 => "nomic-embed-v1.5".to_string(),
189            Self::Custom(spec) => spec.storage_slug(),
190        }
191    }
192}
193
194impl fmt::Display for EmbeddingModel {
195    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
196        f.write_str(&self.config().name)
197    }
198}
199
200/// Vocabulary/tokenizer source for a model.
201#[derive(Debug, Clone)]
202pub enum VocabSource {
203    /// Standard BERT vocab.txt (one token per line, WordPiece).
204    VocabTxt(String),
205    /// HuggingFace tokenizer.json (BPE/Unigram/WordPiece via JSON config).
206    TokenizerJson(String),
207}
208
209impl VocabSource {
210    pub fn filename(&self) -> &str {
211        match self {
212            Self::VocabTxt(f) | Self::TokenizerJson(f) => f,
213        }
214    }
215
216    pub fn is_wordpiece(&self) -> bool {
217        matches!(self, Self::VocabTxt(_))
218    }
219}
220
221/// Complete configuration for a single embedding model.
222#[derive(Debug, Clone)]
223pub struct ModelConfig {
224    pub model: EmbeddingModel,
225    pub name: String,
226    pub hf_repo: String,
227    /// Optional revision pin for custom models (`None` = `main`).
228    pub revision: Option<String>,
229    pub onnx_path: String,
230    pub vocab_file: VocabSource,
231    pub dimensions: usize,
232    pub max_seq_len: usize,
233    pub model_min_bytes: u64,
234    pub vocab_min_bytes: u64,
235    /// Optional prefix prepended to queries before embedding.
236    pub query_prefix: Option<String>,
237    /// Optional prefix prepended to documents/code before embedding.
238    pub document_prefix: Option<String>,
239    /// Whether the model expects token_type_ids input (BERT-style).
240    /// Some models (e.g. nomic-embed) only use input_ids + attention_mask.
241    pub needs_token_type_ids: bool,
242}
243
244impl ModelConfig {
245    fn resolve_base(&self) -> String {
246        format!(
247            "https://huggingface.co/{}/resolve/{}",
248            self.hf_repo,
249            self.revision.as_deref().unwrap_or("main")
250        )
251    }
252
253    /// Full HuggingFace download URL for the ONNX model file.
254    pub fn model_url(&self) -> String {
255        format!("{}/{}", self.resolve_base(), self.onnx_path)
256    }
257
258    /// Full HuggingFace download URL for the vocabulary/tokenizer file.
259    pub fn vocab_url(&self) -> String {
260        format!("{}/{}", self.resolve_base(), self.vocab_file.filename())
261    }
262}
263
264/// Resolve which embedding model to use.
265///
266/// Priority: `LEAN_CTX_EMBEDDING_MODEL` env var > `[embedding].model` in `config.toml` >
267/// the default model. An unrecognized name is skipped (with a warning) so a typo in one
268/// source never silently swaps the model — which would otherwise force a full re-index.
269pub fn resolve_model() -> EmbeddingModel {
270    let env_val = std::env::var("LEAN_CTX_EMBEDDING_MODEL").ok();
271    let embedding_cfg = crate::core::config::Config::load().embedding;
272    resolve_model_from(
273        env_val.as_deref(),
274        embedding_cfg.model.as_deref(),
275        embedding_cfg.dimensions,
276    )
277}
278
279/// Pure model resolution used by [`resolve_model`]; kept separate so the env-var/config
280/// precedence is unit-testable without touching the process environment or the on-disk
281/// `config.toml`.
282fn resolve_model_from(
283    env_val: Option<&str>,
284    config_val: Option<&str>,
285    config_dims: Option<usize>,
286) -> EmbeddingModel {
287    for (source, raw) in [
288        ("LEAN_CTX_EMBEDDING_MODEL", env_val),
289        ("[embedding].model", config_val),
290    ] {
291        let Some(name) = raw.map(str::trim).filter(|s| !s.is_empty()) else {
292            continue;
293        };
294        match EmbeddingModel::from_str_name(name) {
295            Some(EmbeddingModel::Custom(mut spec)) => {
296                spec.dimensions = config_dims;
297                if spec.revision.is_none() {
298                    tracing::warn!(
299                        "Custom embedding model {:?} has no revision pin — supply-chain best \
300                         practice is `hf:{}@<commit-or-tag>` so upstream pushes can never \
301                         silently change your index",
302                        spec.repo,
303                        spec.repo
304                    );
305                }
306                return EmbeddingModel::Custom(spec);
307            }
308            Some(model) => return model,
309            None => {
310                tracing::warn!(
311                    "Unknown embedding model {name:?} from {source}; using {} instead \
312                     (built-ins: minilm, nomic — or hf:org/repo[@rev])",
313                    EmbeddingModel::DEFAULT
314                );
315            }
316        }
317    }
318    EmbeddingModel::DEFAULT
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324
325    #[test]
326    fn default_model_is_minilm() {
327        assert_eq!(EmbeddingModel::DEFAULT, EmbeddingModel::AllMiniLmL6V2);
328    }
329
330    #[test]
331    fn from_str_name_variants() {
332        assert_eq!(
333            EmbeddingModel::from_str_name("minilm"),
334            Some(EmbeddingModel::AllMiniLmL6V2)
335        );
336        assert_eq!(
337            EmbeddingModel::from_str_name("nomic-embed-v1.5"),
338            Some(EmbeddingModel::NomicEmbedV1_5)
339        );
340        assert_eq!(
341            EmbeddingModel::from_str_name("nomic"),
342            Some(EmbeddingModel::NomicEmbedV1_5)
343        );
344        assert_eq!(
345            EmbeddingModel::from_str_name("default"),
346            Some(EmbeddingModel::AllMiniLmL6V2)
347        );
348        assert_eq!(EmbeddingModel::from_str_name("unknown"), None);
349        // The removed jina built-in must no longer resolve as an alias; it is
350        // reachable only via the explicit `hf:` custom scheme.
351        assert_eq!(EmbeddingModel::from_str_name("jina-code-v2"), None);
352        assert_eq!(EmbeddingModel::from_str_name("jina"), None);
353    }
354
355    #[test]
356    fn custom_hf_scheme_parses_repo_and_revision() {
357        let m = EmbeddingModel::from_str_name("hf:jinaai/jina-embeddings-v2-base-code").unwrap();
358        let EmbeddingModel::Custom(spec) = &m else {
359            panic!("expected custom")
360        };
361        assert_eq!(spec.repo, "jinaai/jina-embeddings-v2-base-code");
362        assert_eq!(spec.revision, None);
363
364        let m = EmbeddingModel::from_str_name("hf:org/model@abc123").unwrap();
365        let EmbeddingModel::Custom(spec) = &m else {
366            panic!("expected custom")
367        };
368        assert_eq!(spec.repo, "org/model");
369        assert_eq!(spec.revision.as_deref(), Some("abc123"));
370    }
371
372    #[test]
373    fn custom_hf_scheme_rejects_invalid_repos() {
374        for bad in [
375            "hf:",
376            "hf:no-slash",
377            "hf:too/many/slashes",
378            "hf:with space/repo",
379            "hf:/leading",
380            "hf:trailing/",
381            "hf:org/model@",
382        ] {
383            let parsed = EmbeddingModel::from_str_name(bad);
384            if bad == "hf:org/model@" {
385                // Empty revision degrades to an unpinned spec, not a reject.
386                let Some(EmbeddingModel::Custom(spec)) = parsed else {
387                    panic!("expected custom for {bad}")
388                };
389                assert_eq!(spec.revision, None);
390            } else {
391                assert_eq!(parsed, None, "{bad} should be rejected");
392            }
393        }
394    }
395
396    #[test]
397    fn custom_config_urls_and_storage() {
398        let m = EmbeddingModel::from_str_name("hf:Org/My_Model@v1.2").unwrap();
399        let cfg = m.config();
400        assert_eq!(
401            cfg.model_url(),
402            "https://huggingface.co/Org/My_Model/resolve/v1.2/onnx/model.onnx"
403        );
404        assert_eq!(
405            cfg.vocab_url(),
406            "https://huggingface.co/Org/My_Model/resolve/v1.2/tokenizer.json"
407        );
408        assert!(!cfg.vocab_file.is_wordpiece());
409        assert_eq!(cfg.name, "hf:Org/My_Model@v1.2");
410        assert_eq!(m.storage_dir_name(), "hf-org-my-model-v1-2");
411    }
412
413    #[test]
414    fn custom_storage_slugs_differ_per_revision() {
415        let a = EmbeddingModel::from_str_name("hf:org/model@aaa").unwrap();
416        let b = EmbeddingModel::from_str_name("hf:org/model@bbb").unwrap();
417        let c = EmbeddingModel::from_str_name("hf:org/model").unwrap();
418        let slugs = [
419            a.storage_dir_name(),
420            b.storage_dir_name(),
421            c.storage_dir_name(),
422        ];
423        let unique: std::collections::HashSet<_> = slugs.iter().collect();
424        assert_eq!(unique.len(), 3);
425    }
426
427    #[test]
428    fn all_models_have_valid_configs() {
429        for model in EmbeddingModel::ALL {
430            let cfg = model.config();
431            assert!(!cfg.name.is_empty());
432            assert!(!cfg.hf_repo.is_empty());
433            assert!(cfg.dimensions > 0);
434            assert!(cfg.max_seq_len > 0);
435            assert!(cfg.model_min_bytes > 0);
436            assert!(cfg.vocab_min_bytes > 0);
437        }
438    }
439
440    #[test]
441    fn model_urls_are_valid() {
442        for model in EmbeddingModel::ALL {
443            let cfg = model.config();
444            let model_url = cfg.model_url();
445            let vocab_url = cfg.vocab_url();
446            assert!(model_url.starts_with("https://huggingface.co/"));
447            assert!(vocab_url.starts_with("https://huggingface.co/"));
448            assert!(model_url.contains("resolve/main"));
449        }
450    }
451
452    #[test]
453    fn storage_dir_names_are_unique() {
454        let names: Vec<_> = EmbeddingModel::ALL
455            .iter()
456            .map(EmbeddingModel::storage_dir_name)
457            .collect();
458        let unique: std::collections::HashSet<_> = names.iter().collect();
459        assert_eq!(names.len(), unique.len());
460    }
461
462    #[test]
463    fn display_uses_model_name() {
464        assert_eq!(
465            format!("{}", EmbeddingModel::AllMiniLmL6V2),
466            "all-MiniLM-L6-v2"
467        );
468        assert_eq!(
469            format!("{}", EmbeddingModel::NomicEmbedV1_5),
470            "nomic-embed-text-v1.5"
471        );
472    }
473
474    #[test]
475    fn resolve_defaults_when_nothing_set() {
476        assert_eq!(
477            resolve_model_from(None, None, None),
478            EmbeddingModel::DEFAULT
479        );
480        assert_eq!(
481            resolve_model_from(Some(""), Some("   "), None),
482            EmbeddingModel::DEFAULT
483        );
484    }
485
486    #[test]
487    fn config_selects_model_when_env_unset() {
488        assert_eq!(
489            resolve_model_from(None, Some("nomic"), None),
490            EmbeddingModel::NomicEmbedV1_5
491        );
492        assert_eq!(
493            resolve_model_from(None, Some("minilm"), None),
494            EmbeddingModel::AllMiniLmL6V2
495        );
496    }
497
498    #[test]
499    fn env_var_overrides_config() {
500        assert_eq!(
501            resolve_model_from(Some("minilm"), Some("nomic"), None),
502            EmbeddingModel::AllMiniLmL6V2
503        );
504    }
505
506    #[test]
507    fn unknown_name_falls_through_then_defaults() {
508        // Bad env value → valid config value wins.
509        assert_eq!(
510            resolve_model_from(Some("bogus"), Some("nomic"), None),
511            EmbeddingModel::NomicEmbedV1_5
512        );
513        // Bad everywhere → default (never silently breaks the index).
514        assert_eq!(
515            resolve_model_from(Some("bogus"), Some("nope"), None),
516            EmbeddingModel::DEFAULT
517        );
518        // Empty/whitespace in the higher-priority source is skipped, not treated as a match.
519        assert_eq!(
520            resolve_model_from(Some("   "), Some("nomic"), None),
521            EmbeddingModel::NomicEmbedV1_5
522        );
523    }
524
525    #[test]
526    fn resolve_custom_picks_up_config_dimensions() {
527        let m = resolve_model_from(None, Some("hf:org/model@pin"), Some(1024));
528        let EmbeddingModel::Custom(spec) = m else {
529            panic!("expected custom")
530        };
531        assert_eq!(spec.dimensions, Some(1024));
532        assert_eq!(spec.revision.as_deref(), Some("pin"));
533    }
534
535    #[test]
536    fn nomic_has_prefixes() {
537        let cfg = EmbeddingModel::NomicEmbedV1_5.config();
538        assert!(cfg.query_prefix.is_some());
539        assert!(cfg.document_prefix.is_some());
540        assert!(!cfg.needs_token_type_ids);
541    }
542
543    #[test]
544    fn minilm_is_wordpiece() {
545        let cfg = EmbeddingModel::AllMiniLmL6V2.config();
546        assert!(cfg.vocab_file.is_wordpiece());
547    }
548
549    #[test]
550    fn builtin_models_have_valid_vocab_sources() {
551        // All current built-ins are WordPiece (vocab.txt) models. Custom HF
552        // repos use tokenizer.json, but those are user-defined, not built-ins.
553        for model in EmbeddingModel::ALL {
554            assert!(
555                model.config().vocab_file.is_wordpiece(),
556                "{model} should use WordPiece vocab.txt"
557            );
558        }
559    }
560
561    #[test]
562    fn custom_models_use_tokenizer_json() {
563        let m = EmbeddingModel::from_str_name("hf:org/model").unwrap();
564        assert!(!m.config().vocab_file.is_wordpiece());
565        assert_eq!(m.config().vocab_file.filename(), "tokenizer.json");
566    }
567}
lean_ctx/core/embeddings/model_registry.rs

lean_ctx/core/embeddings/
model_registry.rs
