pub use cortex_store::repo::EmbedRecord;
pub type EmbedResult<T> = Result<T, EmbedError>;
#[derive(Debug, thiserror::Error)]
pub enum EmbedError {
#[error("embed input rejected: {0}")]
InvalidInput(String),
#[error("embed backend failed: {0}")]
Backend(String),
#[error("embed dimension mismatch: backend `{backend_id}` advertises dim {expected} but produced {actual}")]
DimensionMismatch {
backend_id: String,
expected: usize,
actual: usize,
},
}
pub trait Embedder: Send + Sync {
fn backend_id(&self) -> &str;
fn dim(&self) -> usize;
fn embed(&self, text: &str, tags: &[String]) -> EmbedResult<Vec<f32>>;
}
#[must_use]
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
if a.len() != b.len() || a.is_empty() {
return 0.0;
}
let mut dot = 0.0f32;
let mut na = 0.0f32;
let mut nb = 0.0f32;
for (x, y) in a.iter().zip(b.iter()) {
if x.is_nan() || y.is_nan() {
return 0.0;
}
dot += x * y;
na += x * x;
nb += y * y;
}
if na <= 0.0 || nb <= 0.0 {
return 0.0;
}
let denom = na.sqrt() * nb.sqrt();
if denom <= 0.0 || !denom.is_finite() {
return 0.0;
}
let raw = dot / denom;
if raw.is_nan() {
return 0.0;
}
raw.clamp(-1.0, 1.0)
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::{DateTime, TimeZone, Utc};
use cortex_core::MemoryId;
fn at(second: u32) -> DateTime<Utc> {
Utc.with_ymd_and_hms(2026, 1, 1, 12, 0, second).unwrap()
}
fn mem(id: &str) -> MemoryId {
id.parse().expect("valid memory id")
}
#[test]
fn cosine_similarity_handles_orthogonal_zero() {
let a = vec![1.0f32, 0.0, 0.0];
let b = vec![0.0f32, 1.0, 0.0];
let sim = cosine_similarity(&a, &b);
assert!(
(sim - 0.0).abs() < 1e-6,
"orthogonal vectors must produce similarity ~0, got {sim}"
);
}
#[test]
fn cosine_similarity_handles_identical_one() {
let a = vec![0.3f32, -0.7, 0.5, 0.1];
let sim = cosine_similarity(&a, &a);
assert!(
(sim - 1.0).abs() < 1e-6,
"identical vectors must produce similarity 1.0, got {sim}"
);
}
#[test]
fn cosine_similarity_handles_opposite_negative_one() {
let a = vec![0.3f32, -0.7, 0.5, 0.1];
let b: Vec<f32> = a.iter().map(|v| -v).collect();
let sim = cosine_similarity(&a, &b);
assert!(
(sim - (-1.0)).abs() < 1e-6,
"antipodal vectors must produce similarity -1.0, got {sim}"
);
}
#[test]
fn cosine_similarity_handles_nan_safely() {
let a = vec![1.0f32, f32::NAN, 0.0];
let b = vec![1.0f32, 1.0, 1.0];
let sim = cosine_similarity(&a, &b);
assert!(
!sim.is_nan(),
"cosine_similarity must never return NaN, got {sim}"
);
assert_eq!(sim, 0.0);
let zeros = vec![0.0f32; 3];
let sim_zero = cosine_similarity(&zeros, &[1.0, 2.0, 3.0]);
assert_eq!(sim_zero, 0.0);
let sim_mismatch = cosine_similarity(&[1.0, 0.0], &[1.0, 0.0, 0.0]);
assert_eq!(sim_mismatch, 0.0);
let sim_empty: f32 = cosine_similarity(&[], &[]);
assert_eq!(sim_empty, 0.0);
}
#[test]
fn embed_record_serializes_round_trip() {
let record = EmbedRecord::new(
mem("mem_01ARZ3NDEKTSV4RRFFQ69G5FAV"),
"stub:v1",
vec![0.1f32, -0.2, 0.3, 0.4, -0.5],
at(7),
)
.expect("valid record");
let mut bytes = Vec::with_capacity(record.vector.len() * 4);
for v in &record.vector {
bytes.extend_from_slice(&v.to_le_bytes());
}
assert_eq!(bytes.len(), record.vector.len() * 4);
let mut decoded = Vec::with_capacity(record.vector.len());
for chunk in bytes.chunks_exact(4) {
let arr = <[u8; 4]>::try_from(chunk).expect("chunk_exact yields four bytes");
decoded.push(f32::from_le_bytes(arr));
}
assert_eq!(decoded, record.vector);
assert_eq!(record.dim as usize, record.vector.len());
let rebuilt = EmbedRecord::new(
record.memory_id,
record.backend_id.clone(),
decoded,
record.computed_at,
)
.expect("rebuilt record");
assert_eq!(rebuilt, record);
}
#[test]
fn embed_record_new_rejects_oversized_dim() {
let rec = EmbedRecord::new(
mem("mem_01ARZ3NDEKTSV4RRFFQ69G5FAV"),
"stub:v1",
vec![1.0f32; 8],
at(0),
)
.expect("eight-dim record");
assert_eq!(rec.dim, 8);
assert_eq!(rec.vector.len(), 8);
}
}
pub mod local_stub;
pub use local_stub::{LocalStubEmbedder, STUB_BACKEND_ID, STUB_DIM};
pub mod ollama;
pub use ollama::{
OllamaEmbedder, DEFAULT_OLLAMA_EMBED_MODEL, DEFAULT_OLLAMA_ENDPOINT, NOMIC_EMBED_DIM,
OLLAMA_BACKEND_ID_PREFIX,
};