klieo_embed_common/
lib.rs1#![deny(missing_docs)]
2#![deny(rust_2018_idioms)]
3#![deny(rustdoc::broken_intra_doc_links)]
4
5use async_trait::async_trait;
23use klieo_core::error::MemoryError;
24
25#[cfg(feature = "fastembed")]
26mod fastembed;
27#[cfg(feature = "fastembed")]
28pub use fastembed::{FastEmbedEmbedder, DEFAULT_DIM, DEFAULT_MODEL};
29
30#[async_trait]
36pub trait Embedder: Send + Sync {
37 fn dimension(&self) -> usize;
41
42 async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, MemoryError>;
44}
45
46pub struct DummyEmbedder;
51
52#[async_trait]
53impl Embedder for DummyEmbedder {
54 fn dimension(&self) -> usize {
55 384
56 }
57
58 async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, MemoryError> {
59 Ok(texts.iter().map(|_| vec![0.0f32; 384]).collect())
60 }
61}
62
63#[cfg(any(test, feature = "test-utils"))]
67pub struct FakeEmbedder {
68 dim: usize,
69}
70
71#[cfg(any(test, feature = "test-utils"))]
72impl FakeEmbedder {
73 pub fn new(dim: usize) -> Self {
75 Self { dim }
76 }
77}
78
79#[cfg(any(test, feature = "test-utils"))]
80impl Default for FakeEmbedder {
81 fn default() -> Self {
82 Self::new(8)
83 }
84}
85
86#[cfg(any(test, feature = "test-utils"))]
87#[async_trait]
88impl Embedder for FakeEmbedder {
89 fn dimension(&self) -> usize {
90 self.dim
91 }
92
93 async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, MemoryError> {
94 let dim = self.dim;
95 Ok(texts
96 .iter()
97 .map(|text| {
98 let mut v = vec![0.0f32; dim];
101 let bytes = text.as_bytes();
102 for (i, slot) in v.iter_mut().enumerate() {
103 const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
104 const FNV_PRIME: u64 = 0x0000_0001_0000_01b3;
105 let mut h: u64 = FNV_OFFSET;
106 h ^= i as u64;
107 h = h.wrapping_mul(FNV_PRIME);
108 for &b in bytes {
109 h ^= b as u64;
110 h = h.wrapping_mul(FNV_PRIME);
111 }
112 *slot = (h as f32 / u64::MAX as f32) - 0.5;
113 }
114 let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
115 if norm > 0.0 {
116 for x in &mut v {
117 *x /= norm;
118 }
119 }
120 v
121 })
122 .collect())
123 }
124}
125
126#[cfg(test)]
127mod tests {
128 use super::*;
129
130 #[tokio::test]
131 async fn dummy_returns_zero_vectors_of_dim_384() {
132 let e = DummyEmbedder;
133 assert_eq!(e.dimension(), 384);
134 let out = e.embed(&["a".into(), "b".into()]).await.unwrap();
135 assert_eq!(out.len(), 2);
136 assert_eq!(out[0].len(), 384);
137 assert!(out[0].iter().all(|x| *x == 0.0));
138 }
139
140 #[tokio::test]
141 async fn fake_embedder_is_deterministic() {
142 let e = FakeEmbedder::new(16);
143 let a = e.embed(&["hello".into()]).await.unwrap();
144 let b = e.embed(&["hello".into()]).await.unwrap();
145 assert_eq!(a, b);
146 }
147
148 #[tokio::test]
149 async fn fake_embedder_distinguishes_inputs() {
150 let e = FakeEmbedder::new(16);
151 let a = e.embed(&["alpha".into()]).await.unwrap();
152 let b = e.embed(&["beta".into()]).await.unwrap();
153 assert_ne!(a, b);
154 }
155
156 #[tokio::test]
157 async fn fake_embedder_outputs_unit_vectors() {
158 let e = FakeEmbedder::new(8);
159 let v = e.embed(&["hello".into()]).await.unwrap();
160 let norm: f32 = v[0].iter().map(|x| x * x).sum::<f32>().sqrt();
161 assert!(
162 (norm - 1.0).abs() < 1e-5,
163 "fake embedder must produce unit vectors, got norm={norm}"
164 );
165 }
166}