Skip to main content

do_memory_core/embeddings/
similarity.rs

1//! Vector similarity calculations and search utilities
2
3use serde::{Deserialize, Serialize};
4
5/// Result from similarity search containing the item and similarity score
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct SimilaritySearchResult<T> {
8    /// The found item (episode or pattern)
9    pub item: T,
10    /// Similarity score (0.0 to 1.0, higher = more similar)
11    pub similarity: f32,
12    /// Additional metadata about the match
13    pub metadata: SimilarityMetadata,
14}
15
16/// Metadata about a similarity match
17#[derive(Debug, Clone, Serialize, Deserialize, Default)]
18pub struct SimilarityMetadata {
19    /// Which embedding was used for the match
20    #[serde(default)]
21    pub embedding_model: String,
22    /// Timestamp of when the embedding was generated
23    pub embedding_timestamp: Option<chrono::DateTime<chrono::Utc>>,
24    /// Additional context about the match
25    #[serde(default)]
26    pub context: serde_json::Value,
27}
28
29/// Calculate cosine similarity between two vectors
30///
31/// Cosine similarity measures the cosine of the angle between two vectors,
32/// giving a similarity score between -1 and 1 (normalized to 0-1 for convenience).
33/// Higher scores indicate greater similarity.
34///
35/// # Optimization:
36/// 1. Uses f32 accumulators to enable LLVM SIMD autovectorization (AVX/SSE/NEON).
37/// 2. Uses a single pass for dot-product and magnitude calculations.
38/// 3. Maintains dynamic range stability by using individual square roots.
39#[must_use]
40pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
41    if a.len() != b.len() || a.is_empty() {
42        return 0.0;
43    }
44
45    let mut dot_product = 0.0f32;
46    let mut norm_a_sq = 0.0f32;
47    let mut norm_b_sq = 0.0f32;
48
49    // Single pass with f32 enables SIMD autovectorization.
50    // f32 is sufficient for embedding similarity and provides a ~4-8x speedup over f64
51    // when properly vectorized by the compiler.
52    for (&x, &y) in a.iter().zip(b.iter()) {
53        dot_product += x * y;
54        norm_a_sq += x * x;
55        norm_b_sq += y * y;
56    }
57
58    if norm_a_sq <= 0.0 || norm_b_sq <= 0.0 {
59        return 0.0;
60    }
61
62    // We keep individual sqrt calls to maintain dynamic range stability and avoid
63    // potential product overflow/underflow before the sqrt is taken.
64    let similarity = dot_product / (norm_a_sq.sqrt() * norm_b_sq.sqrt());
65
66    // Normalize from [-1, 1] to [0, 1] range for semantic scores
67    (similarity + 1.0) / 2.0
68}
69#[cfg(test)]
70mod tests {
71    use super::*;
72
73    #[test]
74    fn test_cosine_similarity() {
75        let vec1 = vec![1.0, 2.0, 3.0];
76        let vec2 = vec![1.0, 2.0, 3.0];
77        let similarity = cosine_similarity(&vec1, &vec2);
78        assert!((similarity - 1.0).abs() < 0.001);
79
80        let vec3 = vec![1.0, 0.0];
81        let vec4 = vec![0.0, 1.0];
82        let similarity = cosine_similarity(&vec3, &vec4);
83        assert!((similarity - 0.5).abs() < 0.001);
84
85        let vec5 = vec![1.0, 2.0, 3.0];
86        let vec6 = vec![-1.0, -2.0, -3.0];
87        let similarity = cosine_similarity(&vec5, &vec6);
88        assert!((similarity - 0.0).abs() < 0.001);
89
90        let vec7 = vec![1.0, 2.0];
91        let vec8 = vec![1.0, 2.0, 3.0];
92        let similarity = cosine_similarity(&vec7, &vec8);
93        assert_eq!(similarity, 0.0);
94    }
95
96    #[test]
97    fn test_cosine_similarity_empty() {
98        let vec1: Vec<f32> = vec![];
99        let vec2: Vec<f32> = vec![];
100        assert_eq!(cosine_similarity(&vec1, &vec2), 0.0);
101
102        let vec3 = vec![1.0, 2.0, 3.0];
103        assert_eq!(cosine_similarity(&vec1, &vec3), 0.0);
104        assert_eq!(cosine_similarity(&vec3, &vec1), 0.0);
105    }
106
107    #[test]
108    fn test_cosine_similarity_zero_magnitude() {
109        let vec1 = vec![0.0, 0.0, 0.0];
110        let vec2 = vec![1.0, 2.0, 3.0];
111        assert_eq!(cosine_similarity(&vec1, &vec2), 0.0);
112
113        let vec3 = vec![1.0, 2.0, 3.0];
114        let vec4 = vec![0.0, 0.0, 0.0];
115        assert_eq!(cosine_similarity(&vec3, &vec4), 0.0);
116
117        let vec5 = vec![0.0, 0.0, 0.0];
118        let vec6 = vec![0.0, 0.0, 0.0];
119        assert_eq!(cosine_similarity(&vec5, &vec6), 0.0);
120    }
121}