Skip to main content

kimetsu_brain/
embeddings.rs

1//! Embeddings + hybrid retrieval scaffolding (v0.4.2).
2//!
3//! The broker is FTS-only through v0.4.1: right answers with wrong
4//! words get missed. v0.4.2 lays the infrastructure for hybrid
5//! retrieval (lexical + semantic) without binding to a specific
6//! embedder. v0.4.3 wires fastembed-rs as the production default.
7//!
8//! Layers introduced here:
9//!   1. The [`Embedder`] trait — anything that can map a text to a
10//!      fixed-dimension float vector plus identify the model that
11//!      produced it.
12//!   2. [`NoopEmbedder`] — production default when no real embedder
13//!      is wired. `embed()` errors with `NotImplemented`; the write
14//!      path treats that as "store NULL" and the retrieval path
15//!      treats it as "skip the cosine blend, FTS only".
16//!   3. [`StubEmbedder`] — deterministic, dependency-free, test-only
17//!      pseudo-embedder. Lets us exercise the hybrid scoring path
18//!      end-to-end without depending on fastembed-rs or downloading
19//!      a model in CI.
20//!   4. [`cosine_similarity`] + BLOB codec helpers so the brain.db
21//!      schema can store embeddings as little-endian `f32` blobs.
22//!
23//! Wire compatibility:
24//!   * Embeddings are nullable. Pre-v0.4.2 rows have NULL embedding
25//!     + NULL embedding_model. The retrieval blender treats them as
26//!     "lexical-only" — they still score via FTS, they just don't
27//!     contribute to the cosine term.
28//!   * The `embedding_model` column carries an opaque string id
29//!     ("bge-small-en-v1.5", "stub-d8", etc.). Queries blend only
30//!     when the query's embedder id matches the row's stored id,
31//!     so mixing models inside one brain.db is safe (rows with a
32//!     different model fall back to lexical-only).
33//!
34//! Scoring (added in v0.4.2):
35//!   `final_relevance = (1 - alpha) * lexical + alpha * cosine`
36//!   where `alpha = brain.broker.hybrid_alpha` (defaulted to 0.5,
37//!   tuned later in v0.4.3 after live measurements). When no
38//!   cosine signal exists, `alpha = 0` effectively — i.e. pure
39//!   lexical. See [`context::memory_candidates`] for the wiring.
40
41use kimetsu_core::KimetsuResult;
42
43/// Default blend factor for hybrid scoring. 0.0 = pure lexical
44/// (v0.4.1 behavior), 1.0 = pure cosine. v0.4.2 ships 0.5 as a
45/// starting point; live data in v0.4.3 will tune it.
46pub const DEFAULT_HYBRID_ALPHA: f32 = 0.5;
47
48/// Embedder trait. Every implementation maps a text to a fixed
49/// dimension `Vec<f32>` and identifies itself via a stable model id.
50///
51/// `Send + Sync` so a single embedder instance can be shared across
52/// the chat REPL's threads (drainer, REPL, hook runner) without
53/// requiring per-call locking.
54pub trait Embedder: Send + Sync {
55    /// Compute an embedding for `text`. The returned vector MUST have
56    /// length == `self.dim()`. Implementations may normalize.
57    fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedderError>;
58
59    /// Stable identifier for the model used. Stored alongside each
60    /// embedding so retrieval can detect cross-model mismatches and
61    /// skip the cosine blend rather than comparing apples-to-oranges.
62    fn model_id(&self) -> &str;
63
64    /// Embedding dimension. Used by the BLOB codec to validate
65    /// stored vectors against the active model on retrieval.
66    fn dim(&self) -> usize;
67
68    /// Convenience: true when this embedder is the production no-op.
69    /// Callers can short-circuit the write/retrieval cosine path
70    /// instead of allocating a vec only to discard it.
71    fn is_noop(&self) -> bool {
72        false
73    }
74}
75
76/// Implement `Embedder` for `Box<dyn Embedder>` so callers can hold
77/// an owned trait object without ceremony.
78impl Embedder for Box<dyn Embedder> {
79    fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedderError> {
80        (**self).embed(text)
81    }
82    fn model_id(&self) -> &str {
83        (**self).model_id()
84    }
85    fn dim(&self) -> usize {
86        (**self).dim()
87    }
88    fn is_noop(&self) -> bool {
89        (**self).is_noop()
90    }
91}
92
93/// Failure modes for an embedder.
94#[derive(Debug, Clone)]
95pub enum EmbedderError {
96    /// The embedder is intentionally a no-op — no embeddings will be
97    /// produced. Callers should fall back to a NULL embedding /
98    /// lexical-only retrieval.
99    NotImplemented,
100    /// Model failed to load. v0.4.3+ — e.g. the fastembed backend
101    /// can't download the model.
102    LoadFailed(String),
103    /// Inference failed (rare).
104    EmbedFailed(String),
105    /// Dimension mismatch between the embedder and a stored row.
106    /// The retrieval path skips this row's cosine contribution.
107    DimMismatch { expected: usize, got: usize },
108}
109
110impl std::fmt::Display for EmbedderError {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        match self {
113            Self::NotImplemented => write!(f, "embedder not implemented"),
114            Self::LoadFailed(msg) => write!(f, "embedder load failed: {msg}"),
115            Self::EmbedFailed(msg) => write!(f, "embed call failed: {msg}"),
116            Self::DimMismatch { expected, got } => {
117                write!(f, "embedding dim mismatch: expected {expected}, got {got}")
118            }
119        }
120    }
121}
122
123impl std::error::Error for EmbedderError {}
124
125/// Production default when no real embedder is configured.
126/// `embed()` returns `Err(NotImplemented)`; callers interpret that
127/// as "store NULL" / "skip the cosine blend".
128#[derive(Debug, Default, Clone, Copy)]
129pub struct NoopEmbedder;
130
131impl NoopEmbedder {
132    pub const MODEL_ID: &'static str = "noop";
133}
134
135impl Embedder for NoopEmbedder {
136    fn embed(&self, _text: &str) -> Result<Vec<f32>, EmbedderError> {
137        Err(EmbedderError::NotImplemented)
138    }
139
140    fn model_id(&self) -> &str {
141        Self::MODEL_ID
142    }
143
144    fn dim(&self) -> usize {
145        0
146    }
147
148    fn is_noop(&self) -> bool {
149        true
150    }
151}
152
153/// Deterministic, dependency-free pseudo-embedder used in tests.
154///
155/// Hashes each word into a fixed-dim bucket (count-of-hash-buckets),
156/// then L2-normalizes the resulting vector. NOT semantic — texts
157/// that share words will be close; texts that don't share words
158/// will be far. Good enough to exercise the hybrid-scoring code
159/// path without depending on a real ML model.
160///
161/// Default dimension is 8 (small enough to keep tests fast).
162#[derive(Debug, Clone, Copy)]
163pub struct StubEmbedder {
164    dim: usize,
165}
166
167impl StubEmbedder {
168    pub const MODEL_ID: &'static str = "stub-d8";
169
170    pub const fn new() -> Self {
171        Self { dim: 8 }
172    }
173
174    pub const fn with_dim(dim: usize) -> Self {
175        Self { dim }
176    }
177}
178
179impl Default for StubEmbedder {
180    fn default() -> Self {
181        Self::new()
182    }
183}
184
185impl Embedder for StubEmbedder {
186    fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedderError> {
187        let mut bucket = vec![0.0f32; self.dim];
188        for word in text.split_whitespace() {
189            // Cheap, stable hash. Don't use DefaultHasher — its seed
190            // randomizes across processes and tests would be flaky.
191            // FNV-1a over the lowercased UTF-8 bytes is plenty.
192            let normalized = word.to_lowercase();
193            let mut h: u64 = 0xcbf2_9ce4_8422_2325;
194            for byte in normalized.bytes() {
195                h ^= byte as u64;
196                h = h.wrapping_mul(0x0000_0100_0000_01B3);
197            }
198            let idx = (h as usize) % self.dim.max(1);
199            bucket[idx] += 1.0;
200        }
201        // L2-normalize so cosine similarity reduces to a dot product.
202        let norm = bucket.iter().map(|v| v * v).sum::<f32>().sqrt();
203        if norm > 0.0 {
204            for v in &mut bucket {
205                *v /= norm;
206            }
207        }
208        Ok(bucket)
209    }
210
211    fn model_id(&self) -> &str {
212        Self::MODEL_ID
213    }
214
215    fn dim(&self) -> usize {
216        self.dim
217    }
218}
219
220/// Open the production-default embedder.
221///
222/// Resolution (v0.4.3):
223///   1. `KIMETSU_BRAIN_EMBEDDER=noop|off|none` → always `NoopEmbedder`,
224///      regardless of Cargo features. Useful for CI, hooks, and
225///      transient subprocesses that shouldn't pay the model-load
226///      cost.
227///   2. Cargo feature `embeddings` enabled →
228///      [`fastembed_backend::open_cached`] returns a process-wide
229///      cached [`FastembedEmbedder`] for the model picked by
230///      [`pick_builtin_model_from_env`] (default `bge-small-en-v1.5`,
231///      `bge-m3` or `jina-v2-base-code` opt-in via env). On model
232///      load failure (network, disk, ort runtime missing) we log
233///      and fall through to Noop so the brain stays usable on FTS
234///      alone.
235///   3. Cargo feature `embeddings` disabled → `NoopEmbedder`,
236///      identical to v0.4.2 build.
237///
238/// The returned trait object is borrowed from a process-static
239/// `OnceLock`; production callers get model-load cost paid exactly
240/// once over the process lifetime. Tests that need a different
241/// embedder must use [`crate::context::retrieve_context_with_embedder`]
242/// with an explicit [`StubEmbedder`] (or any other [`Embedder`])
243/// instead of going through this function.
244pub fn open_default_embedder() -> &'static (dyn Embedder + Send + Sync) {
245    static CACHE: std::sync::OnceLock<Box<dyn Embedder + Send + Sync>> =
246        std::sync::OnceLock::new();
247    let embedder = CACHE.get_or_init(build_default_embedder);
248    embedder.as_ref()
249}
250
251fn build_default_embedder() -> Box<dyn Embedder + Send + Sync> {
252    if env_disables_embedder() {
253        return Box::new(NoopEmbedder);
254    }
255    #[cfg(feature = "embeddings")]
256    {
257        match fastembed_backend::open_cached() {
258            Ok(handle) => return Box::new(handle),
259            Err(err) => {
260                eprintln!(
261                    "kimetsu-brain: fastembed init failed ({err}); falling back to NoopEmbedder. \
262                     Retrieval will stay FTS-only this session. Re-run with \
263                     KIMETSU_BRAIN_EMBEDDER=noop to silence this warning."
264                );
265            }
266        }
267    }
268    Box::new(NoopEmbedder)
269}
270
271/// v0.4.3: env-driven kill switch. Truthy values (1/true/yes/on)
272/// force-disable the embedder for this process; "noop", "off",
273/// "none" do the same. Anything else (or unset) leaves the
274/// `embeddings` feature in control.
275fn env_disables_embedder() -> bool {
276    match std::env::var("KIMETSU_BRAIN_EMBEDDER") {
277        Ok(value) => {
278            let v = value.trim().to_ascii_lowercase();
279            matches!(v.as_str(), "noop" | "off" | "none" | "0" | "false" | "no")
280        }
281        Err(_) => false,
282    }
283}
284
285/// v0.4.3: pick which builtin model to load from the env, returning
286/// a stable identifier. Used both by the fastembed backend (to map
287/// id → `EmbeddingModel`) and by `kimetsu brain reindex` (to label
288/// new rows with the right `embedding_model`).
289///
290/// Resolution:
291///   * unset / "" / "default" / "bge-small" / "bge-small-en-v1.5"
292///     → `"bge-small-en-v1.5"` (384 dim, ~67 MB int8, English)
293///   * "bge-m3"
294///     → `"bge-m3"` (1024 dim, ~600 MB int8, multilingual)
295///   * "jina-code" / "jina-v2-base-code" /
296///     "jina-embeddings-v2-base-code"
297///     → `"jina-v2-base-code"` (768 dim, ~165 MB int8, English +
298///     code-tuned)
299///   * anything else falls back to bge-small with a warning.
300pub fn pick_builtin_model_from_env() -> &'static str {
301    let raw = std::env::var("KIMETSU_BRAIN_EMBEDDER").ok();
302    let v = raw
303        .as_deref()
304        .map(|s| s.trim().to_ascii_lowercase())
305        .unwrap_or_default();
306    match v.as_str() {
307        "" | "default" | "bge-small" | "bge-small-en-v1.5" => "bge-small-en-v1.5",
308        "bge-m3" | "m3" => "bge-m3",
309        "jina-code" | "jina-v2-base-code" | "jina-embeddings-v2-base-code" => "jina-v2-base-code",
310        // "noop"/etc. are handled by env_disables_embedder; this
311        // function shouldn't be called in those cases but defensively
312        // pick the lean default.
313        "noop" | "off" | "none" | "0" | "false" | "no" => "bge-small-en-v1.5",
314        other => {
315            eprintln!(
316                "kimetsu-brain: unknown KIMETSU_BRAIN_EMBEDDER={other:?}, \
317                 falling back to bge-small-en-v1.5"
318            );
319            "bge-small-en-v1.5"
320        }
321    }
322}
323
324// v0.4.3: real fastembed-backed embedder. Lives behind the
325// `embeddings` Cargo feature so the default build skips the
326// ~50-transitive-crate dep tree (ONNX runtime, tokenizers, etc).
327#[cfg(feature = "embeddings")]
328mod fastembed_backend {
329    use super::{Embedder, EmbedderError, pick_builtin_model_from_env};
330    use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
331    use std::sync::{Arc, Mutex, OnceLock};
332
333    /// fastembed-backed embedder. Wraps the ONNX runtime in a
334    /// `Mutex` because `TextEmbedding::embed` takes `&mut self`.
335    /// The lock window is short (one inference per call); the
336    /// chat REPL's threads serialize cleanly through it.
337    pub struct FastembedEmbedder {
338        model_id: &'static str,
339        dim: usize,
340        engine: Mutex<TextEmbedding>,
341    }
342
343    impl FastembedEmbedder {
344        pub fn try_open(builtin_id: &str) -> Result<Self, EmbedderError> {
345            let (kind, model_id, dim) = match builtin_id {
346                "bge-m3" => (EmbeddingModel::BGEM3, "bge-m3", 1024),
347                "jina-v2-base-code" => (
348                    EmbeddingModel::JinaEmbeddingsV2BaseCode,
349                    "jina-v2-base-code",
350                    768,
351                ),
352                // bge-small-en-v1.5 is the default + fallback.
353                _ => (EmbeddingModel::BGESmallENV15, "bge-small-en-v1.5", 384),
354            };
355            let opts = InitOptions::new(kind).with_show_download_progress(false);
356            let engine = TextEmbedding::try_new(opts)
357                .map_err(|e| EmbedderError::LoadFailed(format!("fastembed init: {e}")))?;
358            Ok(Self {
359                model_id,
360                dim,
361                engine: Mutex::new(engine),
362            })
363        }
364    }
365
366    impl Embedder for FastembedEmbedder {
367        fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedderError> {
368            let mut guard = self
369                .engine
370                .lock()
371                .unwrap_or_else(|poisoned| poisoned.into_inner());
372            let mut out = guard
373                .embed(vec![text], None)
374                .map_err(|e| EmbedderError::EmbedFailed(format!("fastembed embed: {e}")))?;
375            let vec = out
376                .pop()
377                .ok_or_else(|| EmbedderError::EmbedFailed("empty result".into()))?;
378            if vec.len() != self.dim {
379                return Err(EmbedderError::DimMismatch {
380                    expected: self.dim,
381                    got: vec.len(),
382                });
383            }
384            Ok(vec)
385        }
386
387        fn model_id(&self) -> &str {
388            self.model_id
389        }
390
391        fn dim(&self) -> usize {
392            self.dim
393        }
394    }
395
396    /// Shared handle. `open_default_embedder` boxes this into a
397    /// `dyn Embedder` and stashes it in a process-static `OnceLock`,
398    /// so we only call into ONNX once per process.
399    #[derive(Clone)]
400    pub struct EmbedderHandle(Arc<FastembedEmbedder>);
401
402    impl Embedder for EmbedderHandle {
403        fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedderError> {
404            self.0.embed(text)
405        }
406        fn model_id(&self) -> &str {
407            self.0.model_id()
408        }
409        fn dim(&self) -> usize {
410            self.0.dim()
411        }
412    }
413
414    /// Open (or return the cached) fastembed embedder for the model
415    /// picked by `KIMETSU_BRAIN_EMBEDDER`. Errors here propagate up
416    /// to `open_default_embedder`, which falls back to Noop +
417    /// prints a one-line warning.
418    pub fn open_cached() -> Result<EmbedderHandle, EmbedderError> {
419        static CELL: OnceLock<Result<Arc<FastembedEmbedder>, EmbedderError>> = OnceLock::new();
420        let init = CELL.get_or_init(|| {
421            let builtin = pick_builtin_model_from_env();
422            FastembedEmbedder::try_open(builtin).map(Arc::new)
423        });
424        match init {
425            Ok(arc) => Ok(EmbedderHandle(arc.clone())),
426            Err(err) => Err(err.clone()),
427        }
428    }
429}
430
431// --------- math helpers ---------
432
433/// Cosine similarity between two vectors. Returns 0.0 when either
434/// vector is empty or all-zeros. Does NOT assume the vectors are
435/// pre-normalized — divides by both norms.
436pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
437    if a.is_empty() || b.is_empty() || a.len() != b.len() {
438        return 0.0;
439    }
440    let mut dot = 0.0f32;
441    let mut na = 0.0f32;
442    let mut nb = 0.0f32;
443    for (x, y) in a.iter().zip(b.iter()) {
444        dot += x * y;
445        na += x * x;
446        nb += y * y;
447    }
448    if na == 0.0 || nb == 0.0 {
449        return 0.0;
450    }
451    dot / (na.sqrt() * nb.sqrt())
452}
453
454// --------- write-path helper ---------
455
456/// Compute the embedding for `text` and persist it onto an existing
457/// `memories.memory_id` row.
458///
459/// No-op when the embedder is intentionally a [`NoopEmbedder`] (it
460/// returns [`EmbedderError::NotImplemented`] which we silently
461/// swallow — the column stays NULL, retrieval falls back to
462/// FTS-only for the row, exact v0.4.1 behavior).
463///
464/// For other embedder errors we surface them up. The caller
465/// (`add_memory`, `add_user_memory`) can decide whether to fail the
466/// whole insert or log+continue — today they propagate.
467pub fn embed_and_persist(
468    conn: &rusqlite::Connection,
469    memory_id: &str,
470    text: &str,
471    embedder: &dyn Embedder,
472) -> KimetsuResult<()> {
473    if embedder.is_noop() {
474        return Ok(());
475    }
476    let vec = match embedder.embed(text) {
477        Ok(v) => v,
478        // NotImplemented is the contract for "skip silently". Treat
479        // any embedder that signals it the same way as NoopEmbedder.
480        Err(EmbedderError::NotImplemented) => return Ok(()),
481        Err(e) => return Err(format!("embed failed for memory {memory_id}: {e}").into()),
482    };
483    if vec.len() != embedder.dim() {
484        return Err(format!(
485            "embedder {} produced {} dims, expected {}",
486            embedder.model_id(),
487            vec.len(),
488            embedder.dim()
489        )
490        .into());
491    }
492    let blob = encode_embedding(&vec);
493    conn.execute(
494        "UPDATE memories SET embedding = ?1, embedding_model = ?2 WHERE memory_id = ?3",
495        rusqlite::params![blob, embedder.model_id(), memory_id],
496    )?;
497    Ok(())
498}
499
500// --------- BLOB codec ---------
501//
502// Embeddings are stored as little-endian f32 BLOBs. The encoder
503// fixes byte order so brain.db files move between architectures.
504// The decoder is strict: it returns Err if the byte length isn't a
505// multiple of 4, or if the resulting dim doesn't match expectations.
506
507/// Serialize a float vector to little-endian bytes for storage.
508pub fn encode_embedding(vec: &[f32]) -> Vec<u8> {
509    let mut out = Vec::with_capacity(vec.len() * 4);
510    for v in vec {
511        out.extend_from_slice(&v.to_le_bytes());
512    }
513    out
514}
515
516/// Decode a BLOB back into a float vector. Optionally validates the
517/// expected dimension; pass `None` to accept any length.
518pub fn decode_embedding(bytes: &[u8], expected_dim: Option<usize>) -> KimetsuResult<Vec<f32>> {
519    if !bytes.len().is_multiple_of(4) {
520        return Err(format!(
521            "embedding blob length {} not a multiple of 4",
522            bytes.len()
523        )
524        .into());
525    }
526    let dim = bytes.len() / 4;
527    if let Some(expected) = expected_dim
528        && dim != expected
529    {
530        return Err(format!(
531            "embedding blob dim {dim} does not match expected {expected}"
532        )
533        .into());
534    }
535    let mut out = Vec::with_capacity(dim);
536    for chunk in bytes.chunks_exact(4) {
537        let mut buf = [0u8; 4];
538        buf.copy_from_slice(chunk);
539        out.push(f32::from_le_bytes(buf));
540    }
541    Ok(out)
542}
543
544#[cfg(test)]
545mod tests {
546    use super::*;
547
548    #[test]
549    fn noop_embedder_returns_not_implemented_and_is_noop() {
550        let e = NoopEmbedder;
551        assert!(e.is_noop());
552        assert_eq!(e.dim(), 0);
553        assert_eq!(e.model_id(), "noop");
554        assert!(matches!(
555            e.embed("hello").unwrap_err(),
556            EmbedderError::NotImplemented
557        ));
558    }
559
560    #[test]
561    fn stub_embedder_is_deterministic() {
562        let e = StubEmbedder::new();
563        let a = e.embed("hello rust").expect("embed a");
564        let b = e.embed("hello rust").expect("embed b");
565        let c = e.embed("hello RUST").expect("embed c");
566        assert_eq!(a, b, "same input -> same output");
567        assert_eq!(
568            a, c,
569            "lowercasing means case differences collapse to the same vector"
570        );
571        assert_eq!(a.len(), 8);
572        // L2-normalized: norm == 1 within float tolerance.
573        let norm = a.iter().map(|v| v * v).sum::<f32>().sqrt();
574        assert!((norm - 1.0).abs() < 1e-5, "expected unit norm, got {norm}");
575    }
576
577    #[test]
578    fn stub_embedder_distinguishes_disjoint_inputs() {
579        let e = StubEmbedder::new();
580        let a = e.embed("foo bar").expect("a");
581        let b = e.embed("qux quux").expect("b");
582        let sim = cosine_similarity(&a, &b);
583        // Disjoint word sets *can* still collide in the 8-bucket
584        // hash, but on average should be low. Sanity bound: not 1.0.
585        assert!(sim < 0.99, "disjoint inputs should not be near-identical: {sim}");
586    }
587
588    #[test]
589    fn stub_embedder_handles_empty_input() {
590        let e = StubEmbedder::new();
591        let v = e.embed("").expect("empty embed");
592        assert_eq!(v.len(), 8);
593        // All zeros: cosine similarity with self is 0 (we guard against
594        // division by zero), which is exactly the behavior the
595        // retrieval blender wants for content-free queries.
596        assert!(v.iter().all(|&x| x == 0.0));
597    }
598
599    #[test]
600    fn cosine_similarity_handles_edge_cases() {
601        // Identical normalized vectors -> 1.0.
602        let a = [1.0f32, 0.0, 0.0];
603        assert!((cosine_similarity(&a, &a) - 1.0).abs() < 1e-6);
604
605        // Orthogonal -> 0.0.
606        let b = [0.0f32, 1.0, 0.0];
607        assert!((cosine_similarity(&a, &b)).abs() < 1e-6);
608
609        // Anti-parallel -> -1.0.
610        let c = [-1.0f32, 0.0, 0.0];
611        assert!((cosine_similarity(&a, &c) + 1.0).abs() < 1e-6);
612
613        // Empty / mismatched dim -> 0.0 by contract.
614        assert_eq!(cosine_similarity(&[], &a), 0.0);
615        assert_eq!(cosine_similarity(&a, &[0.0]), 0.0);
616
617        // Zero norm -> 0.0 by contract (don't divide by zero).
618        let zeros = [0.0f32, 0.0, 0.0];
619        assert_eq!(cosine_similarity(&zeros, &a), 0.0);
620    }
621
622    #[test]
623    fn cosine_similarity_is_symmetric() {
624        let a = [0.6f32, 0.8, 0.0];
625        let b = [0.0f32, 1.0, 0.0];
626        let ab = cosine_similarity(&a, &b);
627        let ba = cosine_similarity(&b, &a);
628        assert!((ab - ba).abs() < 1e-6);
629        // Dot is 0.8, |a|=1, |b|=1 -> sim = 0.8.
630        assert!((ab - 0.8).abs() < 1e-5);
631    }
632
633    #[test]
634    fn encode_decode_embedding_round_trip() {
635        let vec = vec![0.1f32, -0.2, 3.14, -0.000_001, 42.0];
636        let blob = encode_embedding(&vec);
637        assert_eq!(blob.len(), vec.len() * 4);
638        let back = decode_embedding(&blob, Some(vec.len())).expect("decode");
639        assert_eq!(back.len(), vec.len());
640        for (orig, got) in vec.iter().zip(back.iter()) {
641            assert!(
642                (orig - got).abs() < 1e-7,
643                "f32 round-trip should be bit-exact"
644            );
645        }
646    }
647
648    #[test]
649    fn decode_embedding_rejects_unaligned_blob() {
650        let bad = [0u8, 1, 2]; // 3 bytes - not a multiple of 4
651        let err = decode_embedding(&bad, None).unwrap_err();
652        assert!(err.to_string().contains("not a multiple of 4"));
653    }
654
655    #[test]
656    fn decode_embedding_rejects_dim_mismatch() {
657        let vec = vec![1.0f32, 2.0, 3.0];
658        let blob = encode_embedding(&vec);
659        let err = decode_embedding(&blob, Some(5)).unwrap_err();
660        assert!(err.to_string().contains("does not match expected"));
661    }
662
663    /// v0.4.3: under the default Cargo build (no `embeddings` feature)
664    /// `open_default_embedder` MUST return Noop so a `cargo install
665    /// kimetsu-cli` user doesn't accidentally start downloading a
666    /// model from $HOME. Skip when `--features embeddings` is on —
667    /// that build path has its own integration tests (run with
668    /// `cargo test --features embeddings -- --ignored`).
669    #[cfg(not(feature = "embeddings"))]
670    #[test]
671    fn open_default_embedder_returns_noop_on_default_build() {
672        let e = open_default_embedder();
673        assert!(e.is_noop());
674        assert_eq!(e.dim(), 0);
675        assert!(matches!(
676            e.embed("anything").unwrap_err(),
677            EmbedderError::NotImplemented
678        ));
679    }
680
681    /// v0.4.3: env kill-switch works even when the `embeddings`
682    /// feature is on — `KIMETSU_BRAIN_EMBEDDER=noop` returns Noop
683    /// regardless. Tests the env parser directly rather than going
684    /// through the cached `open_default_embedder`, which would
685    /// otherwise be poisoned by whatever the previous test in the
686    /// process initialized.
687    #[test]
688    fn env_disables_embedder_recognizes_off_values() {
689        let lock = crate::user_brain::test_env_lock()
690            .lock()
691            .unwrap_or_else(|p| p.into_inner());
692        let prev = std::env::var("KIMETSU_BRAIN_EMBEDDER").ok();
693        for value in ["noop", "off", "NONE", "0", "false", "no"] {
694            // SAFETY: serialized via the shared brain test env lock.
695            unsafe {
696                std::env::set_var("KIMETSU_BRAIN_EMBEDDER", value);
697            }
698            assert!(env_disables_embedder(), "value {value:?} must disable");
699        }
700        for value in ["", "default", "bge-small", "bge-m3", "jina-code"] {
701            unsafe {
702                std::env::set_var("KIMETSU_BRAIN_EMBEDDER", value);
703            }
704            assert!(
705                !env_disables_embedder(),
706                "value {value:?} must NOT disable"
707            );
708        }
709        // Restore.
710        unsafe {
711            match prev {
712                Some(v) => std::env::set_var("KIMETSU_BRAIN_EMBEDDER", v),
713                None => std::env::remove_var("KIMETSU_BRAIN_EMBEDDER"),
714            }
715        }
716        drop(lock);
717    }
718
719    /// v0.4.3: model picker maps the user-facing env string onto a
720    /// stable model id used by both fastembed init AND the
721    /// `embedding_model` column on each memory row.
722    #[test]
723    fn pick_builtin_model_from_env_handles_aliases() {
724        let lock = crate::user_brain::test_env_lock()
725            .lock()
726            .unwrap_or_else(|p| p.into_inner());
727        let prev = std::env::var("KIMETSU_BRAIN_EMBEDDER").ok();
728        let cases = [
729            ("", "bge-small-en-v1.5"),
730            ("default", "bge-small-en-v1.5"),
731            ("bge-small", "bge-small-en-v1.5"),
732            ("BGE-SMALL-EN-V1.5", "bge-small-en-v1.5"),
733            ("bge-m3", "bge-m3"),
734            ("M3", "bge-m3"),
735            ("jina-code", "jina-v2-base-code"),
736            ("jina-v2-base-code", "jina-v2-base-code"),
737            ("jina-embeddings-v2-base-code", "jina-v2-base-code"),
738            // Unknown values fall back to bge-small with a warning.
739            ("totally-made-up", "bge-small-en-v1.5"),
740        ];
741        for (input, expected) in cases {
742            // SAFETY: serialized via the shared brain test env lock.
743            unsafe {
744                std::env::set_var("KIMETSU_BRAIN_EMBEDDER", input);
745            }
746            assert_eq!(
747                pick_builtin_model_from_env(),
748                expected,
749                "input {input:?} -> expected {expected}"
750            );
751        }
752        unsafe {
753            match prev {
754                Some(v) => std::env::set_var("KIMETSU_BRAIN_EMBEDDER", v),
755                None => std::env::remove_var("KIMETSU_BRAIN_EMBEDDER"),
756            }
757        }
758        drop(lock);
759    }
760}