Skip to main content

aimux_core/
math.rs

1//! Math helpers ported from the Vercel AI SDK TypeScript `util` package.
2
3use thiserror::Error;
4
5/// Errors raised by the util helpers.
6#[derive(Debug, Clone, PartialEq, Eq, Error)]
7pub enum UtilError {
8    /// Vectors passed to [`cosine_similarity`] did not have the same length.
9    #[error("Vectors must have the same length")]
10    VectorLengthMismatch,
11}
12
13/// Cosine similarity between two vectors.
14///
15/// Port of `cosineSimilarity` in `packages/ai/src/util/cosine-similarity.ts`.
16///
17/// Returns `Ok(0.0)` for empty vectors or when either vector is the zero
18/// vector, and `Err(UtilError::VectorLengthMismatch)` when the lengths differ
19/// (the TS version throws an `InvalidArgumentError`).
20///
21/// # Errors
22///
23/// Returns `UtilError::VectorLengthMismatch` when the two vectors have
24/// different lengths.
25pub fn cosine_similarity(vector1: &[f64], vector2: &[f64]) -> Result<f64, UtilError> {
26    if vector1.len() != vector2.len() {
27        return Err(UtilError::VectorLengthMismatch);
28    }
29
30    let n = vector1.len();
31    if n == 0 {
32        return Ok(0.0);
33    }
34
35    let mut magnitude_squared_1 = 0.0;
36    let mut magnitude_squared_2 = 0.0;
37    let mut dot_product = 0.0;
38
39    for i in 0..n {
40        let value1 = vector1[i];
41        let value2 = vector2[i];
42        magnitude_squared_1 += value1 * value1;
43        magnitude_squared_2 += value2 * value2;
44        dot_product += value1 * value2;
45    }
46
47    if magnitude_squared_1 == 0.0 || magnitude_squared_2 == 0.0 {
48        Ok(0.0)
49    } else {
50        Ok(dot_product / (magnitude_squared_1.sqrt() * magnitude_squared_2.sqrt()))
51    }
52}