blomios/IdeA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(memory): embedders vectoriels réels HTTP + ONNX derrière features (LOT C1)

Browse Source 
Remplace les StubEmbedder pour les stratégies localServer/api/localOnnx par de
vrais moteurs, chacun derrière une feature cargo off-by-default — la posture
fondatrice « rien d'imposé, zéro dépendance » (défaut none → rappel naïf) reste
byte-for-byte inchangée.

C1a (feature vector-http, reqwest rustls optional):
- HttpEmbedder couvrant localServer (Ollama/llama.cpp) et api (OpenAI/Voyage…),
  payload OpenAI-compatible /v1/embeddings, ordre restauré par index, bearer
  token lu via env var (jamais en clair), timeout client 30s.
- detect_ollama() pour la détection de l'existant (C3).

C1b (feature vector-onnx, fastembed v5 optional):
- OnnxEmbedder en-process (e5-small, dim 384), init paresseuse + spawn_blocking,
  cache modèle sous <app_data>/embedders/onnx — aucun download au build ni au
  first-run, uniquement à la demande au 1er embed.
- Catalogue RECOMMENDED_ONNX_MODELS + ONNX_CACHE_SUBDIR + onnx_model_is_cached
  exposés (sans feature) pour la config (C2) et la popup (C3).

embedder_from_profile(profile, onnx_cache_dir) dispatche feature-gated ; sans la
feature, retombe sur StubEmbedder (Unsupported) → fallback naïf via
AdaptiveMemoryRecall. Composition root (build_memory_recall) propage le cache dir.

Tests: 10 HTTP + 6 ONNX (dont 2 #[ignore] download réel) + 26 vectoriels, verts
en défaut, --features vector-http et --features vector-onnx.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Blomios
2026-06-08 19:00:40 +02:00
parent b39c11a64d
commit 32398827fb
10 changed files with 2012 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

475
crates/infrastructure/src/store/embedder.rs
View File

@ -31,28 +31,57 @@ use domain::profile::{EmbedderProfile, EmbedderStrategy};
/// Builds the concrete [`Embedder`] for a profile, or `None` when the strategy is
/// [`EmbedderStrategy::None`] (recall stays naïve).
///
/// The concrete engines are stubs ([`StubEmbedder`]) until their real ONNX/HTTP
/// integration lands; they fail cleanly with [`EmbedderError::Unsupported`] rather
/// than panic, so a composing recall degrades to 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) -> Option<Box<dyn Embedder>> {
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 => Some(Box::new(StubEmbedder::new(
profile.id.clone(),
profile.dimension,
"localOnnx",
))),
EmbedderStrategy::LocalServer => Some(Box::new(StubEmbedder::new(
profile.id.clone(),
profile.dimension,
"localServer",
))),
EmbedderStrategy::Api => Some(Box::new(StubEmbedder::new(
profile.id.clone(),
profile.dimension,
"api",
))),
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,
)))
}
}
}
}
@ -171,3 +200,411 @@ fn fnv1a(bytes: &[u8]) -> u64 {
}
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
}
#[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;