apollo/
embeddings.rs

1//! Vector embeddings for semantic memory search
2//! Uses Gemini text-embedding-004 API (free tier)
3
4use anyhow::Result;
5use serde::{Deserialize, Serialize};
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct Embedding {
9    pub text: String,
10    pub vector: Vec<f32>,
11    pub metadata: serde_json::Value,
12}
13
14#[derive(Deserialize)]
15struct GeminiResponse {
16    embedding: GeminiEmbedding,
17}
18
19#[derive(Deserialize)]
20struct GeminiEmbedding {
21    values: Vec<f32>,
22}
23
24pub struct EmbeddingsClient {
25    api_key: String,
26    client: reqwest::Client,
27}
28
29impl EmbeddingsClient {
30    pub fn new(api_key: String) -> Self {
31        Self {
32            api_key,
33            client: crate::http::shared(),
34        }
35    }
36
37    /// Embed text using Gemini text-embedding-004
38    pub async fn embed(&self, text: &str) -> Result<Vec<f32>> {
39        let url = "https://generativelanguage.googleapis.com/v1beta/models/text-embedding-004:embedContent";
40
41        let request = serde_json::json!({
42            "model": "models/text-embedding-004",
43            "content": {
44                "parts": [{
45                    "text": text
46                }]
47            }
48        });
49
50        let resp = self
51            .client
52            .post(url)
53            .header("x-goog-api-key", &self.api_key)
54            .json(&request)
55            .send()
56            .await?
57            .error_for_status()?;
58
59        let result: GeminiResponse = resp.json().await?;
60        Ok(result.embedding.values)
61    }
62
63    /// Embed multiple texts in batch
64    pub async fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
65        let mut results = Vec::new();
66        for text in texts {
67            results.push(self.embed(text).await?);
68        }
69        Ok(results)
70    }
71}
72
73/// Cosine similarity between two vectors, or `None` when they are not
74/// comparable.
75///
76/// Vectors of different lengths come from different embedding models and have
77/// no meaningful similarity, so no score is produced for them.
78pub fn cosine_similarity(a: &[f32], b: &[f32]) -> Option<f32> {
79    if a.len() != b.len() || a.is_empty() {
80        return None;
81    }
82
83    let dot_product: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
84    let magnitude_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
85    let magnitude_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
86
87    if magnitude_a == 0.0 || magnitude_b == 0.0 {
88        return Some(0.0);
89    }
90
91    Some(dot_product / (magnitude_a * magnitude_b))
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn test_cosine_similarity() {
100        let a = vec![1.0, 0.0, 0.0];
101        let b = vec![1.0, 0.0, 0.0];
102        assert!((cosine_similarity(&a, &b).unwrap() - 1.0).abs() < 1e-6);
103
104        let a = vec![1.0, 0.0];
105        let b = vec![0.0, 1.0];
106        assert!((cosine_similarity(&a, &b).unwrap() - 0.0).abs() < 1e-6);
107    }
108
109    #[test]
110    fn test_cosine_similarity_edge_cases() {
111        // Empty vectors
112        let a: Vec<f32> = vec![];
113        let b: Vec<f32> = vec![];
114        assert_eq!(cosine_similarity(&a, &b), None);
115
116        // Mismatched lengths
117        let a = vec![1.0];
118        let b = vec![1.0, 0.0];
119        assert_eq!(cosine_similarity(&a, &b), None);
120
121        // Zero vectors
122        let a = vec![0.0, 0.0];
123        let b = vec![0.0, 0.0];
124        assert_eq!(cosine_similarity(&a, &b), Some(0.0));
125    }
126}
apollo/embeddings.rs

apollo/
embeddings.rs
