do_memory_core/embeddings/
similarity.rs1use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct SimilaritySearchResult<T> {
8 pub item: T,
10 pub similarity: f32,
12 pub metadata: SimilarityMetadata,
14}
15
16#[derive(Debug, Clone, Serialize, Deserialize, Default)]
18pub struct SimilarityMetadata {
19 #[serde(default)]
21 pub embedding_model: String,
22 pub embedding_timestamp: Option<chrono::DateTime<chrono::Utc>>,
24 #[serde(default)]
26 pub context: serde_json::Value,
27}
28
29#[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 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 let similarity = dot_product / (norm_a_sq.sqrt() * norm_b_sq.sqrt());
65
66 (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}