Skip to main content

context_harness_core/
embedding.rs

1//! Embedding provider trait and vector utilities.
2//!
3//! Defines the [`EmbeddingProvider`] trait that all embedding backends
4//! implement, plus pure helper functions for vector serialization and
5//! similarity computation.
6//!
7//! Concrete provider implementations (OpenAI, Ollama, fastembed, tract)
8//! live in the `context-harness` app crate.
9
10/// Trait for embedding providers.
11///
12/// Defines the interface that all embedding backends must implement.
13/// Implementations are created by the application and passed to core
14/// functions that need embedding metadata.
15pub trait EmbeddingProvider: Send + Sync {
16    /// Returns the model identifier (e.g. `"text-embedding-3-small"`).
17    fn model_name(&self) -> &str;
18    /// Returns the embedding vector dimensionality (e.g. `1536`).
19    fn dims(&self) -> usize;
20}
21
22/// Encode a float vector as a BLOB (little-endian f32 bytes).
23///
24/// Each `f32` is stored as 4 bytes in little-endian order, producing
25/// a BLOB of `vec.len() × 4` bytes.
26///
27/// # Example
28///
29/// ```rust
30/// use context_harness_core::embedding::{vec_to_blob, blob_to_vec};
31///
32/// let v = vec![1.0f32, -2.5, 3.125];
33/// let blob = vec_to_blob(&v);
34/// assert_eq!(blob.len(), 12); // 3 × 4 bytes
35/// assert_eq!(blob_to_vec(&blob), v);
36/// ```
37pub fn vec_to_blob(vec: &[f32]) -> Vec<u8> {
38    let mut bytes = Vec::with_capacity(vec.len() * 4);
39    for &v in vec {
40        bytes.extend_from_slice(&v.to_le_bytes());
41    }
42    bytes
43}
44
45/// Decode a BLOB back into a float vector.
46///
47/// Reverses [`vec_to_blob`]: reads 4-byte little-endian `f32` values
48/// from the byte slice.
49pub fn blob_to_vec(blob: &[u8]) -> Vec<f32> {
50    blob.as_chunks::<4>()
51        .0
52        .iter()
53        .map(|chunk| f32::from_le_bytes(*chunk))
54        .collect()
55}
56
57/// Compute cosine similarity between two embedding vectors.
58///
59/// Returns a value in `[-1.0, 1.0]`:
60/// - `1.0` = identical direction
61/// - `0.0` = orthogonal (unrelated)
62/// - `-1.0` = opposite direction
63///
64/// Returns `0.0` for empty vectors or vectors of different lengths.
65pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
66    if a.len() != b.len() || a.is_empty() {
67        return 0.0;
68    }
69
70    let mut dot = 0.0f32;
71    let mut norm_a = 0.0f32;
72    let mut norm_b = 0.0f32;
73
74    for (x, y) in a.iter().zip(b.iter()) {
75        dot += x * y;
76        norm_a += x * x;
77        norm_b += y * y;
78    }
79
80    let denom = norm_a.sqrt() * norm_b.sqrt();
81    if denom < f32::EPSILON {
82        return 0.0;
83    }
84
85    dot / denom
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn test_vec_blob_roundtrip() {
94        let vec = vec![1.0f32, -2.5, 3.125, 0.0, -0.001];
95        let blob = vec_to_blob(&vec);
96        let restored = blob_to_vec(&blob);
97        assert_eq!(vec, restored);
98    }
99
100    #[test]
101    fn test_cosine_identical() {
102        let v = vec![1.0, 2.0, 3.0];
103        let sim = cosine_similarity(&v, &v);
104        assert!((sim - 1.0).abs() < 1e-6);
105    }
106
107    #[test]
108    fn test_cosine_orthogonal() {
109        let a = vec![1.0, 0.0, 0.0];
110        let b = vec![0.0, 1.0, 0.0];
111        let sim = cosine_similarity(&a, &b);
112        assert!(sim.abs() < 1e-6);
113    }
114
115    #[test]
116    fn test_cosine_opposite() {
117        let a = vec![1.0, 0.0];
118        let b = vec![-1.0, 0.0];
119        let sim = cosine_similarity(&a, &b);
120        assert!((sim + 1.0).abs() < 1e-6);
121    }
122
123    #[test]
124    fn test_cosine_empty() {
125        let sim = cosine_similarity(&[], &[]);
126        assert_eq!(sim, 0.0);
127    }
128
129    #[test]
130    fn test_cosine_different_lengths() {
131        let a = vec![1.0, 2.0];
132        let b = vec![1.0];
133        let sim = cosine_similarity(&a, &b);
134        assert_eq!(sim, 0.0);
135    }
136}