Skip to main content

ijima_core/
embeddings.rs

1// Copyright (C) 2026 Industrial Algebra
2// SPDX-License-Identifier: Apache-2.0
3
4//! Embedding contract — pure trait, no backend.
5//!
6//! The trait lives in `ijima-core` so the store, miner, and client can
7//! depend on it without pulling in a heavy ML backend. Concrete
8//! implementations (candle, remote API, ...) live in `ijima-server`
9//! behind feature gates.
10//!
11//! ## Default dimensionality
12//!
13//! [`DEFAULT_EMBEDDING_DIM`] is 384 to match pi-mempalace's
14//! `all-MiniLM-L6-v2`, so the live `memories.db` corpus migrates without
15//! re-embedding. Configurable dimensions are a future concern (requires a
16//! re-embed pass on dimension change).
17
18use crate::Result;
19
20/// Default embedding dimensionality — 384, matching pi-mempalace's
21/// `all-MiniLM-L6-v2` for migration parity.
22pub const DEFAULT_EMBEDDING_DIM: usize = 384;
23
24/// A dense embedding vector for a piece of text.
25#[derive(Debug, Clone, PartialEq)]
26#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
27#[cfg_attr(feature = "serde", serde(transparent))]
28pub struct Embedding(pub Vec<f32>);
29
30impl Embedding {
31    /// Returns the dimensionality of the vector.
32    pub fn dim(&self) -> usize {
33        self.0.len()
34    }
35
36    /// Returns the underlying slice.
37    pub fn as_slice(&self) -> &[f32] {
38        &self.0
39    }
40}
41
42/// The embedding contract every backend implements.
43///
44/// Backends ship in `ijima-server` (candle is the IA-standard default,
45/// consistent with Quantizon); a remote-API backend may follow.
46///
47/// `Send + Sync` so an `Arc<dyn Embedder>` can be shared across an
48/// async store's worker pool.
49pub trait Embedder: Send + Sync {
50    /// Dimensionality of the vectors this embedder produces.
51    fn dim(&self) -> usize;
52
53    /// Embeds a single piece of text.
54    ///
55    /// # Errors
56    ///
57    /// Returns [`crate::IjimaError::Store`] on a backend failure (model
58    /// load, inference, device error).
59    fn embed(&self, text: &str) -> Result<Embedding>;
60
61    /// Embeds a batch of texts. The default loops over [`embed`];
62    /// backends with batch acceleration override this.
63    ///
64    /// # Errors
65    ///
66    /// Propagates any per-item embedding error.
67    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Embedding>> {
68        texts.iter().map(|t| self.embed(t)).collect()
69    }
70
71    /// The model id that produced these embeddings (e.g.
72    /// `sentence-transformers/all-MiniLM-L6-v2@main`). Used for
73    /// **embedding provenance** (D10): memories stamp the model that
74    /// embedded them, so a model swap is detectable and a re-embed pass
75    /// can be triggered rather than silently producing incomparable
76    /// vectors. Default `"unknown"` for backends that don't track it.
77    fn model_id(&self) -> &str {
78        "unknown"
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    /// A deterministic toy embedder used only to exercise the trait
87    /// contract in core. Real backends live in `ijima-server`.
88    struct ConstEmbedder;
89    impl Embedder for ConstEmbedder {
90        fn dim(&self) -> usize {
91            2
92        }
93        fn embed(&self, text: &str) -> Result<Embedding> {
94            Ok(Embedding(vec![text.len() as f32, 0.0]))
95        }
96    }
97
98    #[test]
99    fn default_dim_matches_mempalace_for_migration() {
100        assert_eq!(DEFAULT_EMBEDDING_DIM, 384);
101    }
102
103    #[test]
104    fn embed_batch_defaults_to_per_item_loop() {
105        let e = ConstEmbedder;
106        let got = e.embed_batch(&["a", "bb", "ccc"]).expect("must embed");
107        assert_eq!(got.len(), 3);
108        assert_eq!(got[0].dim(), 2);
109        assert_eq!(got[2].as_slice(), &[3.0, 0.0]);
110    }
111}