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;