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: reqwest::Client::new(),
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
74pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
75    if a.len() != b.len() || a.is_empty() {
76        return 0.0;
77    }
78
79    let dot_product: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
80    let magnitude_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
81    let magnitude_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
82
83    if magnitude_a == 0.0 || magnitude_b == 0.0 {
84        return 0.0;
85    }
86
87    dot_product / (magnitude_a * magnitude_b)
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn test_cosine_similarity() {
96        let a = vec![1.0, 0.0, 0.0];
97        let b = vec![1.0, 0.0, 0.0];
98        assert!((cosine_similarity(&a, &b) - 1.0).abs() < 1e-6);
99
100        let a = vec![1.0, 0.0];
101        let b = vec![0.0, 1.0];
102        assert!((cosine_similarity(&a, &b) - 0.0).abs() < 1e-6);
103    }
104
105    #[test]
106    fn test_cosine_similarity_edge_cases() {
107        // Empty vectors
108        let a: Vec<f32> = vec![];
109        let b: Vec<f32> = vec![];
110        assert!((cosine_similarity(&a, &b) - 0.0).abs() < 1e-6);
111
112        // Mismatched lengths
113        let a = vec![1.0];
114        let b = vec![1.0, 0.0];
115        assert!((cosine_similarity(&a, &b) - 0.0).abs() < 1e-6);
116
117        // Zero vectors
118        let a = vec![0.0, 0.0];
119        let b = vec![0.0, 0.0];
120        assert!((cosine_similarity(&a, &b) - 0.0).abs() < 1e-6);
121    }
122}
apollo/embeddings.rs

apollo/
embeddings.rs
