diff --git a/CHANGELOG.md b/CHANGELOG.md index b1059ad9..163c698d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,8 @@ ### Fixed +- **The embedding cache no longer serves the vector of another model (#1536)** - the persistent cache (`embed_cache.db`) was keyed by the embedder's constant name (`openai`, `ollama`, `embeddinggemma-q4`), so after changing `[embedding] model`, `dims` or the endpoint, an already cached text came back with the old model's vector, and a vector produced by the cloud fallback was cached under the primary's name. The key is now an identity the embedder owns (backend, model, dimensions, and for `openai` and `ollama` a hash of the endpoint; the API key is not part of it), a cached vector whose length is not the embedder's dimensions is ignored, an empty vector is not cached, and the fallback wraps the cached primary so cloud vectors never enter the cache (with a fallback configured, a repeated text whose primary embedding fails now calls the cloud each time instead of being served from the cache). Rows written under the old key are deleted once when the cache is first opened (`PRAGMA user_version` 1); the first embeddings after the upgrade are recomputed. `Embedder` gains a provided method `identity()` (default `name:dims`), so existing implementations keep compiling. + - **`UTEKE_HOME` pointing at another store no longer routes to a running server (#1445)** - with `[server] enabled = true` and `uteke-serve` running, `UTEKE_HOME=/tmp/iso uteke remember ...` was sent to the server and landed in the server's store instead of `/tmp/iso`. The server now reports which store it serves on the new `GET /store` (LAB route, token required when auth is on): `{"id": "<16 hex digits>"}`, a SHA-256 prefix of the store's canonical path, never the path itself. Before routing to a loopback server the CLI compares that id with the store it resolved and runs locally on a mismatch (with a warning on stderr). The probe sends `UTEKE_AUTH_TOKEN` like any routed command; without it an auth-enabled server refuses the command with 401, so nothing lands in the wrong store. A server without the route (404), an unreachable one or a non-loopback one keeps the previous routing. `docs/core-contract.json` gains `/store` in the `lab` list only. ### Changed diff --git a/crates/uteke-core/src/embed/cache.rs b/crates/uteke-core/src/embed/cache.rs index 5f82f11d..7563cb74 100644 --- a/crates/uteke-core/src/embed/cache.rs +++ b/crates/uteke-core/src/embed/cache.rs @@ -49,6 +49,24 @@ impl EmbeddingCache { ) .map_err(|e| Error::db("create embedding cache schema", e))?; + // The cache used to be keyed by the embedder's constant name + // ("openai", "embeddinggemma-q4"), which could serve the vector of + // another model (#1536). Rows written that way can never be hit again + // and may hold vectors of the wrong model: purge them once. A key now + // is an identity (`name:...:dims`), so an old name has no colon. + let version: i64 = conn + .query_row("PRAGMA user_version", [], |row| row.get(0)) + .unwrap_or(0); + if version < 1 { + conn.execute( + "DELETE FROM embedding_cache WHERE model_name NOT LIKE '%:%'", + [], + ) + .map_err(|e| Error::db("purge embedding cache rows of the old key", e))?; + conn.execute_batch("PRAGMA user_version = 1;") + .map_err(|e| Error::db("mark the embedding cache as migrated", e))?; + } + Ok(Self { conn: Mutex::new(conn), }) @@ -67,8 +85,9 @@ impl EmbeddingCache { .collect() } - /// Look up a cached embedding. Returns `None` on miss. - fn lookup(&self, model_name: &str, text: &str) -> Option> { + /// Look up a cached embedding. Returns `None` on miss, and for a stored + /// vector that does not have `dims` elements (it came from something else). + fn lookup(&self, model_name: &str, text: &str, dims: usize) -> Option> { let key = Self::cache_key(model_name, text); let conn = self.conn.lock().ok()?; @@ -80,11 +99,15 @@ impl EmbeddingCache { ) .ok()?; - deserialize_embedding(&blob) + deserialize_embedding(&blob).filter(|v| v.len() == dims) } /// Store an embedding in the cache. fn store(&self, model_name: &str, text: &str, embedding: &[f32]) { + // An empty vector is "no embedding", not something to remember. + if embedding.is_empty() { + return; + } let key = Self::cache_key(model_name, text); let blob = serialize_embedding(embedding); @@ -97,16 +120,19 @@ impl EmbeddingCache { } } - /// Clear all cached embeddings for a specific model. - /// Called when the model is upgraded or changed. + /// Clear all cached embeddings of a backend or model: the rows whose + /// identity is `model_name` or starts with `model_name:` (an identity is + /// `name:...:dims`). Called when the model is upgraded or changed. #[allow(dead_code)] // Public API for future cache invalidation pub fn invalidate_model(&self, model_name: &str) -> Result<(), Error> { let conn = self .conn .lock() .map_err(|_| Error::lock("embedding cache during invalidation"))?; + // `substr` compare, not LIKE: a name may contain `%` or `_`. conn.execute( - "DELETE FROM embedding_cache WHERE model_name = ?1", + "DELETE FROM embedding_cache + WHERE model_name = ?1 OR substr(model_name, 1, length(?1) + 1) = ?1 || ':'", rusqlite::params![model_name], ) .map_err(|e| Error::db("invalidate embedding cache", e))?; @@ -132,10 +158,12 @@ impl CachingEmbedder { impl Embedder for CachingEmbedder { fn embed(&self, text: &str) -> Result, Error> { - let model_name = self.inner.name(); + // Keyed by what produced the vectors, not by the backend's name (#1536). + let model_name = self.inner.identity(); + let model_name = model_name.as_str(); // Try cache first - if let Some(cached) = self.cache.lookup(model_name, text) { + if let Some(cached) = self.cache.lookup(model_name, text, self.inner.dims()) { tracing::debug!(text_len = text.len(), "embedding cache hit"); return Ok(cached); } @@ -161,6 +189,10 @@ impl Embedder for CachingEmbedder { fn name(&self) -> &str { self.inner.name() } + + fn identity(&self) -> String { + self.inner.identity() + } } // ── Serialization ────────────────────────────────────────────────────────── @@ -220,12 +252,11 @@ mod tests { } fn tmp_cache_path() -> std::path::PathBuf { - let dir = std::env::temp_dir(); - let nanos = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos(); - dir.join(format!("uteke_test_cache_{nanos}.db")) + // Unique per call: tests run in parallel, and two calls in the same + // nanosecond used to share (and delete) one file. + static NEXT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + let n = NEXT.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + std::env::temp_dir().join(format!("uteke_test_cache_{}_{n}.db", std::process::id())) } #[test] @@ -327,14 +358,9 @@ mod tests { .unwrap(); assert_eq!(count, 2); - // Invalidate - let conn_invalidate = Connection::open(&cache_path).unwrap(); - conn_invalidate - .execute( - "DELETE FROM embedding_cache WHERE model_name = ?1", - rusqlite::params!["mock-model"], - ) - .unwrap(); + // Invalidate through the real function: it removes every identity of + // the model (`mock-model:2`), and nothing of another one. + cache.cache.invalidate_model("mock-model").unwrap(); let count_after: i64 = conn_check .query_row("SELECT COUNT(*) FROM embedding_cache", [], |row| row.get(0)) @@ -384,4 +410,269 @@ mod tests { let _ = std::fs::remove_file(&cache_path); } + + // ── Identity (#1536) ───────────────────────────────────────────────── + + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// An embedder whose identity, size and value the test chooses, counting calls. + struct Fake { + identity: String, + dims: usize, + value: f32, + calls: Arc, + fail: bool, + } + + impl Fake { + fn new(identity: &str, dims: usize, value: f32) -> (Self, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + ( + Self { + identity: identity.to_string(), + dims, + value, + calls: calls.clone(), + fail: false, + }, + calls, + ) + } + } + + impl Embedder for Fake { + fn embed(&self, _text: &str) -> Result, Error> { + self.calls.fetch_add(1, Ordering::SeqCst); + if self.fail { + return Err(Error::embed_msg("fake failure")); + } + Ok(vec![self.value; self.dims]) + } + fn dims(&self) -> usize { + self.dims + } + fn max_seq_len(&self) -> usize { + 128 + } + fn name(&self) -> &str { + "fake" + } + fn identity(&self) -> String { + self.identity.clone() + } + } + + /// The same cache file used by two embedders in turn, like two runs of the CLI. + fn run_with(path: &std::path::Path, fake: Fake, text: &str) -> Vec { + let cache = EmbeddingCache::open(path).unwrap(); + CachingEmbedder::with_cache(Box::new(fake), cache) + .embed(text) + .unwrap() + } + + #[test] + fn the_default_identity_is_the_name_and_the_dimensions() { + let c = CountingEmbedder { + dims: 4, + call_count: Mutex::new(0), + }; + assert_eq!(c.identity(), "mock-model:4"); + } + + #[test] + fn a_different_model_size_or_endpoint_never_gets_the_vector_of_another() { + let path = tmp_cache_path(); + let (a, a_calls) = Fake::new("openai:model-a:8:aaaa", 8, 1.0); + assert_eq!(run_with(&path, a, "same text"), vec![1.0; 8]); + assert_eq!(a_calls.load(Ordering::SeqCst), 1); + + // Another model, another size, another endpoint: each must be computed. + for (identity, dims, value) in [ + ("openai:model-b:8:aaaa", 8, 2.0), + ("openai:model-a:4:aaaa", 4, 3.0), + ("openai:model-a:8:bbbb", 8, 4.0), + ] { + let (f, calls) = Fake::new(identity, dims, value); + assert_eq!( + run_with(&path, f, "same text"), + vec![value; dims], + "{identity}" + ); + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "{identity} must reach its own backend" + ); + } + + // The same identity is a hit. + let (again, again_calls) = Fake::new("openai:model-a:8:aaaa", 8, 99.0); + assert_eq!(run_with(&path, again, "same text"), vec![1.0; 8]); + assert_eq!(again_calls.load(Ordering::SeqCst), 0); + let _ = std::fs::remove_file(&path); + } + + #[test] + fn a_stored_vector_of_the_wrong_length_is_ignored() { + let path = tmp_cache_path(); + let (f, _) = Fake::new("x:8", 8, 1.0); + run_with(&path, f, "text"); + // Corrupt the stored vector: four elements where the embedder makes eight. + { + let conn = Connection::open(&path).unwrap(); + conn.execute( + "UPDATE embedding_cache SET embedding = ?1", + rusqlite::params![serialize_embedding(&[9.0; 4])], + ) + .unwrap(); + } + let (g, calls) = Fake::new("x:8", 8, 5.0); + assert_eq!(run_with(&path, g, "text"), vec![5.0; 8]); + assert_eq!(calls.load(Ordering::SeqCst), 1, "the bad row was a miss"); + let _ = std::fs::remove_file(&path); + } + + #[test] + fn an_empty_vector_is_not_remembered() { + let path = tmp_cache_path(); + let (f, _) = Fake::new("x:0", 0, 0.0); + run_with(&path, f, "text"); + let conn = Connection::open(&path).unwrap(); + let n: i64 = conn + .query_row("SELECT COUNT(*) FROM embedding_cache", [], |r| r.get(0)) + .unwrap(); + assert_eq!(n, 0); + let _ = std::fs::remove_file(&path); + } + + #[test] + fn rows_written_under_the_old_key_are_purged_once() { + let path = tmp_cache_path(); + // A cache as the previous version left it: keyed by a constant name, + // user_version 0. + { + let conn = Connection::open(&path).unwrap(); + conn.execute_batch( + "CREATE TABLE embedding_cache ( + text_hash TEXT PRIMARY KEY, model_name TEXT NOT NULL, + embedding BLOB NOT NULL, created_at TEXT NOT NULL DEFAULT '');", + ) + .unwrap(); + for (hash, name) in [ + ("h1", "openai"), + ("h2", "embeddinggemma-q4"), + ("h3", "ollama"), + ] { + conn.execute( + "INSERT INTO embedding_cache (text_hash, model_name, embedding) VALUES (?1, ?2, x'00000000')", + rusqlite::params![hash, name], + ) + .unwrap(); + } + } + let count = |path: &std::path::Path| -> i64 { + Connection::open(path) + .unwrap() + .query_row("SELECT COUNT(*) FROM embedding_cache", [], |r| r.get(0)) + .unwrap() + }; + assert_eq!(count(&path), 3); + drop(EmbeddingCache::open(&path).unwrap()); + assert_eq!( + count(&path), + 0, + "the old rows are gone after the first open" + ); + + // A row of the new kind survives every later open. + let (f, _) = Fake::new("openai:m:8:cccc", 8, 1.0); + run_with(&path, f, "kept"); + drop(EmbeddingCache::open(&path).unwrap()); + assert_eq!(count(&path), 1, "the purge runs once, not at every open"); + let _ = std::fs::remove_file(&path); + } + + #[test] + fn what_the_cloud_fallback_answers_is_never_cached() { + use crate::embed::FallbackEmbedder; + let path = tmp_cache_path(); + let (mut primary, _) = Fake::new("onnx:8", 8, 1.0); + primary.fail = true; + let (cloud, cloud_calls) = Fake::new("openai:cloud:8:dddd", 8, 7.0); + // The wrapping order of the core: the cache around the primary, the + // fallback around the cached primary. + let cached_primary = + CachingEmbedder::with_cache(Box::new(primary), EmbeddingCache::open(&path).unwrap()); + let fb = FallbackEmbedder::new(Box::new(cached_primary), Some(Box::new(cloud))).unwrap(); + assert_eq!(fb.embed("text").unwrap(), vec![7.0; 8]); + assert_eq!(fb.embed("text").unwrap(), vec![7.0; 8]); + assert_eq!( + cloud_calls.load(Ordering::SeqCst), + 2, + "no cache served the cloud vector" + ); + let conn = Connection::open(&path).unwrap(); + let n: i64 = conn + .query_row("SELECT COUNT(*) FROM embedding_cache", [], |r| r.get(0)) + .unwrap(); + assert_eq!(n, 0, "nothing the cloud produced is in the cache"); + let _ = std::fs::remove_file(&path); + } + + #[test] + fn the_remote_embedders_identify_model_dimensions_and_endpoint_but_not_the_key() { + use crate::embed::{OllamaEmbedder, OpenAiEmbedder}; + let o = |model: &str, dims: usize, url: &str, key: &str| { + OpenAiEmbedder::new(key, model, url, "", dims) + .unwrap() + .identity() + }; + let base = o("m", 8, "https://a.example/v1", "key-one"); + assert_eq!( + base, + o("m", 8, "https://a.example/v1", "key-two"), + "the key is not part of it" + ); + assert!( + !base.contains("key-one") && !base.contains("a.example"), + "{base}" + ); + assert_ne!(base, o("m2", 8, "https://a.example/v1", "k")); + assert_ne!(base, o("m", 16, "https://a.example/v1", "k")); + assert_ne!(base, o("m", 8, "https://b.example/v1", "k")); + + let l = |model: &str, dims: usize, url: &str| { + OllamaEmbedder::new(url, model, dims).unwrap().identity() + }; + let ollama = l("nomic", 8, "http://localhost:11434"); + assert_ne!(ollama, l("other", 8, "http://localhost:11434")); + assert_ne!(ollama, l("nomic", 4, "http://localhost:11434")); + assert_ne!(ollama, l("nomic", 8, "http://other:11434")); + assert_ne!(ollama, base, "different backends never collide"); + } + + #[test] + fn invalidating_a_model_leaves_the_others_alone() { + let path = tmp_cache_path(); + for identity in ["mock-model:8", "mock-model:4", "mock-model-2:8", "other:8"] { + let (f, _) = Fake::new(identity, 8, 1.0); + run_with(&path, f, "text"); + } + EmbeddingCache::open(&path) + .unwrap() + .invalidate_model("mock-model") + .unwrap(); + let conn = Connection::open(&path).unwrap(); + let mut left: Vec = conn + .prepare("SELECT model_name FROM embedding_cache ORDER BY model_name") + .unwrap() + .query_map([], |r| r.get(0)) + .unwrap() + .map(|r| r.unwrap()) + .collect(); + left.sort(); + assert_eq!(left, vec!["mock-model-2:8", "other:8"]); + let _ = std::fs::remove_file(&path); + } } diff --git a/crates/uteke-core/src/embed/embed_trait.rs b/crates/uteke-core/src/embed/embed_trait.rs index 392ba5db..4e7b1588 100644 --- a/crates/uteke-core/src/embed/embed_trait.rs +++ b/crates/uteke-core/src/embed/embed_trait.rs @@ -30,4 +30,14 @@ pub trait Embedder: Send + Sync { /// Human-readable name for this backend. fn name(&self) -> &str; + + /// What produced the vectors: the backend, the model, the dimensions and + /// (for a remote backend) a fingerprint of the endpoint. Two embedders + /// with the same identity return interchangeable vectors; the embedding + /// cache is keyed by it, so a changed model, size or endpoint never serves + /// the vector of the old one. The default (`name:dims`) suits a backend + /// with one fixed model. + fn identity(&self) -> String { + format!("{}:{}", self.name(), self.dims()) + } } diff --git a/crates/uteke-core/src/embed/fallback.rs b/crates/uteke-core/src/embed/fallback.rs index 11cc7403..fac4f90b 100644 --- a/crates/uteke-core/src/embed/fallback.rs +++ b/crates/uteke-core/src/embed/fallback.rs @@ -101,6 +101,13 @@ impl Embedder for FallbackEmbedder { None => self.primary.name(), } } + + fn identity(&self) -> String { + match &self.fallback { + Some(cloud) => format!("fallback({}|{})", self.primary.identity(), cloud.identity()), + None => self.primary.identity(), + } + } } #[cfg(test)] diff --git a/crates/uteke-core/src/embed/mod.rs b/crates/uteke-core/src/embed/mod.rs index 1bd54781..8a485bc4 100644 --- a/crates/uteke-core/src/embed/mod.rs +++ b/crates/uteke-core/src/embed/mod.rs @@ -23,6 +23,24 @@ pub use openai::OpenAiEmbedder; #[allow(dead_code)] pub type EmbeddingEngine = OnnxEmbedder; +/// A short, stable fingerprint of the parts that identify an endpoint (for an +/// embedder identity). Hashed so a URL with credentials in it never lands in +/// the cache database in the clear. +pub(crate) fn fingerprint(parts: &[&str]) -> String { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + for p in parts { + hasher.update(p.as_bytes()); + hasher.update(b"\0"); + } + hasher + .finalize() + .iter() + .take(4) + .map(|b| format!("{b:02x}")) + .collect() +} + /// Validate a user-supplied base_url: must be http(s):// and absolute. /// /// Rejects empty / schemeless / malformed URLs early so callers get a diff --git a/crates/uteke-core/src/embed/ollama.rs b/crates/uteke-core/src/embed/ollama.rs index 4096e97c..7703446c 100644 --- a/crates/uteke-core/src/embed/ollama.rs +++ b/crates/uteke-core/src/embed/ollama.rs @@ -105,6 +105,15 @@ impl Embedder for OllamaEmbedder { fn name(&self) -> &str { "ollama" } + + fn identity(&self) -> String { + format!( + "ollama:{}:{}:{}", + self.model, + self.dims, + crate::embed::fingerprint(&[&self.base_url]) + ) + } } #[derive(serde::Deserialize)] diff --git a/crates/uteke-core/src/embed/openai.rs b/crates/uteke-core/src/embed/openai.rs index af1b7716..9125005d 100644 --- a/crates/uteke-core/src/embed/openai.rs +++ b/crates/uteke-core/src/embed/openai.rs @@ -157,6 +157,15 @@ impl Embedder for OpenAiEmbedder { fn name(&self) -> &str { "openai" } + + fn identity(&self) -> String { + format!( + "openai:{}:{}:{}", + self.model, + self.dims, + crate::embed::fingerprint(&[&self.base_url, &self.endpoint_path]) + ) + } } #[derive(serde::Deserialize)] diff --git a/crates/uteke-core/src/lib.rs b/crates/uteke-core/src/lib.rs index ce103c82..993a1ee8 100644 --- a/crates/uteke-core/src/lib.rs +++ b/crates/uteke-core/src/lib.rs @@ -1349,34 +1349,9 @@ impl Uteke { *guard = Some(embedder); - // Wrap with fallback if configured. - // Must happen after dim mismatch check so primary dims are validated - // against the index before we potentially add a fallback. - if self.fallback_settings.is_configured() { - let fb = &self.fallback_settings; - tracing::info!( - "Embedding fallback configured — wrapping primary with cloud backup" - ); - let model = fb.model.clone(); - let base_url = fb.base_url.clone(); - let endpoint_path = fb.endpoint_path.clone(); - let dims = backend_dims; // use validated primary dims - let cloud_embedder = crate::embed::OpenAiEmbedder::new( - &fb.api_key, - &model, - &base_url, - &endpoint_path, - dims, - )?; - let fallback_embedder = crate::embed::FallbackEmbedder::new( - // Take the primary out, wrap it, put back - guard.take().unwrap(), - Some(Box::new(cloud_embedder)), - )?; - *guard = Some(Box::new(fallback_embedder)); - } - - // Wrap with persistent embedding cache (#896). + // Wrap the PRIMARY embedder with the persistent embedding cache (#896). + // It goes inside the fallback wrapper below, so vectors the cloud + // fallback produced are never cached under the primary's identity (#1536). // CachingEmbedder intercepts embed() calls and serves from SQLite // cache on hit, avoiding the ~2s ONNX model load on repeated queries. // Skipped for in-memory stores (tests) where no cache path exists. @@ -1389,7 +1364,9 @@ impl Uteke { cache_path = %cache_path.display(), "Persistent embedding cache enabled" ); - let inner = guard.take().unwrap(); + let inner = guard.take().ok_or_else(|| { + Error::embed_msg("embedder missing while adding the cache") + })?; *guard = Some(Box::new(crate::embed::CachingEmbedder::with_cache( inner, cache, ))); @@ -1403,6 +1380,36 @@ impl Uteke { } } } + + // Wrap with fallback if configured. + // Must happen after dim mismatch check so primary dims are validated + // against the index before we potentially add a fallback. + if self.fallback_settings.is_configured() { + let fb = &self.fallback_settings; + tracing::info!( + "Embedding fallback configured — wrapping primary with cloud backup" + ); + let model = fb.model.clone(); + let base_url = fb.base_url.clone(); + let endpoint_path = fb.endpoint_path.clone(); + let dims = backend_dims; // use validated primary dims + let cloud_embedder = crate::embed::OpenAiEmbedder::new( + &fb.api_key, + &model, + &base_url, + &endpoint_path, + dims, + )?; + let primary = guard.take().ok_or_else(|| { + Error::embed_msg("embedder missing while adding the fallback") + })?; + let fallback_embedder = crate::embed::FallbackEmbedder::new( + // Wrap the primary (cached when a cache is configured) + primary, + Some(Box::new(cloud_embedder)), + )?; + *guard = Some(Box::new(fallback_embedder)); + } } Ok(()) }