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}