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    /// potion-code-16M — static (model2vec) code embeddings (256d, ~64MB),
28    /// run in pure Rust: no ONNX Runtime, microseconds per chunk.
29    #[serde(rename = "potion-code-16m")]
30    PotionCode16M,
31    /// Any HuggingFace repo with an ONNX export + tokenizer.json
32    /// (`hf:org/repo[@revision]`, GL #397).
33    Custom(CustomModelSpec),
34}
35
36/// How a model's weights are stored and run.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum ModelFormat {
39    /// An ONNX graph run by ONNX Runtime (needs the `embeddings` feature).
40    Onnx,
41    /// A model2vec token-embedding table in safetensors, run in pure Rust.
42    StaticSafetensors,
43}
44
45impl ModelFormat {
46    /// File name of the weights inside the model directory.
47    #[must_use]
48    pub fn local_file(self) -> &'static str {
49        match self {
50            Self::Onnx => "model.onnx",
51            Self::StaticSafetensors => "model.safetensors",
52        }
53    }
54}
55
56/// potion-code-16M at a fixed revision with its file hashes (MIT,
57/// <https://huggingface.co/minishlab/potion-code-16M>). A built-in model is
58/// pinned outright rather than trusted on first use.
59const POTION_CODE_REVISION: &str = "1b0ff71095656b23306542bbad34a09109673720";
60const POTION_CODE_SHA256: &[(&str, &str)] = &[
61    (
62        "model.safetensors",
63        "ca6159081a6e96cebe4ad878e5e8437bfccc761e8db16223370149cd2faa6c0b",
64    ),
65    (
66        "tokenizer.json",
67        "8e84217af15e70e8127c855435fc3d8a4cd91d7bbe686f72e75f188118ec78ae",
68    ),
69];
70
71/// A user-supplied HuggingFace embedding model (`hf:org/repo[@revision]`).
72#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
73pub struct CustomModelSpec {
74    /// HuggingFace repo id, e.g. `jinaai/jina-embeddings-v2-base-code`.
75    pub repo: String,
76    /// Optional revision pin (tag/branch/commit). `None` resolves `main` —
77    /// supply-chain-wise a pin is strongly recommended and the resolver warns
78    /// without one.
79    pub revision: Option<String>,
80    /// Embedding dimensions (`[embedding].dimensions`). When unset, the real
81    /// value is detected from a probe inference at load time; this is only the
82    /// declared fallback.
83    pub dimensions: Option<usize>,
84}
85
86impl CustomModelSpec {
87    /// Parse `org/repo[@revision]` (the part after the `hf:` scheme).
88    /// Returns `None` when the repo id is not plausibly a HuggingFace repo.
89    fn parse(s: &str) -> Option<Self> {
90        let (repo, revision) = match s.split_once('@') {
91            Some((r, rev)) => (
92                r.trim(),
93                Some(rev.trim().to_string()).filter(|v| !v.is_empty()),
94            ),
95            None => (s.trim(), None),
96        };
97        // A HF repo id is exactly `owner/name` with no whitespace.
98        let mut parts = repo.split('/');
99        let (owner, name) = (parts.next()?, parts.next()?);
100        if parts.next().is_some()
101            || owner.is_empty()
102            || name.is_empty()
103            || repo.chars().any(char::is_whitespace)
104        {
105            return None;
106        }
107        Some(Self {
108            repo: repo.to_string(),
109            revision,
110            dimensions: None,
111        })
112    }
113
114    /// Filesystem-safe storage slug, unique per repo+revision.
115    fn storage_slug(&self) -> String {
116        let mut slug = String::from("hf-");
117        for c in self.repo.chars() {
118            slug.push(match c {
119                'a'..='z' | '0'..='9' | '-' => c,
120                'A'..='Z' => c.to_ascii_lowercase(),
121                _ => '-',
122            });
123        }
124        if let Some(rev) = &self.revision {
125            slug.push('-');
126            for c in rev.chars().take(16) {
127                slug.push(match c {
128                    'a'..='z' | '0'..='9' | '-' => c,
129                    'A'..='Z' => c.to_ascii_lowercase(),
130                    _ => '-',
131                });
132            }
133        }
134        slug
135    }
136}
137
138impl EmbeddingModel {
139    pub const DEFAULT: Self = Self::AllMiniLmL6V2;
140
141    pub fn config(&self) -> ModelConfig {
142        match self {
143            Self::AllMiniLmL6V2 => ModelConfig {
144                model: self.clone(),
145                name: "all-MiniLM-L6-v2".into(),
146                hf_repo: "sentence-transformers/all-MiniLM-L6-v2".into(),
147                revision: None,
148                onnx_path: "onnx/model.onnx".into(),
149                vocab_file: VocabSource::VocabTxt("vocab.txt".into()),
150                dimensions: 384,
151                max_seq_len: 256,
152                model_min_bytes: 1_000_000,
153                vocab_min_bytes: 100_000,
154                query_prefix: None,
155                document_prefix: None,
156                needs_token_type_ids: true,
157                format: ModelFormat::Onnx,
158                pinned_sha256: &[],
159            },
160            Self::PotionCode16M => ModelConfig {
161                model: self.clone(),
162                name: "potion-code-16M".into(),
163                hf_repo: "minishlab/potion-code-16M".into(),
164                revision: Some(POTION_CODE_REVISION.into()),
165                onnx_path: "model.safetensors".into(),
166                vocab_file: VocabSource::TokenizerJson("tokenizer.json".into()),
167                dimensions: 256,
168                // Static models average token vectors and see no positions, so a
169                // long chunk costs only its token count; cap it like model2vec.
170                max_seq_len: 512,
171                model_min_bytes: 10_000_000,
172                vocab_min_bytes: 100_000,
173                query_prefix: None,
174                document_prefix: None,
175                needs_token_type_ids: false,
176                format: ModelFormat::StaticSafetensors,
177                pinned_sha256: POTION_CODE_SHA256,
178            },
179            Self::NomicEmbedV1_5 => ModelConfig {
180                model: self.clone(),
181                name: "nomic-embed-text-v1.5".into(),
182                hf_repo: "nomic-ai/nomic-embed-text-v1.5".into(),
183                revision: None,
184                onnx_path: "onnx/model.onnx".into(),
185                vocab_file: VocabSource::VocabTxt("vocab.txt".into()),
186                dimensions: 768,
187                max_seq_len: 512,
188                model_min_bytes: 100_000_000,
189                vocab_min_bytes: 100_000,
190                query_prefix: Some("search_query: ".into()),
191                document_prefix: Some("search_document: ".into()),
192                needs_token_type_ids: false,
193                format: ModelFormat::Onnx,
194                pinned_sha256: &[],
195            },
196            Self::Custom(spec) => ModelConfig {
197                model: self.clone(),
198                // The canonical name doubles as the index `model_id`, so a
199                // repo or revision change triggers the one-shot re-index.
200                name: match &spec.revision {
201                    Some(rev) => format!("hf:{}@{rev}", spec.repo),
202                    None => format!("hf:{}", spec.repo),
203                },
204                hf_repo: spec.repo.clone(),
205                revision: spec.revision.clone(),
206                onnx_path: "onnx/model.onnx".into(),
207                // Custom repos must ship a HuggingFace tokenizer.json — the
208                // universal format (WordPiece/BPE/Unigram all serialize to it).
209                vocab_file: VocabSource::TokenizerJson("tokenizer.json".into()),
210                // Declared fallback; the probe inference at load time detects
211                // the real width (`detect_dimensions`) and wins.
212                dimensions: spec.dimensions.unwrap_or(768),
213                max_seq_len: 512,
214                model_min_bytes: 1_000_000,
215                vocab_min_bytes: 1_000,
216                query_prefix: None,
217                document_prefix: None,
218                // Probed from the ONNX graph at load time; BERT-style models
219                // with a third input still get token_type_ids wired up.
220                needs_token_type_ids: false,
221                format: ModelFormat::Onnx,
222                pinned_sha256: &[],
223            },
224        }
225    }
226
227    /// Parse model name from string (env var / config file).
228    ///
229    /// Accepts the built-in aliases plus the `hf:org/repo[@revision]` scheme
230    /// for custom HuggingFace models (GL #397).
231    pub fn from_str_name(s: &str) -> Option<Self> {
232        let trimmed = s.trim();
233        if let Some(rest) = trimmed.strip_prefix("hf:") {
234            return CustomModelSpec::parse(rest).map(Self::Custom);
235        }
236        match trimmed.to_lowercase().replace('_', "-").as_str() {
237            "all-minilm-l6-v2" | "minilm" | "default" => Some(Self::AllMiniLmL6V2),
238            "nomic-embed-v1.5" | "nomic-embed-text-v1.5" | "nomic" | "nomic-embed" => {
239                Some(Self::NomicEmbedV1_5)
240            }
241            "potion-code-16m" | "potion-code" | "potion" | "static" => Some(Self::PotionCode16M),
242            _ => None,
243        }
244    }
245
246    /// All built-in model variants (custom models are user-defined).
247    pub const ALL: &'static [Self] = &[
248        Self::AllMiniLmL6V2,
249        Self::NomicEmbedV1_5,
250        Self::PotionCode16M,
251    ];
252
253    /// Unique subdirectory name for model storage isolation.
254    pub fn storage_dir_name(&self) -> String {
255        match self {
256            Self::AllMiniLmL6V2 => "all-minilm-l6-v2".to_string(),
257            Self::NomicEmbedV1_5 => "nomic-embed-v1.5".to_string(),
258            Self::PotionCode16M => "potion-code-16m".to_string(),
259            Self::Custom(spec) => spec.storage_slug(),
260        }
261    }
262}
263
264impl fmt::Display for EmbeddingModel {
265    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
266        f.write_str(&self.config().name)
267    }
268}
269
270/// Vocabulary/tokenizer source for a model.
271#[derive(Debug, Clone)]
272pub enum VocabSource {
273    /// Standard BERT vocab.txt (one token per line, WordPiece).
274    VocabTxt(String),
275    /// HuggingFace tokenizer.json (BPE/Unigram/WordPiece via JSON config).
276    TokenizerJson(String),
277}
278
279impl VocabSource {
280    pub fn filename(&self) -> &str {
281        match self {
282            Self::VocabTxt(f) | Self::TokenizerJson(f) => f,
283        }
284    }
285
286    pub fn is_wordpiece(&self) -> bool {
287        matches!(self, Self::VocabTxt(_))
288    }
289}
290
291/// Complete configuration for a single embedding model.
292#[derive(Debug, Clone)]
293pub struct ModelConfig {
294    pub model: EmbeddingModel,
295    pub name: String,
296    pub hf_repo: String,
297    /// Optional revision pin for custom models (`None` = `main`).
298    pub revision: Option<String>,
299    pub onnx_path: String,
300    pub vocab_file: VocabSource,
301    pub dimensions: usize,
302    pub max_seq_len: usize,
303    pub model_min_bytes: u64,
304    pub vocab_min_bytes: u64,
305    /// Optional prefix prepended to queries before embedding.
306    pub query_prefix: Option<String>,
307    /// Optional prefix prepended to documents/code before embedding.
308    pub document_prefix: Option<String>,
309    /// Whether the model expects token_type_ids input (BERT-style).
310    /// Some models (e.g. nomic-embed) only use input_ids + attention_mask.
311    pub needs_token_type_ids: bool,
312    /// How the weights are stored and run.
313    pub format: ModelFormat,
314    /// `(local file, sha256-hex)` every download must match; empty means
315    /// trust-on-first-use via `model.lock.json`.
316    pub pinned_sha256: &'static [(&'static str, &'static str)],
317}
318
319impl ModelConfig {
320    fn resolve_base(&self) -> String {
321        format!(
322            "https://huggingface.co/{}/resolve/{}",
323            self.hf_repo,
324            self.revision.as_deref().unwrap_or("main")
325        )
326    }
327
328    /// File name of the weights inside the model directory.
329    #[must_use]
330    pub fn local_model_file(&self) -> &'static str {
331        self.format.local_file()
332    }
333
334    /// Full HuggingFace download URL for the model weights.
335    pub fn model_url(&self) -> String {
336        format!("{}/{}", self.resolve_base(), self.onnx_path)
337    }
338
339    /// Full HuggingFace download URL for the vocabulary/tokenizer file.
340    pub fn vocab_url(&self) -> String {
341        format!("{}/{}", self.resolve_base(), self.vocab_file.filename())
342    }
343}
344
345/// Resolve which embedding model to use.
346///
347/// Priority: `LEAN_CTX_EMBEDDING_MODEL` env var > `[embedding].model` in `config.toml` >
348/// the default model. An unrecognized name is skipped (with a warning) so a typo in one
349/// source never silently swaps the model — which would otherwise force a full re-index.
350pub fn resolve_model() -> EmbeddingModel {
351    resolve_model_with_source().0
352}
353
354/// Which setting selected the embedding model (#2049: `embeddings status`
355/// says so, because the env var and `config.toml` are easy to confuse).
356#[derive(Debug, Clone, PartialEq, Eq)]
357pub enum ModelSource {
358    /// `LEAN_CTX_EMBEDDING_MODEL`, with its raw value.
359    Env(String),
360    /// `[embedding].model` in `config.toml`, with its raw value.
361    Config(String),
362    /// Nothing (valid) set.
363    Default,
364}
365
366/// [`resolve_model`] plus the setting it came from.
367pub fn resolve_model_with_source() -> (EmbeddingModel, ModelSource) {
368    let env_val = std::env::var("LEAN_CTX_EMBEDDING_MODEL").ok();
369    let embedding_cfg = crate::core::config::Config::load().embedding;
370    resolve_model_with_source_from(
371        env_val.as_deref(),
372        embedding_cfg.model.as_deref(),
373        embedding_cfg.dimensions,
374        default_model(),
375    )
376}
377
378/// The model used when none is configured: MiniLM where ONNX Runtime is
379/// installed, the static `potion-code-16m` everywhere else.
380///
381/// Measured 2026-10-13 on 8,086 documented functions of this repository
382/// (docstring → code): MiniLM MRR 0.348 / recall@10 0.570, potion-code MRR
383/// 0.247 / recall@10 0.431 — but potion-code indexed them in 0.07 s instead
384/// of 557 s (CPU, `eval.rs`), answers a query in 9 µs instead of 26 ms, and
385/// needs no `libonnxruntime` — without which an installation has no semantic
386/// search at all. So the better model wins where it can run, and the one that
387/// runs everywhere covers the rest. Resolved once per process; it only
388/// searches the file system, it never loads the library.
389#[must_use]
390pub fn default_model() -> EmbeddingModel {
391    #[cfg(feature = "embeddings")]
392    {
393        static ORT_PRESENT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
394        let ort_present = *ORT_PRESENT
395            .get_or_init(|| crate::core::ort_environment::resolved_ort_dylib_path().is_ok());
396        if !ort_present {
397            return EmbeddingModel::PotionCode16M;
398        }
399    }
400    EmbeddingModel::DEFAULT
401}
402
403/// Pure model resolution used by [`resolve_model`]; kept separate so the env-var/config
404/// precedence is unit-testable without touching the process environment or the on-disk
405/// `config.toml`.
406#[cfg(test)]
407fn resolve_model_from(
408    env_val: Option<&str>,
409    config_val: Option<&str>,
410    config_dims: Option<usize>,
411) -> EmbeddingModel {
412    resolve_model_with_source_from(env_val, config_val, config_dims, EmbeddingModel::DEFAULT).0
413}
414
415fn resolve_model_with_source_from(
416    env_val: Option<&str>,
417    config_val: Option<&str>,
418    config_dims: Option<usize>,
419    default: EmbeddingModel,
420) -> (EmbeddingModel, ModelSource) {
421    for (source, raw) in [
422        ("LEAN_CTX_EMBEDDING_MODEL", env_val),
423        ("[embedding].model", config_val),
424    ] {
425        let Some(name) = raw.map(str::trim).filter(|s| !s.is_empty()) else {
426            continue;
427        };
428        let origin = if source == "LEAN_CTX_EMBEDDING_MODEL" {
429            ModelSource::Env(name.to_string())
430        } else {
431            ModelSource::Config(name.to_string())
432        };
433        match EmbeddingModel::from_str_name(name) {
434            Some(EmbeddingModel::Custom(mut spec)) => {
435                spec.dimensions = config_dims;
436                if spec.revision.is_none() {
437                    tracing::warn!(
438                        "Custom embedding model {:?} has no revision pin — supply-chain best \
439                         practice is `hf:{}@<commit-or-tag>` so upstream pushes can never \
440                         silently change your index",
441                        spec.repo,
442                        spec.repo
443                    );
444                }
445                return (EmbeddingModel::Custom(spec), origin);
446            }
447            Some(model) => return (model, origin),
448            None => {
449                tracing::warn!(
450                    "Unknown embedding model {name:?} from {source}; using {default} instead \
451                     (built-ins: minilm, nomic, potion-code — or hf:org/repo[@rev])"
452                );
453            }
454        }
455    }
456    (default, ModelSource::Default)
457}
458
459#[cfg(test)]
460mod tests {
461    use super::*;
462
463    #[test]
464    fn default_model_is_minilm() {
465        assert_eq!(EmbeddingModel::DEFAULT, EmbeddingModel::AllMiniLmL6V2);
466    }
467
468    #[test]
469    fn from_str_name_variants() {
470        assert_eq!(
471            EmbeddingModel::from_str_name("minilm"),
472            Some(EmbeddingModel::AllMiniLmL6V2)
473        );
474        assert_eq!(
475            EmbeddingModel::from_str_name("nomic-embed-v1.5"),
476            Some(EmbeddingModel::NomicEmbedV1_5)
477        );
478        assert_eq!(
479            EmbeddingModel::from_str_name("nomic"),
480            Some(EmbeddingModel::NomicEmbedV1_5)
481        );
482        assert_eq!(
483            EmbeddingModel::from_str_name("default"),
484            Some(EmbeddingModel::AllMiniLmL6V2)
485        );
486        assert_eq!(EmbeddingModel::from_str_name("unknown"), None);
487        // The removed jina built-in must no longer resolve as an alias; it is
488        // reachable only via the explicit `hf:` custom scheme.
489        assert_eq!(EmbeddingModel::from_str_name("jina-code-v2"), None);
490        assert_eq!(EmbeddingModel::from_str_name("jina"), None);
491    }
492
493    #[test]
494    fn custom_hf_scheme_parses_repo_and_revision() {
495        let m = EmbeddingModel::from_str_name("hf:jinaai/jina-embeddings-v2-base-code").unwrap();
496        let EmbeddingModel::Custom(spec) = &m else {
497            panic!("expected custom")
498        };
499        assert_eq!(spec.repo, "jinaai/jina-embeddings-v2-base-code");
500        assert_eq!(spec.revision, None);
501
502        let m = EmbeddingModel::from_str_name("hf:org/model@abc123").unwrap();
503        let EmbeddingModel::Custom(spec) = &m else {
504            panic!("expected custom")
505        };
506        assert_eq!(spec.repo, "org/model");
507        assert_eq!(spec.revision.as_deref(), Some("abc123"));
508    }
509
510    #[test]
511    fn custom_hf_scheme_rejects_invalid_repos() {
512        for bad in [
513            "hf:",
514            "hf:no-slash",
515            "hf:too/many/slashes",
516            "hf:with space/repo",
517            "hf:/leading",
518            "hf:trailing/",
519            "hf:org/model@",
520        ] {
521            let parsed = EmbeddingModel::from_str_name(bad);
522            if bad == "hf:org/model@" {
523                // Empty revision degrades to an unpinned spec, not a reject.
524                let Some(EmbeddingModel::Custom(spec)) = parsed else {
525                    panic!("expected custom for {bad}")
526                };
527                assert_eq!(spec.revision, None);
528            } else {
529                assert_eq!(parsed, None, "{bad} should be rejected");
530            }
531        }
532    }
533
534    #[test]
535    fn custom_config_urls_and_storage() {
536        let m = EmbeddingModel::from_str_name("hf:Org/My_Model@v1.2").unwrap();
537        let cfg = m.config();
538        assert_eq!(
539            cfg.model_url(),
540            "https://huggingface.co/Org/My_Model/resolve/v1.2/onnx/model.onnx"
541        );
542        assert_eq!(
543            cfg.vocab_url(),
544            "https://huggingface.co/Org/My_Model/resolve/v1.2/tokenizer.json"
545        );
546        assert!(!cfg.vocab_file.is_wordpiece());
547        assert_eq!(cfg.name, "hf:Org/My_Model@v1.2");
548        assert_eq!(m.storage_dir_name(), "hf-org-my-model-v1-2");
549    }
550
551    #[test]
552    fn custom_storage_slugs_differ_per_revision() {
553        let a = EmbeddingModel::from_str_name("hf:org/model@aaa").unwrap();
554        let b = EmbeddingModel::from_str_name("hf:org/model@bbb").unwrap();
555        let c = EmbeddingModel::from_str_name("hf:org/model").unwrap();
556        let slugs = [
557            a.storage_dir_name(),
558            b.storage_dir_name(),
559            c.storage_dir_name(),
560        ];
561        let unique: std::collections::HashSet<_> = slugs.iter().collect();
562        assert_eq!(unique.len(), 3);
563    }
564
565    #[test]
566    fn all_models_have_valid_configs() {
567        for model in EmbeddingModel::ALL {
568            let cfg = model.config();
569            assert!(!cfg.name.is_empty(), "`cfg.name` is empty");
570            assert!(!cfg.hf_repo.is_empty(), "`cfg.hf_repo` is empty");
571            assert!(cfg.dimensions > 0);
572            assert!(cfg.max_seq_len > 0);
573            assert!(cfg.model_min_bytes > 0);
574            assert!(cfg.vocab_min_bytes > 0);
575        }
576    }
577
578    #[test]
579    fn model_urls_are_valid() {
580        for model in EmbeddingModel::ALL {
581            let cfg = model.config();
582            let model_url = cfg.model_url();
583            let vocab_url = cfg.vocab_url();
584            assert!(model_url.starts_with("https://huggingface.co/"));
585            assert!(vocab_url.starts_with("https://huggingface.co/"));
586            // Pinned built-ins resolve their exact revision, the rest `main`.
587            let rev = cfg.revision.as_deref().unwrap_or("main");
588            assert!(
589                model_url.contains(&format!("resolve/{rev}/")),
590                "{model_url}"
591            );
592        }
593    }
594
595    #[test]
596    fn storage_dir_names_are_unique() {
597        let names: Vec<_> = EmbeddingModel::ALL
598            .iter()
599            .map(EmbeddingModel::storage_dir_name)
600            .collect();
601        let unique: std::collections::HashSet<_> = names.iter().collect();
602        assert_eq!(names.len(), unique.len());
603    }
604
605    #[test]
606    fn display_uses_model_name() {
607        assert_eq!(
608            format!("{}", EmbeddingModel::AllMiniLmL6V2),
609            "all-MiniLM-L6-v2"
610        );
611        assert_eq!(
612            format!("{}", EmbeddingModel::NomicEmbedV1_5),
613            "nomic-embed-text-v1.5"
614        );
615    }
616
617    #[test]
618    fn resolve_defaults_when_nothing_set() {
619        assert_eq!(
620            resolve_model_from(None, None, None),
621            EmbeddingModel::DEFAULT
622        );
623        assert_eq!(
624            resolve_model_from(Some(""), Some("   "), None),
625            EmbeddingModel::DEFAULT
626        );
627    }
628
629    #[test]
630    fn config_selects_model_when_env_unset() {
631        assert_eq!(
632            resolve_model_from(None, Some("nomic"), None),
633            EmbeddingModel::NomicEmbedV1_5
634        );
635        assert_eq!(
636            resolve_model_from(None, Some("minilm"), None),
637            EmbeddingModel::AllMiniLmL6V2
638        );
639    }
640
641    #[test]
642    fn resolution_reports_the_setting_that_won() {
643        assert_eq!(
644            resolve_model_with_source_from(
645                Some(" nomic "),
646                Some("minilm"),
647                None,
648                EmbeddingModel::DEFAULT
649            )
650            .1,
651            ModelSource::Env("nomic".into())
652        );
653        assert_eq!(
654            resolve_model_with_source_from(
655                Some("bogus"),
656                Some("nomic"),
657                None,
658                EmbeddingModel::DEFAULT
659            )
660            .1,
661            ModelSource::Config("nomic".into())
662        );
663        assert_eq!(
664            resolve_model_with_source_from(None, Some("nope"), None, EmbeddingModel::DEFAULT).1,
665            ModelSource::Default
666        );
667    }
668
669    #[test]
670    fn env_var_overrides_config() {
671        assert_eq!(
672            resolve_model_from(Some("minilm"), Some("nomic"), None),
673            EmbeddingModel::AllMiniLmL6V2
674        );
675    }
676
677    #[test]
678    fn unknown_name_falls_through_then_defaults() {
679        // Bad env value → valid config value wins.
680        assert_eq!(
681            resolve_model_from(Some("bogus"), Some("nomic"), None),
682            EmbeddingModel::NomicEmbedV1_5
683        );
684        // Bad everywhere → default (never silently breaks the index).
685        assert_eq!(
686            resolve_model_from(Some("bogus"), Some("nope"), None),
687            EmbeddingModel::DEFAULT
688        );
689        // Empty/whitespace in the higher-priority source is skipped, not treated as a match.
690        assert_eq!(
691            resolve_model_from(Some("   "), Some("nomic"), None),
692            EmbeddingModel::NomicEmbedV1_5
693        );
694    }
695
696    #[test]
697    fn resolve_custom_picks_up_config_dimensions() {
698        let m = resolve_model_from(None, Some("hf:org/model@pin"), Some(1024));
699        let EmbeddingModel::Custom(spec) = m else {
700            panic!("expected custom")
701        };
702        assert_eq!(spec.dimensions, Some(1024));
703        assert_eq!(spec.revision.as_deref(), Some("pin"));
704    }
705
706    #[test]
707    fn nomic_has_prefixes() {
708        let cfg = EmbeddingModel::NomicEmbedV1_5.config();
709        assert!(cfg.query_prefix.is_some());
710        assert!(cfg.document_prefix.is_some());
711        assert!(!cfg.needs_token_type_ids);
712    }
713
714    #[test]
715    fn minilm_is_wordpiece() {
716        let cfg = EmbeddingModel::AllMiniLmL6V2.config();
717        assert!(cfg.vocab_file.is_wordpiece());
718    }
719
720    #[test]
721    fn builtin_models_have_valid_vocab_sources() {
722        // ONNX built-ins are BERT WordPiece (vocab.txt) models; the static
723        // model2vec built-in ships its WordPiece vocabulary as tokenizer.json.
724        for model in EmbeddingModel::ALL {
725            let cfg = model.config();
726            match cfg.format {
727                ModelFormat::Onnx => assert!(
728                    cfg.vocab_file.is_wordpiece(),
729                    "{model} should use WordPiece vocab.txt"
730                ),
731                ModelFormat::StaticSafetensors => assert_eq!(
732                    cfg.vocab_file.filename(),
733                    "tokenizer.json",
734                    "{model} ships tokenizer.json"
735                ),
736            }
737        }
738    }
739
740    #[test]
741    fn custom_models_use_tokenizer_json() {
742        let m = EmbeddingModel::from_str_name("hf:org/model").unwrap();
743        assert!(!m.config().vocab_file.is_wordpiece());
744        assert_eq!(m.config().vocab_file.filename(), "tokenizer.json");
745    }
746}
lean_ctx/core/embeddings/model_registry.rs

lean_ctx/core/embeddings/
model_registry.rs
