Skip to main content

remem/retrieval/embedding/
local_semantic.rs

1use std::path::{Component, Path, PathBuf};
2
3use anyhow::{bail, Context, Result};
4use serde::{Deserialize, Serialize};
5use sha2::{Digest, Sha256};
6
7use super::{EmbeddingConfig, TextEmbedding};
8
9#[cfg(feature = "local-onnx")]
10mod download;
11#[cfg(feature = "local-onnx")]
12mod fs_cleanup;
13#[cfg(test)]
14mod hash_counter;
15mod manifest;
16mod model_path;
17#[cfg(feature = "local-onnx")]
18mod runtime;
19#[cfg(test)]
20mod test_support;
21#[cfg(all(windows, feature = "local-onnx"))]
22mod windows_cleanup;
23#[cfg(windows)]
24mod windows_model_root;
25#[cfg(windows)]
26mod windows_security;
27
28#[cfg(not(windows))]
29use model_path::install_dir_for_preset;
30pub(super) use model_path::model_root;
31
32#[cfg(test)]
33use hash_counter::ModelFileHashCounter;
34use manifest::read_verified_manifest_compatible;
35#[cfg(feature = "local-onnx")]
36use manifest::with_model_read_lock;
37#[cfg(test)]
38use manifest::{collect_model_artifacts, write_manifest};
39#[cfg(feature = "local-onnx")]
40use manifest::{open_or_create_model_lock, read_verified_manifest_unlocked};
41#[cfg(test)]
42pub(crate) use test_support::install_test_model;
43#[cfg(all(test, feature = "local-onnx"))]
44pub(crate) use test_support::{
45    fail_next_test_model_embed_generic, fail_next_test_model_embed_unavailable,
46    fail_test_model_runtime_readiness, install_test_model_v1, install_untrusted_test_model,
47    test_model_runtime_file,
48};
49
50pub(super) const DEFAULT_LOCAL_SEMANTIC_DIMENSIONS: usize = 384;
51pub(super) const DEFAULT_LOCAL_SEMANTIC_MODEL: &str = "fastembed-intfloat-multilingual-e5-small-v1";
52
53const MANIFEST_FILE: &str = "remem-model-manifest.json";
54const MANIFEST_SCHEMA_VERSION: u32 = 2;
55const FASTEMBED_RUNTIME: &str = "fastembed-rs/onnxruntime";
56const HUGGING_FACE_BASE_URL: &str = "https://huggingface.co";
57#[cfg(feature = "local-onnx")]
58const HUGGING_FACE_ENDPOINT_ENV: &str = "HF_ENDPOINT";
59#[cfg(feature = "local-onnx")]
60pub(super) const AUTO_EVALUATED_DEFAULT_ARTIFACT_SHA256: &str =
61    "3970612d6f31b81d1dc30ddac0099da273b5753d1a07412e8390cf799e7836a6";
62const MODEL_DOWNLOAD_LOCK_FILE: &str = ".remem-model-download.lock";
63const MODEL_STATE_LOCK_FILE: &str = ".remem-model-state.lock";
64const TOKENIZER_RUNTIME_FILES: &[&str] = &[
65    "tokenizer.json",
66    "config.json",
67    "special_tokens_map.json",
68    "tokenizer_config.json",
69];
70#[derive(Debug)]
71struct LocalEmbeddingModelUnavailableError(String);
72
73impl std::fmt::Display for LocalEmbeddingModelUnavailableError {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        f.write_str(&self.0)
76    }
77}
78
79impl std::error::Error for LocalEmbeddingModelUnavailableError {}
80
81pub(super) fn is_model_unavailable_error(error: &anyhow::Error) -> bool {
82    error
83        .downcast_ref::<LocalEmbeddingModelUnavailableError>()
84        .is_some()
85}
86
87pub(super) fn model_unavailable_error(reason: impl Into<String>) -> anyhow::Error {
88    LocalEmbeddingModelUnavailableError(reason.into()).into()
89}
90
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92pub(super) enum LocalEmbeddingInputKind {
93    Query,
94    Passage,
95    Generic,
96}
97
98#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
99pub(super) enum LocalEmbeddingPreset {
100    MultilingualE5Small,
101    BgeM3,
102}
103
104impl LocalEmbeddingPreset {
105    fn all() -> &'static [Self] {
106        &[Self::MultilingualE5Small, Self::BgeM3]
107    }
108
109    fn default() -> Self {
110        Self::MultilingualE5Small
111    }
112
113    fn parse(raw: &str) -> Result<Self> {
114        match raw.trim().to_ascii_lowercase().as_str() {
115            "" => Ok(Self::default()),
116            "multilingual-e5-small"
117            | "intfloat/multilingual-e5-small"
118            | DEFAULT_LOCAL_SEMANTIC_MODEL => Ok(Self::MultilingualE5Small),
119            "bge-m3" | "baai/bge-m3" | "fastembed-bge-m3-v1" => Ok(Self::BgeM3),
120            other => bail!(
121                "unsupported local embedding model preset {other}; recognized local presets: multilingual-e5-small, bge-m3"
122            ),
123        }
124    }
125
126    fn label(self) -> &'static str {
127        match self {
128            Self::MultilingualE5Small => "multilingual-e5-small",
129            Self::BgeM3 => "bge-m3",
130        }
131    }
132
133    fn model_id(self) -> &'static str {
134        match self {
135            Self::MultilingualE5Small => DEFAULT_LOCAL_SEMANTIC_MODEL,
136            Self::BgeM3 => "fastembed-bge-m3-v1",
137        }
138    }
139
140    fn upstream_model(self) -> &'static str {
141        match self {
142            Self::MultilingualE5Small => "intfloat/multilingual-e5-small",
143            Self::BgeM3 => "BAAI/bge-m3",
144        }
145    }
146
147    fn source_url(self) -> String {
148        format!("{HUGGING_FACE_BASE_URL}/{}", self.upstream_model())
149    }
150
151    fn dimensions(self) -> usize {
152        match self {
153            Self::MultilingualE5Small => DEFAULT_LOCAL_SEMANTIC_DIMENSIONS,
154            Self::BgeM3 => 1024,
155        }
156    }
157
158    fn cache_repo_dir(self) -> String {
159        format!("models--{}", self.upstream_model()).replace('/', "--")
160    }
161
162    fn model_file(self) -> &'static str {
163        "onnx/model.onnx"
164    }
165
166    fn additional_model_files(self) -> &'static [&'static str] {
167        match self {
168            Self::MultilingualE5Small => &[],
169            Self::BgeM3 => &["onnx/model.onnx_data", "onnx/Constant_7_attr__value"],
170        }
171    }
172
173    fn required_runtime_files(self) -> impl Iterator<Item = &'static str> {
174        std::iter::once(self.model_file())
175            .chain(self.additional_model_files().iter().copied())
176            .chain(TOKENIZER_RUNTIME_FILES.iter().copied())
177    }
178
179    #[cfg(feature = "local-onnx")]
180    fn prefix_input(self, text: &str, kind: LocalEmbeddingInputKind) -> String {
181        match (self, kind) {
182            (Self::MultilingualE5Small, LocalEmbeddingInputKind::Query) => {
183                format!("query: {text}")
184            }
185            (Self::MultilingualE5Small, LocalEmbeddingInputKind::Passage) => {
186                format!("passage: {text}")
187            }
188            _ => text.to_string(),
189        }
190    }
191
192    #[cfg(feature = "local-onnx")]
193    fn fastembed_model(self) -> fastembed::EmbeddingModel {
194        match self {
195            Self::MultilingualE5Small => fastembed::EmbeddingModel::MultilingualE5Small,
196            Self::BgeM3 => fastembed::EmbeddingModel::BGEM3,
197        }
198    }
199}
200
201#[derive(Debug, Clone, PartialEq, Eq)]
202pub(super) struct LocalModelProfile {
203    pub(super) model: String,
204    pub(super) dimensions: usize,
205    pub(super) install_dir: PathBuf,
206    pub(super) artifact_sha256: String,
207}
208
209#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
210pub struct LocalEmbeddingDownloadReport {
211    pub preset: String,
212    pub model_id: String,
213    pub upstream_model: String,
214    pub dimensions: usize,
215    pub install_dir: String,
216    pub files_verified: usize,
217    pub artifact_sha256: String,
218}
219
220#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
221pub struct LocalEmbeddingModelInventory {
222    pub preset: String,
223    pub model_id: String,
224    pub upstream_model: String,
225    pub dimensions: usize,
226    pub install_dir: String,
227    pub installed: bool,
228    pub checksum_verified: bool,
229    pub unavailable_reason: Option<String>,
230}
231
232#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
233pub struct LocalEmbeddingInventoryReport {
234    pub model_root: String,
235    pub configured_preset: String,
236    pub models: Vec<LocalEmbeddingModelInventory>,
237}
238
239#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
240struct LocalModelManifest {
241    schema_version: u32,
242    preset: String,
243    model_id: String,
244    upstream_model: String,
245    dimensions: usize,
246    runtime: String,
247    source_url: Option<String>,
248    downloaded_at_epoch: i64,
249    files: Vec<LocalModelFile>,
250    #[serde(default)]
251    symlinks: Vec<LocalModelSymlink>,
252}
253
254#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
255struct LocalModelFile {
256    path: String,
257    sha256: String,
258    #[serde(default, skip_serializing_if = "Option::is_none")]
259    source_sha256: Option<String>,
260    bytes: u64,
261}
262
263#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
264struct LocalModelSymlink {
265    path: String,
266    link_target: String,
267    resolved_path: String,
268}
269
270#[cfg(feature = "local-onnx")]
271pub(super) fn installed_model_profile(config: &EmbeddingConfig) -> Result<LocalModelProfile> {
272    let preset = configured_local_preset_or_default(config)?;
273    verified_profile_for_preset(config, preset)
274}
275
276#[cfg(feature = "local-onnx")]
277pub(crate) fn with_configured_model_read_lock<T>(
278    config: &EmbeddingConfig,
279    operation: impl FnOnce() -> Result<T>,
280) -> Result<T> {
281    let preset = configured_local_preset_or_default(config)?;
282    #[cfg(windows)]
283    let _windows_install = if config.provider == super::EmbeddingProvider::Auto {
284        windows_model_root::create_managed_install(config, preset)?
285    } else {
286        windows_model_root::open_managed_install(config, preset, false)?
287            .ok_or_else(|| windows_model_root::missing_install_error())?
288    };
289    #[cfg(windows)]
290    let install_dir = _windows_install.install_dir().to_path_buf();
291    #[cfg(not(windows))]
292    let install_dir = install_dir_for_preset(config, preset)?;
293    #[cfg(not(windows))]
294    if config.provider == super::EmbeddingProvider::Auto {
295        std::fs::create_dir_all(&install_dir)
296            .with_context(|| format!("create model-state pin dir {}", install_dir.display()))?;
297    }
298    with_model_read_lock(&install_dir, operation)
299}
300
301#[cfg(not(feature = "local-onnx"))]
302pub(super) fn installed_model_profile(config: &EmbeddingConfig) -> Result<LocalModelProfile> {
303    let preset = configured_local_preset_or_default(config)?;
304    #[cfg(windows)]
305    windows_model_root::checked_model_root(config)?;
306    Err(model_unavailable_error(format!(
307        "local semantic embedding runtime is not built; rebuild remem with the local-onnx feature to use {}",
308        preset.label()
309    )))
310}
311
312#[cfg(not(feature = "local-onnx"))]
313pub(crate) fn with_configured_model_read_lock<T>(
314    config: &EmbeddingConfig,
315    _operation: impl FnOnce() -> Result<T>,
316) -> Result<T> {
317    let _ = installed_model_profile(config)?;
318    Err(model_unavailable_error(
319        "local semantic embedding runtime is not built",
320    ))
321}
322
323pub(super) fn auto_installed_model_profile(
324    config: &EmbeddingConfig,
325) -> Result<Option<LocalModelProfile>> {
326    let preset = configured_local_preset_or_default(config)?;
327    #[cfg(windows)]
328    let _windows_install = match windows_model_root::open_managed_install(config, preset, true)? {
329        Some(install) => install,
330        None => return Ok(None),
331    };
332    #[cfg(windows)]
333    let install_dir = _windows_install.install_dir().to_path_buf();
334    #[cfg(not(windows))]
335    let install_dir = install_dir_for_preset(config, preset)?;
336    match std::fs::symlink_metadata(&install_dir) {
337        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
338        Ok(metadata) if metadata.file_type().is_symlink() => Err(model_unavailable_error(format!(
339            "local embedding model install path is a symlink: {}",
340            install_dir.display()
341        ))),
342        Ok(metadata) if !metadata.file_type().is_dir() => Err(model_unavailable_error(format!(
343            "local embedding model install path is not a directory: {}",
344            install_dir.display()
345        ))),
346        Ok(_) if model_state_coordination_only(&install_dir)? => Ok(None),
347        Ok(_) => auto_verified_model_profile(config, preset).map(Some),
348        Err(error) => Err(model_unavailable_error(format!(
349            "inspect local embedding model {} in {}: {error:#}",
350            preset.label(),
351            install_dir.display()
352        ))),
353    }
354}
355
356fn model_state_coordination_only(install_dir: &Path) -> Result<bool> {
357    let entries = std::fs::read_dir(install_dir)
358        .context("read model-state coordination directory")?
359        .collect::<std::io::Result<Vec<_>>>()?;
360    Ok(entries.len() == 1
361        && entries[0].file_name() == MODEL_STATE_LOCK_FILE
362        && entries[0].file_type()?.is_file())
363}
364
365pub(super) fn download_model(model: Option<&str>) -> Result<LocalEmbeddingDownloadReport> {
366    let config = super::resolve_embedding_config()?;
367    let preset = match model {
368        Some(raw) => LocalEmbeddingPreset::parse(raw)?,
369        None => configured_local_preset_or_default(&config)?,
370    };
371
372    #[cfg(not(feature = "local-onnx"))]
373    {
374        #[cfg(windows)]
375        windows_model_root::checked_model_root(&config)?;
376        bail!(
377            "local semantic embedding runtime is not built; rebuild remem with the local-onnx feature to download {}",
378            preset.label()
379        );
380    }
381
382    #[cfg(feature = "local-onnx")]
383    {
384        #[cfg(windows)]
385        let _windows_install = windows_model_root::create_managed_install(&config, preset)?;
386        #[cfg(windows)]
387        let install_dir = _windows_install.install_dir().to_path_buf();
388        #[cfg(not(windows))]
389        let install_dir = install_dir_for_preset(&config, preset)?;
390        #[cfg(not(windows))]
391        std::fs::create_dir_all(&install_dir).with_context(|| {
392            format!("create local embedding model dir {}", install_dir.display())
393        })?;
394        let (lock_path, download_lock) =
395            open_or_create_model_lock(&install_dir, MODEL_DOWNLOAD_LOCK_FILE)
396                .context("open local model download serialization lock")?;
397        fs2::FileExt::lock_exclusive(&download_lock)
398            .with_context(|| format!("lock local model download {}", lock_path.display()))?;
399        let (_state_lock_path, state_lock_file) =
400            open_or_create_model_lock(&install_dir, MODEL_STATE_LOCK_FILE)
401                .context("initialize local model state lock")?;
402        drop(state_lock_file);
403        let staging = download::materialize_hugging_face_artifacts(preset, &install_dir)?;
404        let candidate = (|| {
405            let prepared = download::prepare_downloaded_model(
406                preset,
407                staging.path(),
408                chrono::Utc::now().timestamp(),
409            )?;
410            let imported = download::import_immutable_candidate(
411                staging.path(),
412                &install_dir,
413                preset,
414                &prepared.manifest,
415            )?;
416            Ok((prepared, imported))
417        })();
418        let (prepared, imported) = match candidate {
419            Ok(candidate) => candidate,
420            Err(error) => return download::cleanup_staging_after_error(staging, error),
421        };
422        staging.cleanup()?;
423        let (state_lock_path, state_lock) =
424            open_or_create_model_lock(&install_dir, MODEL_STATE_LOCK_FILE)
425                .context("open local model state lock for publish")?;
426        fs2::FileExt::lock_exclusive(&state_lock).with_context(|| {
427            format!(
428                "lock local model state for publish {}",
429                state_lock_path.display()
430            )
431        })?;
432        let verified = download::activate_candidate_manifest(
433            &install_dir,
434            preset,
435            prepared.manifest,
436            prepared.artifact_sha256,
437            imported,
438        )?;
439        let artifact_sha256 = verified.artifact_sha256;
440        let manifest = verified.manifest;
441        Ok(LocalEmbeddingDownloadReport {
442            preset: manifest.preset,
443            model_id: manifest.model_id,
444            upstream_model: manifest.upstream_model,
445            dimensions: manifest.dimensions,
446            install_dir: install_dir.display().to_string(),
447            files_verified: manifest.files.len(),
448            artifact_sha256,
449        })
450    }
451}
452
453pub(super) fn inventory() -> Result<LocalEmbeddingInventoryReport> {
454    let config = super::resolve_embedding_config()?;
455    #[cfg(windows)]
456    let root = windows_model_root::checked_model_root(&config)?;
457    #[cfg(not(windows))]
458    let root = model_root(&config)?;
459    let configured = configured_local_preset_or_default(&config)?;
460    let models = LocalEmbeddingPreset::all()
461        .iter()
462        .copied()
463        .map(|preset| inventory_for_preset(&config, preset))
464        .collect::<Result<Vec<_>>>()?;
465    Ok(LocalEmbeddingInventoryReport {
466        model_root: root.display().to_string(),
467        configured_preset: configured.label().to_string(),
468        models,
469    })
470}
471
472#[cfg(feature = "local-onnx")]
473pub(super) fn embed_text(
474    text: &str,
475    config: &EmbeddingConfig,
476    kind: LocalEmbeddingInputKind,
477) -> Result<TextEmbedding> {
478    let preset = configured_local_preset_or_default(config)?;
479    #[cfg(windows)]
480    let _windows_install = windows_model_root::open_managed_install(config, preset, false)?
481        .ok_or_else(|| windows_model_root::missing_install_error())?;
482    #[cfg(windows)]
483    let install_dir = _windows_install.install_dir().to_path_buf();
484    #[cfg(not(windows))]
485    let install_dir = install_dir_for_preset(config, preset)?;
486    #[cfg(test)]
487    if let Some(failure) = test_support::take_next_embed_failure(&install_dir)? {
488        return match failure {
489            test_support::TestEmbedFailure::ModelUnavailable(reason) => {
490                Err(model_unavailable_error(reason))
491            }
492            test_support::TestEmbedFailure::Generic(reason) => Err(anyhow::anyhow!(reason)),
493        };
494    }
495    read_verified_manifest_compatible(&install_dir, Some(preset)).map_err(|error| {
496        model_unavailable_error(format!(
497            "local embedding model {} is not ready in {}: {error:#}",
498            preset.label(),
499            install_dir.display()
500        ))
501    })?;
502    with_model_read_lock(&install_dir, || {
503        let verified =
504            read_verified_manifest_unlocked(&install_dir, Some(preset)).map_err(|error| {
505                model_unavailable_error(format!(
506                    "local embedding model {} is not ready in {}: {error:#}",
507                    preset.label(),
508                    install_dir.display()
509                ))
510            })?;
511        if config.provider == super::EmbeddingProvider::Auto {
512            require_auto_evaluated_artifact(&install_dir, preset, &verified.artifact_sha256)?;
513        }
514        let profile = profile_from_verified_manifest(&install_dir, &verified);
515        let values = runtime::embed_with_verified_model(
516            preset,
517            &install_dir,
518            &verified.manifest,
519            &profile.artifact_sha256,
520            text,
521            kind,
522        )?;
523        if values.len() != profile.dimensions {
524            bail!(
525                "local embedding model {} returned {} dimensions, expected {}",
526                profile.model,
527                values.len(),
528                profile.dimensions
529            );
530        }
531        TextEmbedding::new(profile.model, values)
532    })
533}
534
535#[cfg(not(feature = "local-onnx"))]
536pub(super) fn embed_text(
537    _text: &str,
538    config: &EmbeddingConfig,
539    _kind: LocalEmbeddingInputKind,
540) -> Result<TextEmbedding> {
541    let preset = configured_local_preset_or_default(config)?;
542    #[cfg(windows)]
543    windows_model_root::checked_model_root(config)?;
544    Err(model_unavailable_error(format!(
545        "local semantic embedding runtime is not built; rebuild remem with the local-onnx feature to use {}",
546        preset.label()
547    )))
548}
549
550fn configured_preset(config: &EmbeddingConfig) -> Result<LocalEmbeddingPreset> {
551    let raw = config.model.trim();
552    if raw.is_empty() || raw == super::OPENAI_DEFAULT_MODEL {
553        return Ok(LocalEmbeddingPreset::default());
554    }
555    LocalEmbeddingPreset::parse(raw)
556}
557
558pub(super) fn configured_model_id(config: &EmbeddingConfig) -> Result<String> {
559    Ok(configured_preset(config)?.model_id().to_string())
560}
561
562fn configured_local_preset_or_default(config: &EmbeddingConfig) -> Result<LocalEmbeddingPreset> {
563    if config.provider == super::EmbeddingProvider::Local {
564        configured_preset(config)
565    } else {
566        Ok(LocalEmbeddingPreset::default())
567    }
568}
569
570#[cfg(feature = "local-onnx")]
571fn verified_profile_for_preset(
572    config: &EmbeddingConfig,
573    preset: LocalEmbeddingPreset,
574) -> Result<LocalModelProfile> {
575    verified_profile_for_preset_with_policy(config, preset, false)
576}
577
578#[cfg(feature = "local-onnx")]
579fn auto_verified_model_profile(
580    config: &EmbeddingConfig,
581    preset: LocalEmbeddingPreset,
582) -> Result<LocalModelProfile> {
583    verified_profile_for_preset_with_policy(config, preset, true)
584}
585
586#[cfg(not(feature = "local-onnx"))]
587fn auto_verified_model_profile(
588    config: &EmbeddingConfig,
589    _preset: LocalEmbeddingPreset,
590) -> Result<LocalModelProfile> {
591    installed_model_profile(config)
592}
593
594#[cfg(feature = "local-onnx")]
595fn verified_profile_for_preset_with_policy(
596    config: &EmbeddingConfig,
597    preset: LocalEmbeddingPreset,
598    enforce_auto_evaluated_artifact: bool,
599) -> Result<LocalModelProfile> {
600    #[cfg(windows)]
601    let _windows_install = windows_model_root::open_managed_install(config, preset, false)?
602        .ok_or_else(|| windows_model_root::missing_install_error())?;
603    #[cfg(windows)]
604    let install_dir = _windows_install.install_dir().to_path_buf();
605    #[cfg(not(windows))]
606    let install_dir = install_dir_for_preset(config, preset)?;
607    read_verified_manifest_compatible(&install_dir, Some(preset)).map_err(|error| {
608        model_unavailable_error(format!(
609            "local embedding model {} is not ready in {}: {error:#}",
610            preset.label(),
611            install_dir.display()
612        ))
613    })?;
614    with_model_read_lock(&install_dir, || {
615        let verified = read_verified_manifest_unlocked(&install_dir, Some(preset))?;
616        if enforce_auto_evaluated_artifact {
617            require_auto_evaluated_artifact(&install_dir, preset, &verified.artifact_sha256)?;
618        }
619        runtime::ensure_verified_model_ready(
620            preset,
621            &install_dir,
622            &verified.manifest,
623            &verified.artifact_sha256,
624        )?;
625        Ok(profile_from_verified_manifest(&install_dir, &verified))
626    })
627    .map_err(|error| {
628        model_unavailable_error(format!(
629            "local embedding model {} is not ready in {}: {error:#}",
630            preset.label(),
631            install_dir.display()
632        ))
633    })
634}
635
636#[cfg(feature = "local-onnx")]
637fn auto_artifact_is_trusted(_install_dir: &Path, artifact_sha256: &str) -> Result<bool> {
638    if artifact_sha256 == AUTO_EVALUATED_DEFAULT_ARTIFACT_SHA256 {
639        return Ok(true);
640    }
641    #[cfg(test)]
642    if test_support::is_test_auto_artifact_trusted(_install_dir)? {
643        return Ok(true);
644    }
645    Ok(false)
646}
647
648#[cfg(feature = "local-onnx")]
649fn require_auto_evaluated_artifact(
650    install_dir: &Path,
651    preset: LocalEmbeddingPreset,
652    artifact_sha256: &str,
653) -> Result<()> {
654    if auto_artifact_is_trusted(install_dir, artifact_sha256)? {
655        return Ok(());
656    }
657    Err(model_unavailable_error(format!(
658        "automatic local embedding requires evaluated {} artifact sha256:{}; installed artifact sha256:{} is not trusted for Auto. Upgrade remem or redownload the model from {}",
659        preset.label(),
660        AUTO_EVALUATED_DEFAULT_ARTIFACT_SHA256,
661        artifact_sha256,
662        HUGGING_FACE_BASE_URL
663    )))
664}
665
666#[cfg(feature = "local-onnx")]
667fn profile_from_verified_manifest(
668    install_dir: &Path,
669    verified: &manifest::VerifiedLocalManifest,
670) -> LocalModelProfile {
671    LocalModelProfile {
672        model: format!(
673            "{}@sha256:{}",
674            verified.manifest.model_id, verified.artifact_sha256
675        ),
676        dimensions: verified.manifest.dimensions,
677        install_dir: install_dir.to_path_buf(),
678        artifact_sha256: verified.artifact_sha256.clone(),
679    }
680}
681
682fn inventory_for_preset(
683    config: &EmbeddingConfig,
684    preset: LocalEmbeddingPreset,
685) -> Result<LocalEmbeddingModelInventory> {
686    #[cfg(windows)]
687    let _windows_install = match windows_model_root::open_managed_install(config, preset, true)? {
688        Some(install) => install,
689        None => {
690            return Ok(LocalEmbeddingModelInventory {
691                preset: preset.label().to_string(),
692                model_id: preset.model_id().to_string(),
693                upstream_model: preset.upstream_model().to_string(),
694                dimensions: preset.dimensions(),
695                install_dir: model_root(config)?
696                    .join(preset.model_id())
697                    .display()
698                    .to_string(),
699                installed: false,
700                checksum_verified: false,
701                unavailable_reason: Some("local embedding model is not installed".to_string()),
702            });
703        }
704    };
705    #[cfg(windows)]
706    let install_dir = _windows_install.install_dir().to_path_buf();
707    #[cfg(not(windows))]
708    let install_dir = install_dir_for_preset(config, preset)?;
709    match read_verified_manifest_compatible(&install_dir, Some(preset)) {
710        Ok(verified) => Ok(LocalEmbeddingModelInventory {
711            preset: verified.manifest.preset,
712            model_id: verified.manifest.model_id,
713            upstream_model: verified.manifest.upstream_model,
714            dimensions: verified.manifest.dimensions,
715            install_dir: install_dir.display().to_string(),
716            installed: true,
717            checksum_verified: is_sha256_hex(&verified.artifact_sha256),
718            unavailable_reason: None,
719        }),
720        Err(error) => Ok(LocalEmbeddingModelInventory {
721            preset: preset.label().to_string(),
722            model_id: preset.model_id().to_string(),
723            upstream_model: preset.upstream_model().to_string(),
724            dimensions: preset.dimensions(),
725            install_dir: install_dir.display().to_string(),
726            installed: false,
727            checksum_verified: false,
728            unavailable_reason: Some(error.to_string()),
729        }),
730    }
731}
732
733fn checked_relative_path(raw: &str) -> Result<PathBuf> {
734    let path = PathBuf::from(raw);
735    if path.is_absolute() {
736        bail!("manifest path must be relative: {raw}");
737    }
738    if path
739        .components()
740        .any(|component| !matches!(component, Component::Normal(_)))
741    {
742        bail!("manifest path must not contain parent/current components: {raw}");
743    }
744    Ok(path)
745}
746
747fn source_sha256_from_hf_blob_path(relative: &str, actual_sha256: &str) -> Result<Option<String>> {
748    let parts = relative.split('/').collect::<Vec<_>>();
749    let Some(file_name) = parts.last().copied() else {
750        return Ok(None);
751    };
752    if parts.len() < 2 || parts[parts.len() - 2] != "blobs" || !is_sha256_hex(file_name) {
753        return Ok(None);
754    }
755    if file_name != actual_sha256 {
756        bail!(
757            "source checksum mismatch for Hugging Face cache blob {relative}: expected {file_name}, got {actual_sha256}"
758        );
759    }
760    Ok(Some(file_name.to_string()))
761}
762
763fn is_sha256_hex(value: &str) -> bool {
764    value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit())
765}
766
767fn sha256_file(path: &Path) -> Result<String> {
768    #[cfg(test)]
769    let pending_hash = hash_counter::PendingModelFileHash::for_path(path)?;
770    let mut file = std::fs::File::open(path).with_context(|| format!("open {}", path.display()))?;
771    let mut hasher = Sha256::new();
772    let mut buffer = [0_u8; 64 * 1024];
773    loop {
774        let read = std::io::Read::read(&mut file, &mut buffer)
775            .with_context(|| format!("read {}", path.display()))?;
776        if read == 0 {
777            break;
778        }
779        hasher.update(&buffer[..read]);
780    }
781    let sha256 = hasher
782        .finalize()
783        .iter()
784        .map(|byte| format!("{byte:02x}"))
785        .collect();
786    #[cfg(test)]
787    pending_hash.record()?;
788    Ok(sha256)
789}
790
791#[cfg(test)]
792mod tests;