use std::fmt;
#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum EmbeddingModel {
AllMiniLmL6V2,
NomicEmbedV1_5,
#[serde(rename = "potion-code-16m")]
PotionCode16M,
Custom(CustomModelSpec),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelFormat {
Onnx,
StaticSafetensors,
}
impl ModelFormat {
#[must_use]
pub fn local_file(self) -> &'static str {
match self {
Self::Onnx => "model.onnx",
Self::StaticSafetensors => "model.safetensors",
}
}
}
const POTION_CODE_REVISION: &str = "1b0ff71095656b23306542bbad34a09109673720";
const POTION_CODE_SHA256: &[(&str, &str)] = &[
(
"model.safetensors",
"ca6159081a6e96cebe4ad878e5e8437bfccc761e8db16223370149cd2faa6c0b",
),
(
"tokenizer.json",
"8e84217af15e70e8127c855435fc3d8a4cd91d7bbe686f72e75f188118ec78ae",
),
];
#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub struct CustomModelSpec {
pub repo: String,
pub revision: Option<String>,
pub dimensions: Option<usize>,
}
impl CustomModelSpec {
fn parse(s: &str) -> Option<Self> {
let (repo, revision) = match s.split_once('@') {
Some((r, rev)) => (
r.trim(),
Some(rev.trim().to_string()).filter(|v| !v.is_empty()),
),
None => (s.trim(), None),
};
let mut parts = repo.split('/');
let (owner, name) = (parts.next()?, parts.next()?);
if parts.next().is_some()
|| owner.is_empty()
|| name.is_empty()
|| repo.chars().any(char::is_whitespace)
{
return None;
}
Some(Self {
repo: repo.to_string(),
revision,
dimensions: None,
})
}
fn storage_slug(&self) -> String {
let mut slug = String::from("hf-");
for c in self.repo.chars() {
slug.push(match c {
'a'..='z' | '0'..='9' | '-' => c,
'A'..='Z' => c.to_ascii_lowercase(),
_ => '-',
});
}
if let Some(rev) = &self.revision {
slug.push('-');
for c in rev.chars().take(16) {
slug.push(match c {
'a'..='z' | '0'..='9' | '-' => c,
'A'..='Z' => c.to_ascii_lowercase(),
_ => '-',
});
}
}
slug
}
}
impl EmbeddingModel {
pub const DEFAULT: Self = Self::AllMiniLmL6V2;
pub fn config(&self) -> ModelConfig {
match self {
Self::AllMiniLmL6V2 => ModelConfig {
model: self.clone(),
name: "all-MiniLM-L6-v2".into(),
hf_repo: "sentence-transformers/all-MiniLM-L6-v2".into(),
revision: None,
onnx_path: "onnx/model.onnx".into(),
vocab_file: VocabSource::VocabTxt("vocab.txt".into()),
dimensions: 384,
max_seq_len: 256,
model_min_bytes: 1_000_000,
vocab_min_bytes: 100_000,
query_prefix: None,
document_prefix: None,
needs_token_type_ids: true,
format: ModelFormat::Onnx,
pinned_sha256: &[],
},
Self::PotionCode16M => ModelConfig {
model: self.clone(),
name: "potion-code-16M".into(),
hf_repo: "minishlab/potion-code-16M".into(),
revision: Some(POTION_CODE_REVISION.into()),
onnx_path: "model.safetensors".into(),
vocab_file: VocabSource::TokenizerJson("tokenizer.json".into()),
dimensions: 256,
max_seq_len: 512,
model_min_bytes: 10_000_000,
vocab_min_bytes: 100_000,
query_prefix: None,
document_prefix: None,
needs_token_type_ids: false,
format: ModelFormat::StaticSafetensors,
pinned_sha256: POTION_CODE_SHA256,
},
Self::NomicEmbedV1_5 => ModelConfig {
model: self.clone(),
name: "nomic-embed-text-v1.5".into(),
hf_repo: "nomic-ai/nomic-embed-text-v1.5".into(),
revision: None,
onnx_path: "onnx/model.onnx".into(),
vocab_file: VocabSource::VocabTxt("vocab.txt".into()),
dimensions: 768,
max_seq_len: 512,
model_min_bytes: 100_000_000,
vocab_min_bytes: 100_000,
query_prefix: Some("search_query: ".into()),
document_prefix: Some("search_document: ".into()),
needs_token_type_ids: false,
format: ModelFormat::Onnx,
pinned_sha256: &[],
},
Self::Custom(spec) => ModelConfig {
model: self.clone(),
name: match &spec.revision {
Some(rev) => format!("hf:{}@{rev}", spec.repo),
None => format!("hf:{}", spec.repo),
},
hf_repo: spec.repo.clone(),
revision: spec.revision.clone(),
onnx_path: "onnx/model.onnx".into(),
vocab_file: VocabSource::TokenizerJson("tokenizer.json".into()),
dimensions: spec.dimensions.unwrap_or(768),
max_seq_len: 512,
model_min_bytes: 1_000_000,
vocab_min_bytes: 1_000,
query_prefix: None,
document_prefix: None,
needs_token_type_ids: false,
format: ModelFormat::Onnx,
pinned_sha256: &[],
},
}
}
pub fn from_str_name(s: &str) -> Option<Self> {
let trimmed = s.trim();
if let Some(rest) = trimmed.strip_prefix("hf:") {
return CustomModelSpec::parse(rest).map(Self::Custom);
}
match trimmed.to_lowercase().replace('_', "-").as_str() {
"all-minilm-l6-v2" | "minilm" | "default" => Some(Self::AllMiniLmL6V2),
"nomic-embed-v1.5" | "nomic-embed-text-v1.5" | "nomic" | "nomic-embed" => {
Some(Self::NomicEmbedV1_5)
}
"potion-code-16m" | "potion-code" | "potion" | "static" => Some(Self::PotionCode16M),
_ => None,
}
}
pub const ALL: &'static [Self] = &[
Self::AllMiniLmL6V2,
Self::NomicEmbedV1_5,
Self::PotionCode16M,
];
pub fn storage_dir_name(&self) -> String {
match self {
Self::AllMiniLmL6V2 => "all-minilm-l6-v2".to_string(),
Self::NomicEmbedV1_5 => "nomic-embed-v1.5".to_string(),
Self::PotionCode16M => "potion-code-16m".to_string(),
Self::Custom(spec) => spec.storage_slug(),
}
}
}
impl fmt::Display for EmbeddingModel {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.config().name)
}
}
#[derive(Debug, Clone)]
pub enum VocabSource {
VocabTxt(String),
TokenizerJson(String),
}
impl VocabSource {
pub fn filename(&self) -> &str {
match self {
Self::VocabTxt(f) | Self::TokenizerJson(f) => f,
}
}
pub fn is_wordpiece(&self) -> bool {
matches!(self, Self::VocabTxt(_))
}
}
#[derive(Debug, Clone)]
pub struct ModelConfig {
pub model: EmbeddingModel,
pub name: String,
pub hf_repo: String,
pub revision: Option<String>,
pub onnx_path: String,
pub vocab_file: VocabSource,
pub dimensions: usize,
pub max_seq_len: usize,
pub model_min_bytes: u64,
pub vocab_min_bytes: u64,
pub query_prefix: Option<String>,
pub document_prefix: Option<String>,
pub needs_token_type_ids: bool,
pub format: ModelFormat,
pub pinned_sha256: &'static [(&'static str, &'static str)],
}
impl ModelConfig {
fn resolve_base(&self) -> String {
format!(
"https://huggingface.co/{}/resolve/{}",
self.hf_repo,
self.revision.as_deref().unwrap_or("main")
)
}
#[must_use]
pub fn local_model_file(&self) -> &'static str {
self.format.local_file()
}
pub fn model_url(&self) -> String {
format!("{}/{}", self.resolve_base(), self.onnx_path)
}
pub fn vocab_url(&self) -> String {
format!("{}/{}", self.resolve_base(), self.vocab_file.filename())
}
}
pub fn resolve_model() -> EmbeddingModel {
resolve_model_with_source().0
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelSource {
Env(String),
Config(String),
Default,
}
pub fn resolve_model_with_source() -> (EmbeddingModel, ModelSource) {
let env_val = std::env::var("LEAN_CTX_EMBEDDING_MODEL").ok();
let embedding_cfg = crate::core::config::Config::load().embedding;
resolve_model_with_source_from(
env_val.as_deref(),
embedding_cfg.model.as_deref(),
embedding_cfg.dimensions,
default_model(),
)
}
#[must_use]
pub fn default_model() -> EmbeddingModel {
#[cfg(feature = "embeddings")]
{
static ORT_PRESENT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
let ort_present = *ORT_PRESENT
.get_or_init(|| crate::core::ort_environment::resolved_ort_dylib_path().is_ok());
if !ort_present {
return EmbeddingModel::PotionCode16M;
}
}
EmbeddingModel::DEFAULT
}
#[cfg(test)]
fn resolve_model_from(
env_val: Option<&str>,
config_val: Option<&str>,
config_dims: Option<usize>,
) -> EmbeddingModel {
resolve_model_with_source_from(env_val, config_val, config_dims, EmbeddingModel::DEFAULT).0
}
fn resolve_model_with_source_from(
env_val: Option<&str>,
config_val: Option<&str>,
config_dims: Option<usize>,
default: EmbeddingModel,
) -> (EmbeddingModel, ModelSource) {
for (source, raw) in [
("LEAN_CTX_EMBEDDING_MODEL", env_val),
("[embedding].model", config_val),
] {
let Some(name) = raw.map(str::trim).filter(|s| !s.is_empty()) else {
continue;
};
let origin = if source == "LEAN_CTX_EMBEDDING_MODEL" {
ModelSource::Env(name.to_string())
} else {
ModelSource::Config(name.to_string())
};
match EmbeddingModel::from_str_name(name) {
Some(EmbeddingModel::Custom(mut spec)) => {
spec.dimensions = config_dims;
if spec.revision.is_none() {
tracing::warn!(
"Custom embedding model {:?} has no revision pin — supply-chain best \
practice is `hf:{}@<commit-or-tag>` so upstream pushes can never \
silently change your index",
spec.repo,
spec.repo
);
}
return (EmbeddingModel::Custom(spec), origin);
}
Some(model) => return (model, origin),
None => {
tracing::warn!(
"Unknown embedding model {name:?} from {source}; using {default} instead \
(built-ins: minilm, nomic, potion-code — or hf:org/repo[@rev])"
);
}
}
}
(default, ModelSource::Default)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_model_is_minilm() {
assert_eq!(EmbeddingModel::DEFAULT, EmbeddingModel::AllMiniLmL6V2);
}
#[test]
fn from_str_name_variants() {
assert_eq!(
EmbeddingModel::from_str_name("minilm"),
Some(EmbeddingModel::AllMiniLmL6V2)
);
assert_eq!(
EmbeddingModel::from_str_name("nomic-embed-v1.5"),
Some(EmbeddingModel::NomicEmbedV1_5)
);
assert_eq!(
EmbeddingModel::from_str_name("nomic"),
Some(EmbeddingModel::NomicEmbedV1_5)
);
assert_eq!(
EmbeddingModel::from_str_name("default"),
Some(EmbeddingModel::AllMiniLmL6V2)
);
assert_eq!(EmbeddingModel::from_str_name("unknown"), None);
assert_eq!(EmbeddingModel::from_str_name("jina-code-v2"), None);
assert_eq!(EmbeddingModel::from_str_name("jina"), None);
}
#[test]
fn custom_hf_scheme_parses_repo_and_revision() {
let m = EmbeddingModel::from_str_name("hf:jinaai/jina-embeddings-v2-base-code").unwrap();
let EmbeddingModel::Custom(spec) = &m else {
panic!("expected custom")
};
assert_eq!(spec.repo, "jinaai/jina-embeddings-v2-base-code");
assert_eq!(spec.revision, None);
let m = EmbeddingModel::from_str_name("hf:org/model@abc123").unwrap();
let EmbeddingModel::Custom(spec) = &m else {
panic!("expected custom")
};
assert_eq!(spec.repo, "org/model");
assert_eq!(spec.revision.as_deref(), Some("abc123"));
}
#[test]
fn custom_hf_scheme_rejects_invalid_repos() {
for bad in [
"hf:",
"hf:no-slash",
"hf:too/many/slashes",
"hf:with space/repo",
"hf:/leading",
"hf:trailing/",
"hf:org/model@",
] {
let parsed = EmbeddingModel::from_str_name(bad);
if bad == "hf:org/model@" {
let Some(EmbeddingModel::Custom(spec)) = parsed else {
panic!("expected custom for {bad}")
};
assert_eq!(spec.revision, None);
} else {
assert_eq!(parsed, None, "{bad} should be rejected");
}
}
}
#[test]
fn custom_config_urls_and_storage() {
let m = EmbeddingModel::from_str_name("hf:Org/My_Model@v1.2").unwrap();
let cfg = m.config();
assert_eq!(
cfg.model_url(),
"https://huggingface.co/Org/My_Model/resolve/v1.2/onnx/model.onnx"
);
assert_eq!(
cfg.vocab_url(),
"https://huggingface.co/Org/My_Model/resolve/v1.2/tokenizer.json"
);
assert!(!cfg.vocab_file.is_wordpiece());
assert_eq!(cfg.name, "hf:Org/My_Model@v1.2");
assert_eq!(m.storage_dir_name(), "hf-org-my-model-v1-2");
}
#[test]
fn custom_storage_slugs_differ_per_revision() {
let a = EmbeddingModel::from_str_name("hf:org/model@aaa").unwrap();
let b = EmbeddingModel::from_str_name("hf:org/model@bbb").unwrap();
let c = EmbeddingModel::from_str_name("hf:org/model").unwrap();
let slugs = [
a.storage_dir_name(),
b.storage_dir_name(),
c.storage_dir_name(),
];
let unique: std::collections::HashSet<_> = slugs.iter().collect();
assert_eq!(unique.len(), 3);
}
#[test]
fn all_models_have_valid_configs() {
for model in EmbeddingModel::ALL {
let cfg = model.config();
assert!(!cfg.name.is_empty(), "`cfg.name` is empty");
assert!(!cfg.hf_repo.is_empty(), "`cfg.hf_repo` is empty");
assert!(cfg.dimensions > 0);
assert!(cfg.max_seq_len > 0);
assert!(cfg.model_min_bytes > 0);
assert!(cfg.vocab_min_bytes > 0);
}
}
#[test]
fn model_urls_are_valid() {
for model in EmbeddingModel::ALL {
let cfg = model.config();
let model_url = cfg.model_url();
let vocab_url = cfg.vocab_url();
assert!(model_url.starts_with("https://huggingface.co/"));
assert!(vocab_url.starts_with("https://huggingface.co/"));
let rev = cfg.revision.as_deref().unwrap_or("main");
assert!(
model_url.contains(&format!("resolve/{rev}/")),
"{model_url}"
);
}
}
#[test]
fn storage_dir_names_are_unique() {
let names: Vec<_> = EmbeddingModel::ALL
.iter()
.map(EmbeddingModel::storage_dir_name)
.collect();
let unique: std::collections::HashSet<_> = names.iter().collect();
assert_eq!(names.len(), unique.len());
}
#[test]
fn display_uses_model_name() {
assert_eq!(
format!("{}", EmbeddingModel::AllMiniLmL6V2),
"all-MiniLM-L6-v2"
);
assert_eq!(
format!("{}", EmbeddingModel::NomicEmbedV1_5),
"nomic-embed-text-v1.5"
);
}
#[test]
fn resolve_defaults_when_nothing_set() {
assert_eq!(
resolve_model_from(None, None, None),
EmbeddingModel::DEFAULT
);
assert_eq!(
resolve_model_from(Some(""), Some(" "), None),
EmbeddingModel::DEFAULT
);
}
#[test]
fn config_selects_model_when_env_unset() {
assert_eq!(
resolve_model_from(None, Some("nomic"), None),
EmbeddingModel::NomicEmbedV1_5
);
assert_eq!(
resolve_model_from(None, Some("minilm"), None),
EmbeddingModel::AllMiniLmL6V2
);
}
#[test]
fn resolution_reports_the_setting_that_won() {
assert_eq!(
resolve_model_with_source_from(
Some(" nomic "),
Some("minilm"),
None,
EmbeddingModel::DEFAULT
)
.1,
ModelSource::Env("nomic".into())
);
assert_eq!(
resolve_model_with_source_from(
Some("bogus"),
Some("nomic"),
None,
EmbeddingModel::DEFAULT
)
.1,
ModelSource::Config("nomic".into())
);
assert_eq!(
resolve_model_with_source_from(None, Some("nope"), None, EmbeddingModel::DEFAULT).1,
ModelSource::Default
);
}
#[test]
fn env_var_overrides_config() {
assert_eq!(
resolve_model_from(Some("minilm"), Some("nomic"), None),
EmbeddingModel::AllMiniLmL6V2
);
}
#[test]
fn unknown_name_falls_through_then_defaults() {
assert_eq!(
resolve_model_from(Some("bogus"), Some("nomic"), None),
EmbeddingModel::NomicEmbedV1_5
);
assert_eq!(
resolve_model_from(Some("bogus"), Some("nope"), None),
EmbeddingModel::DEFAULT
);
assert_eq!(
resolve_model_from(Some(" "), Some("nomic"), None),
EmbeddingModel::NomicEmbedV1_5
);
}
#[test]
fn resolve_custom_picks_up_config_dimensions() {
let m = resolve_model_from(None, Some("hf:org/model@pin"), Some(1024));
let EmbeddingModel::Custom(spec) = m else {
panic!("expected custom")
};
assert_eq!(spec.dimensions, Some(1024));
assert_eq!(spec.revision.as_deref(), Some("pin"));
}
#[test]
fn nomic_has_prefixes() {
let cfg = EmbeddingModel::NomicEmbedV1_5.config();
assert!(cfg.query_prefix.is_some());
assert!(cfg.document_prefix.is_some());
assert!(!cfg.needs_token_type_ids);
}
#[test]
fn minilm_is_wordpiece() {
let cfg = EmbeddingModel::AllMiniLmL6V2.config();
assert!(cfg.vocab_file.is_wordpiece());
}
#[test]
fn builtin_models_have_valid_vocab_sources() {
for model in EmbeddingModel::ALL {
let cfg = model.config();
match cfg.format {
ModelFormat::Onnx => assert!(
cfg.vocab_file.is_wordpiece(),
"{model} should use WordPiece vocab.txt"
),
ModelFormat::StaticSafetensors => assert_eq!(
cfg.vocab_file.filename(),
"tokenizer.json",
"{model} ships tokenizer.json"
),
}
}
}
#[test]
fn custom_models_use_tokenizer_json() {
let m = EmbeddingModel::from_str_name("hf:org/model").unwrap();
assert!(!m.config().vocab_file.is_wordpiece());
assert_eq!(m.config().vocab_file.filename(), "tokenizer.json");
}
}