//! [`Embedder`] adapters (LOT C, étage 2 vectoriel — §14.5.3). //! //! The étage-2 semantic recall is driven by a declarative //! [`EmbedderProfile`]; this module turns such a profile into a concrete //! [`Embedder`] via [`embedder_from_profile`]. //! //! ## Zero heavy dependency by design //! //! IdeA's founding posture is **default `none`, nothing imposed, zero //! dependency** (ARCHITECTURE §14.5.3). The concrete engines therefore ship as //! **documented stubs**: //! //! - [`EmbedderStrategy::LocalOnnx`] / [`EmbedderStrategy::LocalServer`] / //! [`EmbedderStrategy::Api`] map to [`StubEmbedder`], which returns a clean //! [`EmbedderError::Unsupported`] (never a panic) — the real ONNX/HTTP //! integration is an explicit follow-up that would add the heavy crate behind a //! feature, here and only here; //! - [`EmbedderStrategy::None`] yields no embedder (`None`): recall stays the //! dependency-free naïve étage 1. //! //! [`HashEmbedder`] is a small **deterministic, dependency-free** embedder used to //! exercise the whole vector pipeline (it is not a stub: it produces stable, //! repeatable vectors from a hashing bag-of-words, good enough for tests and for a //! trivial offline fallback). use std::path::PathBuf; use std::sync::Arc; use async_trait::async_trait; use serde::{Deserialize, Serialize}; use domain::ports::{ Embedder, EmbedderEnvInspector, EmbedderEnvReport, EmbedderError, EmbedderPromptDismissal, EmbedderPromptStore, FileSystem, FsError, RemotePath, StoreError, }; use domain::profile::{EmbedderProfile, EmbedderStrategy}; use domain::project::ProjectPath; /// Whether this binary was compiled with the HTTP embedding capability /// (`localServer`/`api` real engines + Ollama detection). Reported honestly to the /// C2/C3 UI so it only offers strategies actually wired into *this* build. pub const VECTOR_HTTP_ENABLED: bool = cfg!(feature = "vector-http"); /// Whether this binary was compiled with the in-process ONNX embedding capability /// (`localOnnx` real engine via `fastembed`). Reported honestly to the C2/C3 UI. pub const VECTOR_ONNX_ENABLED: bool = cfg!(feature = "vector-onnx"); /// Builds the concrete [`Embedder`] for a profile, or `None` when the strategy is /// [`EmbedderStrategy::None`] (recall stays naïve). /// /// The remaining concrete engines are stubs ([`StubEmbedder`]) until their real /// HTTP/ONNX integration is enabled via the matching feature; they fail cleanly /// with [`EmbedderError::Unsupported`] rather than panic, so a composing recall /// degrades to naïve. /// /// `onnx_cache_dir` is the directory under which a `localOnnx` engine caches its /// model files (`/embedders/onnx`, see [`ONNX_CACHE_SUBDIR`]). It is /// ignored by every other strategy. #[must_use] pub fn embedder_from_profile( profile: &EmbedderProfile, onnx_cache_dir: &std::path::Path, ) -> Option> { let _ = onnx_cache_dir; // used only under `vector-onnx`; silence the unused warning otherwise. match profile.strategy { EmbedderStrategy::None => None, EmbedderStrategy::LocalOnnx => { #[cfg(feature = "vector-onnx")] { Some(Box::new(OnnxEmbedder::from_profile( profile, onnx_cache_dir, ))) } #[cfg(not(feature = "vector-onnx"))] { Some(Box::new(StubEmbedder::new( profile.id.clone(), profile.dimension, "localOnnx", ))) } } // Real HTTP engines under the `vector-http` feature; the dependency-free // build keeps the documented stub so the default posture is unchanged. EmbedderStrategy::LocalServer | EmbedderStrategy::Api => { #[cfg(feature = "vector-http")] { Some(Box::new(HttpEmbedder::from_profile(profile))) } #[cfg(not(feature = "vector-http"))] { let strategy = if profile.strategy == EmbedderStrategy::Api { "api" } else { "localServer" }; Some(Box::new(StubEmbedder::new( profile.id.clone(), profile.dimension, strategy, ))) } } } } /// A placeholder [`Embedder`] for a not-yet-implemented concrete strategy. /// /// Every [`embed`](Embedder::embed) call returns [`EmbedderError::Unsupported`] /// (never a panic), naming the strategy. It exists so the wiring, profiles, and /// adapters are all in place and testable today; swapping in the real engine is a /// localised follow-up. pub struct StubEmbedder { id: String, dimension: usize, strategy: &'static str, } impl StubEmbedder { /// Builds a stub for `strategy` advertising `dimension`. #[must_use] pub fn new(id: impl Into, dimension: usize, strategy: &'static str) -> Self { Self { id: id.into(), dimension, strategy, } } } #[async_trait] impl Embedder for StubEmbedder { fn id(&self) -> &str { &self.id } async fn embed(&self, _texts: &[String]) -> Result>, EmbedderError> { Err(EmbedderError::Unsupported(format!( "embedder strategy `{}` is not implemented yet (follow-up: real ONNX/HTTP backend)", self.strategy ))) } fn dimension(&self) -> usize { self.dimension } } /// A deterministic, dependency-free [`Embedder`] backed by a hashing /// bag-of-words. Pure std (FNV-1a hashing of whitespace tokens into `dimension` /// buckets, then L2-normalised). Stable across runs and platforms. /// /// It is **not** a stub: it lets the étage-2 pipeline run end-to-end with no heavy /// dependency. Vector quality is crude (lexical overlap only), but the /// [`Embedder`] contract is fully honoured, so it is a valid offline embedder and /// the substrate for deterministic tests of [`crate::store::VectorMemoryRecall`]. #[derive(Clone)] pub struct HashEmbedder { id: String, dimension: usize, } impl HashEmbedder { /// Builds a hashing embedder producing `dimension`-length vectors. /// /// # Panics /// Panics if `dimension` is `0` (an embedder must produce fixed-length /// non-empty vectors). #[must_use] pub fn new(id: impl Into, dimension: usize) -> Self { assert!(dimension > 0, "HashEmbedder dimension must be non-zero"); Self { id: id.into(), dimension, } } /// Embeds one text into an L2-normalised `dimension`-length vector. fn embed_one(&self, text: &str) -> Vec { let mut v = vec![0f32; self.dimension]; for token in text.split_whitespace() { let lower = token.to_ascii_lowercase(); let bucket = (fnv1a(lower.as_bytes()) as usize) % self.dimension; v[bucket] += 1.0; } let norm = v.iter().map(|x| x * x).sum::().sqrt(); if norm > 0.0 { for x in &mut v { *x /= norm; } } v } } #[async_trait] impl Embedder for HashEmbedder { fn id(&self) -> &str { &self.id } async fn embed(&self, texts: &[String]) -> Result>, EmbedderError> { Ok(texts.iter().map(|t| self.embed_one(t)).collect()) } fn dimension(&self) -> usize { self.dimension } } /// FNV-1a 64-bit hash of a byte slice (deterministic, dependency-free). fn fnv1a(bytes: &[u8]) -> u64 { const OFFSET: u64 = 0xcbf2_9ce4_8422_2325; const PRIME: u64 = 0x0000_0100_0000_01b3; let mut hash = OFFSET; for &b in bytes { hash ^= u64::from(b); hash = hash.wrapping_mul(PRIME); } hash } // --------------------------------------------------------------------------- // HttpEmbedder — real `localServer` / `api` engine (LOT C1a, feature `vector-http`). // --------------------------------------------------------------------------- /// Default Ollama (and llama.cpp `--api`) OpenAI-compatible embeddings endpoint, /// used when a `localServer` profile leaves [`EmbedderProfile::endpoint`] empty. #[cfg(feature = "vector-http")] pub const DEFAULT_LOCAL_EMBED_ENDPOINT: &str = "http://localhost:11434/v1/embeddings"; /// A real [`Embedder`] talking to an **OpenAI-compatible** `/v1/embeddings` /// endpoint over HTTP — the shape served by Ollama, llama.cpp's server, HuggingFace /// text-embeddings-inference, and the OpenAI/Voyage/Together APIs alike. One adapter /// covers both [`EmbedderStrategy::LocalServer`] (no auth) and /// [`EmbedderStrategy::Api`] (a `Bearer` token read from an env var — never the key /// itself in config). /// /// **Best-effort by contract** (see [`Embedder`]): an unreachable host, a non-2xx /// status, a missing API key, or a malformed body all map to a clean /// [`EmbedderError`] (never a panic), so a composing [`crate::store::AdaptiveMemoryRecall`] /// degrades to the naïve recall. Construction is cheap and infallible; nothing /// hits the network until [`embed`](Embedder::embed) is called (so the /// zero-download posture holds until the embedder is actually used). #[cfg(feature = "vector-http")] pub struct HttpEmbedder { id: String, dimension: usize, endpoint: String, model: Option, /// `Some(var)` for the `api` strategy: the env var carrying the bearer token. api_key_env: Option, client: reqwest::Client, } #[cfg(feature = "vector-http")] #[derive(serde::Serialize)] struct EmbeddingsRequest<'a> { #[serde(skip_serializing_if = "Option::is_none")] model: Option<&'a str>, input: &'a [String], } #[cfg(feature = "vector-http")] #[derive(serde::Deserialize)] struct EmbeddingsResponse { data: Vec, } #[cfg(feature = "vector-http")] #[derive(serde::Deserialize)] struct EmbeddingDatum { embedding: Vec, /// Position in the input batch; the OpenAI shape guarantees it, and we sort by /// it to restore input order regardless of how the server returns the array. #[serde(default)] index: usize, } #[cfg(feature = "vector-http")] impl HttpEmbedder { /// Builds the embedder from a declarative profile. `localServer` profiles with /// no `endpoint` fall back to [`DEFAULT_LOCAL_EMBED_ENDPOINT`]; `api` profiles /// carry their endpoint explicitly. `api_key_env` is retained only for the /// `api` strategy. #[must_use] pub fn from_profile(profile: &EmbedderProfile) -> Self { let is_api = profile.strategy == EmbedderStrategy::Api; let endpoint = profile .endpoint .clone() .filter(|e| !e.is_empty()) .unwrap_or_else(|| DEFAULT_LOCAL_EMBED_ENDPOINT.to_owned()); Self { id: profile.id.clone(), dimension: profile.dimension, endpoint, model: profile.model.clone(), api_key_env: if is_api { profile.api_key_env.clone() } else { None }, // A bounded timeout keeps a hung/slow endpoint from stalling a recall // forever; the composing `AdaptiveMemoryRecall` then degrades to naïve. // Fall back to a default client if the builder cannot be constructed. client: reqwest::Client::builder() .timeout(std::time::Duration::from_secs(30)) .build() .unwrap_or_else(|_| reqwest::Client::new()), } } } #[cfg(feature = "vector-http")] #[async_trait] impl Embedder for HttpEmbedder { fn id(&self) -> &str { &self.id } async fn embed(&self, texts: &[String]) -> Result>, EmbedderError> { if texts.is_empty() { return Ok(Vec::new()); } let mut request = self.client.post(&self.endpoint).json(&EmbeddingsRequest { model: self.model.as_deref(), input: texts, }); // `api` strategy: attach the bearer token read from the configured env var. // A configured-but-unset var is an "unavailable" condition, not a panic. if let Some(var) = &self.api_key_env { match std::env::var(var) { Ok(key) if !key.is_empty() => { request = request.bearer_auth(key); } _ => { return Err(EmbedderError::Unavailable(format!( "API key env var `{var}` is not set" ))); } } } let response = request.send().await.map_err(|e| { EmbedderError::Unavailable(format!("request to `{}` failed: {e}", self.endpoint)) })?; if !response.status().is_success() { return Err(EmbedderError::Unavailable(format!( "embeddings endpoint `{}` returned status {}", self.endpoint, response.status() ))); } let body: EmbeddingsResponse = response .json() .await .map_err(|e| EmbedderError::Io(format!("malformed embeddings response: {e}")))?; if body.data.len() != texts.len() { return Err(EmbedderError::Io(format!( "embeddings count mismatch: expected {}, got {}", texts.len(), body.data.len() ))); } // Restore input order, then validate each vector's dimension. let mut data = body.data; data.sort_by_key(|d| d.index); let mut vectors = Vec::with_capacity(data.len()); for datum in data { if datum.embedding.len() != self.dimension { return Err(EmbedderError::Io(format!( "embedding dimension mismatch: profile declares {}, server returned {}", self.dimension, datum.embedding.len() ))); } vectors.push(datum.embedding); } Ok(vectors) } fn dimension(&self) -> usize { self.dimension } } /// Best-effort probe of whether an Ollama-style local embedding server is reachable /// at `base_url` (e.g. `http://localhost:11434`). Used by the contextual /// "configure an embedder?" prompt (C3) to detect an already-installed local engine /// before suggesting any download — the Linux-spirit "detect the existing first" /// rule. Never errors: a failure to reach the host simply means "not detected". #[cfg(feature = "vector-http")] pub async fn detect_ollama(base_url: &str) -> bool { let base = base_url.trim_end_matches('/'); let url = format!("{base}/api/tags"); let client = match reqwest::Client::builder() .timeout(std::time::Duration::from_millis(800)) .build() { Ok(c) => c, Err(_) => return false, }; matches!(client.get(&url).send().await, Ok(r) if r.status().is_success()) } // --------------------------------------------------------------------------- // OnnxEmbedder — real in-process `localOnnx` engine (LOT C1b, feature `vector-onnx`). // --------------------------------------------------------------------------- // // The data/inspection helpers below (`OnnxModelInfo`, `RECOMMENDED_ONNX_MODELS`, // `ONNX_CACHE_SUBDIR`, `onnx_model_is_cached`) are **always** compiled — they are // pure data and pure-FS inspection, with no dependency on `fastembed`, so the C2/C3 // UI can describe and probe the ONNX models even in a build without the feature. /// Cache subdirectory (under the app data dir) where the `localOnnx` engine stores /// its downloaded model files. Forced explicitly so nothing ever lands in the /// global hf-hub cache outside IdeA's data directory. pub const ONNX_CACHE_SUBDIR: &str = "embedders/onnx"; /// Static, dependency-free description of a recommendable local ONNX model, for the /// "configure an embedder?" UI (C2/C3): display name, vector dimension, approximate /// on-disk size, and whether it is the recommended default. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct OnnxModelInfo { /// Stable model id accepted by a `localOnnx` profile's `model` field. pub id: &'static str, /// Human-readable name for the UI. pub display_name: &'static str, /// Length of the vectors this model produces. pub dimension: usize, /// Approximate download/disk size in megabytes (for the UI download hint). pub approx_size_mb: u32, /// Whether this is the recommended default model. pub recommended: bool, } /// The curated list of local ONNX models IdeA can offer to download. Kept tiny on /// purpose (the Linux-spirit "small, multilingual, good enough" default). pub const RECOMMENDED_ONNX_MODELS: &[OnnxModelInfo] = &[OnnxModelInfo { id: "multilingual-e5-small", display_name: "Multilingual E5 Small", dimension: 384, approx_size_mb: 118, recommended: true, }]; /// Best-effort, **pure-FS** probe of whether a `localOnnx` model already lives in /// `cache_dir` (no network, no `fastembed` dependency — available even without the /// `vector-onnx` feature). Heuristic: a non-empty subdirectory whose name contains /// the model token, mirroring hf-hub's `models----` cache layout. /// /// A `false` here only means "not detected"; it never blocks anything. #[must_use] pub fn onnx_model_is_cached(cache_dir: &std::path::Path, model: &str) -> bool { let needle = model.replace(['/', '_'], "-").to_ascii_lowercase(); let Ok(entries) = std::fs::read_dir(cache_dir) else { return false; }; for entry in entries.flatten() { if !entry.file_type().map(|t| t.is_dir()).unwrap_or(false) { continue; } let name = entry.file_name().to_string_lossy().to_ascii_lowercase(); if !name.contains(&needle) { continue; } // Non-empty directory ⇒ treat as a present cached model. if std::fs::read_dir(entry.path()) .map(|mut d| d.next().is_some()) .unwrap_or(false) { return true; } } false } // --------------------------------------------------------------------------- // EmbedderEnvProbe — best-effort local-environment inspector (LOT C2/C3). // --------------------------------------------------------------------------- /// Default base URL of a local Ollama-style embedding server, probed by /// [`EmbedderEnvProbe`] to detect an already-installed local engine (the /// Linux-spirit "detect the existing first" rule). The base (no path): the probe /// appends `/api/tags` via [`detect_ollama`]. pub const DEFAULT_OLLAMA_BASE_URL: &str = "http://localhost:11434"; /// Concrete [`EmbedderEnvInspector`]: a **best-effort, never-failing** probe of the /// local embedding environment. /// /// - `onnx_cached_models`: for each model in [`RECOMMENDED_ONNX_MODELS`], reports /// the ids already present in `onnx_cache_dir` (pure FS, no `fastembed` /// dependency). The blocking `read_dir` runs on a `spawn_blocking` thread so it /// never stalls the async runtime. /// - `ollama_detected`: `true` only when built with the `vector-http` capability /// **and** an Ollama-style server answers at `ollama_base_url`; `false` otherwise. /// /// Construction is cheap and infallible; nothing touches disk or the network until /// [`inspect`](EmbedderEnvInspector::inspect) is called. #[derive(Debug, Clone)] pub struct EmbedderEnvProbe { onnx_cache_dir: PathBuf, /// Only read under `vector-http` (by [`detect_ollama`]); retained unconditionally /// so the constructor signature is stable across build configs. #[cfg_attr(not(feature = "vector-http"), allow(dead_code))] ollama_base_url: String, } impl EmbedderEnvProbe { /// Builds the probe from the ONNX cache directory (`/embedders/onnx`, /// see [`ONNX_CACHE_SUBDIR`]) and the base URL of the local Ollama-style server. #[must_use] pub fn new(onnx_cache_dir: PathBuf, ollama_base_url: impl Into) -> Self { Self { onnx_cache_dir, ollama_base_url: ollama_base_url.into(), } } } #[async_trait] impl EmbedderEnvInspector for EmbedderEnvProbe { async fn inspect(&self) -> EmbedderEnvReport { // Pure-FS cache scan off the runtime: `onnx_model_is_cached` calls blocking // `read_dir`, so run the whole scan on a blocking thread. let cache_dir = self.onnx_cache_dir.clone(); let onnx_cached_models = tokio::task::spawn_blocking(move || { RECOMMENDED_ONNX_MODELS .iter() .filter(|m| onnx_model_is_cached(&cache_dir, m.id)) .map(|m| m.id.to_owned()) .collect::>() }) .await .unwrap_or_default(); #[cfg(feature = "vector-http")] let ollama_detected = detect_ollama(&self.ollama_base_url).await; #[cfg(not(feature = "vector-http"))] let ollama_detected = false; EmbedderEnvReport { ollama_detected, onnx_cached_models, } } } // --------------------------------------------------------------------------- // FsEmbedderPromptStore — per-project embedder-suggestion state (LOT C3). // --------------------------------------------------------------------------- /// `.ideai/` directory name inside a project root. const PROMPT_IDEAI_DIR: &str = ".ideai"; /// Memory sub-dir (the suggestion state lives beside the memory it concerns). const PROMPT_MEMORY_DIR: &str = "memory"; /// State file name (dot-prefixed: derived/local state, like the vector index). const PROMPT_FILE: &str = ".embedder-prompt.json"; /// Wire form of the suggestion-dismissal choice (camelCase). #[derive(Debug, Clone, Copy, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] enum DismissalWire { Later, Never, } impl From for DismissalWire { fn from(d: EmbedderPromptDismissal) -> Self { match d { EmbedderPromptDismissal::Later => Self::Later, EmbedderPromptDismissal::Never => Self::Never, } } } impl From for EmbedderPromptDismissal { fn from(w: DismissalWire) -> Self { match w { DismissalWire::Later => Self::Later, DismissalWire::Never => Self::Never, } } } /// On-disk shape of `.ideai/memory/.embedder-prompt.json`: /// `{ "dismissed": "later" | "never" }`. #[derive(Debug, Clone, Copy, Serialize, Deserialize)] struct PromptDoc { dismissed: DismissalWire, } /// File-backed [`EmbedderPromptStore`] (LOT C3): persists the per-project response /// to the embedder suggestion under `.ideai/memory/.embedder-prompt.json`, through /// the [`FileSystem`] port (so it is location-neutral, SSH/WSL unchanged). /// /// Reads are **best-effort**: a missing file ⇒ `Ok(None)` (never answered); a /// malformed file also degrades to `Ok(None)` rather than surfacing an error, so a /// hand-corrupted state never blocks the suggestion logic. #[derive(Clone)] pub struct FsEmbedderPromptStore { fs: Arc, } impl FsEmbedderPromptStore { /// Builds the store from the [`FileSystem`] port. #[must_use] pub fn new(fs: Arc) -> Self { Self { fs } } /// `/.ideai/memory`. fn memory_dir(root: &ProjectPath) -> String { let base = root.as_str().trim_end_matches(['/', '\\']); format!("{base}/{PROMPT_IDEAI_DIR}/{PROMPT_MEMORY_DIR}") } /// `/.embedder-prompt.json`. fn prompt_path(root: &ProjectPath) -> RemotePath { RemotePath::new(format!("{}/{PROMPT_FILE}", Self::memory_dir(root))) } } #[async_trait] impl EmbedderPromptStore for FsEmbedderPromptStore { async fn read( &self, root: &ProjectPath, ) -> Result, StoreError> { match self.fs.read(&Self::prompt_path(root)).await { Ok(bytes) => Ok(serde_json::from_slice::(&bytes) .ok() .map(|doc| doc.dismissed.into())), // Never answered (or unreadable) ⇒ no recorded dismissal. Err(FsError::NotFound(_)) | Err(_) => Ok(None), } } async fn write( &self, root: &ProjectPath, dismissal: EmbedderPromptDismissal, ) -> Result<(), StoreError> { let dir = RemotePath::new(Self::memory_dir(root)); self.fs .create_dir_all(&dir) .await .map_err(|e| StoreError::Io(e.to_string()))?; let doc = PromptDoc { dismissed: dismissal.into(), }; let bytes = serde_json::to_vec_pretty(&doc) .map_err(|e| StoreError::Serialization(e.to_string()))?; self.fs .write(&Self::prompt_path(root), &bytes) .await .map_err(|e| StoreError::Io(e.to_string())) } } #[cfg(feature = "vector-onnx")] mod onnx { use std::path::{Path, PathBuf}; use std::sync::Arc; use async_trait::async_trait; use fastembed::{EmbeddingModel, InitOptions, TextEmbedding}; use tokio::sync::{Mutex, OnceCell}; use domain::ports::{Embedder, EmbedderError}; use domain::profile::EmbedderProfile; /// Resolves a profile `model` string to a concrete fastembed model + its /// dimension. `None` model ⇒ the recommended default (Multilingual E5 Small, /// 384). An unknown, non-empty string ⇒ `None` (caller maps to `Unsupported`). pub(super) fn resolve_onnx_model(model: Option<&str>) -> Option<(EmbeddingModel, usize)> { match model { None => Some((EmbeddingModel::MultilingualE5Small, 384)), Some(m) => match m.trim().to_ascii_lowercase().as_str() { "multilingual-e5-small" | "e5-small" => { Some((EmbeddingModel::MultilingualE5Small, 384)) } _ => None, }, } } /// A real in-process [`Embedder`] running a quantised ONNX model via `fastembed` /// (`localOnnx` strategy). The model is loaded **lazily** on first /// [`embed`](Embedder::embed) (downloaded once into `cache_dir` if missing), then /// reused; all CPU-bound work (model load + inference) runs on a blocking thread. /// /// **Best-effort by contract** (see [`Embedder`]): a failed download/load, /// an unknown model, or a profile/model dimension mismatch all map to a clean /// [`EmbedderError`] (never a panic), so a composing /// [`crate::store::AdaptiveMemoryRecall`] degrades to naïve. Construction is /// cheap and infallible — nothing touches disk or the network until `embed`. pub struct OnnxEmbedder { id: String, dimension: usize, /// `Some(model)` when the profile's model string resolved to a known model; /// `None` for an unknown string ⇒ `embed` returns [`EmbedderError::Unsupported`]. model: Option, /// The original (unknown) model string, for the `Unsupported` message. requested_model: Option, cache_dir: PathBuf, cell: OnceCell>>, } impl OnnxEmbedder { /// Builds the embedder from a declarative profile. **Cheap and infallible**: /// no I/O, no download. An unknown model is detected lazily at `embed` time /// (mapped to [`EmbedderError::Unsupported`]); here it is stored as-is and the /// fallback model is recorded so the struct stays valid. #[must_use] pub fn from_profile(profile: &EmbedderProfile, cache_dir: &Path) -> Self { // Record the resolved model when known; an unknown string is kept so // construction never fails — `embed` returns `Unsupported` for it. let model = resolve_onnx_model(profile.model.as_deref()).map(|(m, _)| m); Self { id: profile.id.clone(), dimension: profile.dimension, model, requested_model: profile.model.clone(), cache_dir: cache_dir.to_path_buf(), cell: OnceCell::new(), } } /// Lazily loads (and caches) the `TextEmbedding`, returning the shared /// `Mutex`-guarded handle. Heavy work runs on a blocking thread. async fn engine(&self) -> Result>, EmbedderError> { let Some(model) = self.model.clone() else { return Err(EmbedderError::Unsupported(format!( "unknown ONNX model `{}` (supported: multilingual-e5-small)", self.requested_model.as_deref().unwrap_or("") ))); }; self.cell .get_or_try_init(|| async { let cache_dir = self.cache_dir.clone(); let built = tokio::task::spawn_blocking(move || { TextEmbedding::try_new(InitOptions::new(model).with_cache_dir(cache_dir)) }) .await .map_err(|e| EmbedderError::Io(format!("onnx init task failed: {e}")))? .map_err(|e| { EmbedderError::Unavailable(format!( "failed to load/download ONNX model: {e}" )) })?; Ok(Arc::new(Mutex::new(built))) }) .await .cloned() } } #[async_trait] impl Embedder for OnnxEmbedder { fn id(&self) -> &str { &self.id } async fn embed(&self, texts: &[String]) -> Result>, EmbedderError> { // Short-circuit BEFORE any model load: an empty batch never downloads. if texts.is_empty() { return Ok(Vec::new()); } // `engine()` surfaces `Unsupported` for an unknown model and // `Unavailable` for a failed load/download — never a panic. let expected = self.dimension; let engine = self.engine().await?; let texts_owned = texts.to_vec(); let vectors = tokio::task::spawn_blocking(move || { let mut guard = engine.blocking_lock(); guard.embed(texts_owned, None) }) .await .map_err(|e| EmbedderError::Io(format!("onnx inference task failed: {e}")))? .map_err(|e| EmbedderError::Io(format!("onnx inference failed: {e}")))?; for v in &vectors { if v.len() != expected { return Err(EmbedderError::Unavailable(format!( "embedding dimension mismatch: profile declares {expected}, model produces {}", v.len() ))); } } Ok(vectors) } fn dimension(&self) -> usize { self.dimension } } } #[cfg(feature = "vector-onnx")] pub use onnx::OnnxEmbedder;