- LOT C2 (§14.5.3) : use cases de configuration des embedders déclaratifs (List/Save/Delete + DescribeEmbedderEngines : modèles ONNX recommandés, environnement local détecté, stratégies compilées). UI EmbedderSettings. - LOT C3 (§14.5.5) : suggestion contextuelle best-effort à l'activation quand la mémoire dépasse le budget de recall sans embedder configuré (event EmbedderSuggested, anti-spam 1×/session, « ne plus demander »). - Contexte projet partagé .ideai/CONTEXT.md (model-agnostic) injecté à tous les agents/profils au lancement, avant la persona. UI ProjectContextPanel. Tests : backend workspace vert (0 échec) ; frontend 306/306. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
811 lines
31 KiB
Rust
811 lines
31 KiB
Rust
//! [`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 (`<app_data_dir>/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<Box<dyn Embedder>> {
|
|
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<String>, 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<Vec<Vec<f32>>, 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<String>, 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<f32> {
|
|
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::<f32>().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<Vec<Vec<f32>>, 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<String>,
|
|
/// `Some(var)` for the `api` strategy: the env var carrying the bearer token.
|
|
api_key_env: Option<String>,
|
|
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<EmbeddingDatum>,
|
|
}
|
|
|
|
#[cfg(feature = "vector-http")]
|
|
#[derive(serde::Deserialize)]
|
|
struct EmbeddingDatum {
|
|
embedding: Vec<f32>,
|
|
/// 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<Vec<Vec<f32>>, 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--<org>--<name>` 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 (`<app_data_dir>/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<String>) -> 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::<Vec<String>>()
|
|
})
|
|
.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<EmbedderPromptDismissal> for DismissalWire {
|
|
fn from(d: EmbedderPromptDismissal) -> Self {
|
|
match d {
|
|
EmbedderPromptDismissal::Later => Self::Later,
|
|
EmbedderPromptDismissal::Never => Self::Never,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl From<DismissalWire> 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<dyn FileSystem>,
|
|
}
|
|
|
|
impl FsEmbedderPromptStore {
|
|
/// Builds the store from the [`FileSystem`] port.
|
|
#[must_use]
|
|
pub fn new(fs: Arc<dyn FileSystem>) -> Self {
|
|
Self { fs }
|
|
}
|
|
|
|
/// `<root>/.ideai/memory`.
|
|
fn memory_dir(root: &ProjectPath) -> String {
|
|
let base = root.as_str().trim_end_matches(['/', '\\']);
|
|
format!("{base}/{PROMPT_IDEAI_DIR}/{PROMPT_MEMORY_DIR}")
|
|
}
|
|
|
|
/// `<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<Option<EmbedderPromptDismissal>, StoreError> {
|
|
match self.fs.read(&Self::prompt_path(root)).await {
|
|
Ok(bytes) => Ok(serde_json::from_slice::<PromptDoc>(&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<EmbeddingModel>,
|
|
/// The original (unknown) model string, for the `Unsupported` message.
|
|
requested_model: Option<String>,
|
|
cache_dir: PathBuf,
|
|
cell: OnceCell<Arc<Mutex<TextEmbedding>>>,
|
|
}
|
|
|
|
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<Arc<Mutex<TextEmbedding>>, 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<Vec<Vec<f32>>, 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;
|