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

lean_ctx/core/embeddings/
model_registry.rs
