rectilinear_core/embedding/
mod.rs

1#[cfg(feature = "local-embeddings")]
2mod local;
3
4use anyhow::{Context, Result};
5use reqwest::StatusCode;
6use serde::Deserialize;
7use sha2::{Digest, Sha256};
8
9use crate::config::{Config, EmbeddingBackend};
10
11const GEMINI_EMBEDDING_MODEL: &str = "gemini-embedding-2-preview";
12
13enum Backend {
14    Gemini(GeminiBackend),
15    #[cfg(feature = "local-embeddings")]
16    Local(local::LocalBackend),
17}
18
19pub struct Embedder {
20    backend: Backend,
21    dimensions: usize,
22}
23
24/// Stable fingerprint of the exact issue fields passed to `chunk_text`.
25pub fn issue_content_hash(title: &str, description: Option<&str>) -> String {
26    let mut hasher = Sha256::new();
27    hasher.update(title.as_bytes());
28    hasher.update([0]);
29    hasher.update(description.unwrap_or_default().as_bytes());
30    hex::encode(hasher.finalize())
31}
32
33// --- Gemini API backend ---
34
35struct GeminiBackend {
36    client: reqwest::Client,
37    api_key: String,
38}
39
40#[derive(Deserialize)]
41struct GeminiModelsResponse {
42    models: Vec<GeminiModel>,
43}
44
45#[derive(Deserialize)]
46struct GeminiModel {
47    #[allow(dead_code)]
48    name: String,
49}
50
51#[derive(Deserialize)]
52struct GeminiErrorEnvelope {
53    error: GeminiErrorBody,
54}
55
56#[derive(Deserialize)]
57struct GeminiErrorBody {
58    #[allow(dead_code)]
59    code: Option<u16>,
60    message: Option<String>,
61    status: Option<String>,
62}
63
64impl GeminiBackend {
65    fn new(api_key: &str) -> Self {
66        Self {
67            client: reqwest::Client::new(),
68            api_key: api_key.to_string(),
69        }
70    }
71
72    fn with_http_client(client: reqwest::Client, api_key: &str) -> Self {
73        Self {
74            client,
75            api_key: api_key.to_string(),
76        }
77    }
78
79    async fn test_api_key(&self) -> Result<()> {
80        let resp = self
81            .client
82            .get("https://generativelanguage.googleapis.com/v1beta/models?pageSize=1")
83            .header("x-goog-api-key", &self.api_key)
84            .send()
85            .await
86            .context("Failed to call Gemini models API")?;
87
88        let status = resp.status();
89        let body = resp
90            .text()
91            .await
92            .context("Failed to read Gemini models response")?;
93
94        if !status.is_success() {
95            anyhow::bail!("{}", summarize_gemini_error(status, &body));
96        }
97
98        let response: GeminiModelsResponse =
99            serde_json::from_str(&body).context("Failed to parse Gemini models response")?;
100        if response.models.is_empty() {
101            anyhow::bail!("Gemini returned no models");
102        }
103
104        Ok(())
105    }
106
107    async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
108        let url = format!(
109            "https://generativelanguage.googleapis.com/v1beta/models/{GEMINI_EMBEDDING_MODEL}:batchEmbedContents?key={}",
110            self.api_key,
111        );
112
113        let mut all_embeddings = Vec::new();
114        for batch in texts.chunks(100) {
115            let requests: Vec<_> = batch
116                .iter()
117                .map(|text| {
118                    serde_json::json!({
119                        "model": format!("models/{GEMINI_EMBEDDING_MODEL}"),
120                        "content": {
121                            "parts": [{"text": text}]
122                        },
123                        "outputDimensionality": 768
124                    })
125                })
126                .collect();
127
128            let body = serde_json::json!({ "requests": requests });
129
130            let resp = self
131                .client
132                .post(&url)
133                .json(&body)
134                .send()
135                .await
136                .context("Failed to call Gemini embedding API")?;
137
138            let status = resp.status();
139            if !status.is_success() {
140                let text = resp.text().await.unwrap_or_default();
141                anyhow::bail!("Gemini API returned {}: {}", status, text);
142            }
143
144            let data: serde_json::Value = resp.json().await?;
145            let embeddings = data["embeddings"]
146                .as_array()
147                .context("No embeddings in response")?;
148
149            for emb in embeddings {
150                let values: Vec<f32> = emb["values"]
151                    .as_array()
152                    .context("No values in embedding")?
153                    .iter()
154                    .map(|v| v.as_f64().unwrap_or(0.0) as f32)
155                    .collect();
156                all_embeddings.push(values);
157            }
158        }
159
160        Ok(all_embeddings)
161    }
162}
163
164fn summarize_gemini_error(status: StatusCode, body: &str) -> String {
165    if let Ok(error) = serde_json::from_str::<GeminiErrorEnvelope>(body) {
166        return compact_gemini_error(
167            status,
168            error.error.message.as_deref(),
169            error.error.status.as_deref(),
170        );
171    }
172
173    compact_gemini_error(status, None, None)
174}
175
176fn compact_gemini_error(
177    status: StatusCode,
178    message: Option<&str>,
179    api_status: Option<&str>,
180) -> String {
181    let normalized = message.unwrap_or("").trim().to_lowercase();
182
183    if normalized.contains("api key not valid") || normalized.contains("invalid api key") {
184        return "Invalid API key".into();
185    }
186
187    if normalized.contains("reported as leaked") || normalized.contains("disabled") {
188        return "Key blocked".into();
189    }
190
191    if normalized.contains("billing")
192        || matches!(api_status, Some("FAILED_PRECONDITION" | "SERVICE_DISABLED"))
193    {
194        return "Setup required".into();
195    }
196
197    if status == StatusCode::UNAUTHORIZED || matches!(api_status, Some("UNAUTHENTICATED")) {
198        return "Unauthorized".into();
199    }
200
201    if status == StatusCode::FORBIDDEN || matches!(api_status, Some("PERMISSION_DENIED")) {
202        return "Access denied".into();
203    }
204
205    if status == StatusCode::TOO_MANY_REQUESTS || matches!(api_status, Some("RESOURCE_EXHAUSTED")) {
206        return "Rate limited".into();
207    }
208
209    if status.is_server_error() {
210        return "Gemini unavailable".into();
211    }
212
213    if let Some(message) = message {
214        let trimmed = message.trim();
215        if !trimmed.is_empty() && trimmed.len() <= 48 {
216            return trimmed.to_string();
217        }
218    }
219
220    status
221        .canonical_reason()
222        .unwrap_or("Request failed")
223        .to_string()
224}
225
226// --- Embedder (main interface) ---
227
228impl Embedder {
229    pub fn new(config: &Config) -> Result<Self> {
230        let gemini_key = std::env::var("GEMINI_API_KEY")
231            .ok()
232            .or_else(|| config.embedding.gemini_api_key.clone());
233
234        match config.embedding.backend {
235            EmbeddingBackend::Api => {
236                let key = gemini_key.context(
237                    "Gemini API key required for API backend. Set GEMINI_API_KEY or configure in config.",
238                )?;
239                Self::new_api(&key)
240            }
241            #[cfg(feature = "local-embeddings")]
242            EmbeddingBackend::Local => {
243                let backend = local::LocalBackend::new(config)?;
244                let dimensions = backend.dimensions();
245                Ok(Self {
246                    dimensions,
247                    backend: Backend::Local(backend),
248                })
249            }
250            #[cfg(not(feature = "local-embeddings"))]
251            EmbeddingBackend::Local => {
252                anyhow::bail!(
253                    "Local embeddings not available — compile with `local-embeddings` feature"
254                )
255            }
256        }
257    }
258
259    /// Create an embedder using the Gemini API backend.
260    pub fn new_api(api_key: &str) -> Result<Self> {
261        Ok(Self {
262            dimensions: 768,
263            backend: Backend::Gemini(GeminiBackend::new(api_key)),
264        })
265    }
266
267    /// Create an embedder using the Gemini API backend with a pre-built HTTP client.
268    pub fn new_api_with_http_client(client: reqwest::Client, api_key: &str) -> Result<Self> {
269        Ok(Self {
270            dimensions: 768,
271            backend: Backend::Gemini(GeminiBackend::with_http_client(client, api_key)),
272        })
273    }
274
275    pub async fn test_api_key(&self) -> Result<()> {
276        match &self.backend {
277            Backend::Gemini(b) => b.test_api_key().await,
278            #[cfg(feature = "local-embeddings")]
279            Backend::Local(_) => anyhow::bail!("Gemini API key not in use"),
280        }
281    }
282
283    /// Create an embedder using the local GGUF backend.
284    #[cfg(feature = "local-embeddings")]
285    pub fn new_local(_models_dir: &std::path::Path) -> Result<Self> {
286        // TODO: pass models_dir through to LocalBackend instead of using Config default
287        let config = Config {
288            embedding: crate::config::EmbeddingConfig {
289                backend: EmbeddingBackend::Local,
290                gemini_api_key: None,
291            },
292            ..Config::default()
293        };
294        let backend = local::LocalBackend::new(&config)?;
295        let dimensions = backend.dimensions();
296        Ok(Self {
297            dimensions,
298            backend: Backend::Local(backend),
299        })
300    }
301
302    pub async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
303        match &self.backend {
304            Backend::Gemini(b) => b.embed_batch(texts).await,
305            #[cfg(feature = "local-embeddings")]
306            Backend::Local(b) => b.embed_batch(texts),
307        }
308    }
309
310    pub async fn embed_single(&self, text: &str) -> Result<Vec<f32>> {
311        let results: Vec<Vec<f32>> = self.embed_batch(&[text.to_string()]).await?;
312        results.into_iter().next().context("No embedding returned")
313    }
314
315    pub fn dimensions(&self) -> usize {
316        self.dimensions
317    }
318
319    pub fn backend_name(&self) -> &str {
320        match &self.backend {
321            Backend::Gemini(_) => GEMINI_EMBEDDING_MODEL,
322            #[cfg(feature = "local-embeddings")]
323            Backend::Local(_) => "embeddinggemma-300m-qat-q8_0",
324        }
325    }
326}
327
328// --- Text chunking ---
329
330/// Chunk text into segments of approximately `max_tokens` tokens with `overlap` token overlap.
331pub fn chunk_text(title: &str, text: &str, max_tokens: usize, overlap: usize) -> Vec<String> {
332    let prefix = format!("title: {}\n\n", title);
333
334    if text.is_empty() {
335        return vec![format!("{}(no description)", prefix)];
336    }
337
338    let max_chars = max_tokens * 4;
339    let overlap_chars = overlap * 4;
340
341    if text.len() <= max_chars {
342        return vec![format!("{}{}", prefix, text)];
343    }
344
345    // Snap a byte offset to the nearest char boundary (rounding down)
346    let floor_char = |s: &str, pos: usize| {
347        let pos = pos.min(s.len());
348        let mut i = pos;
349        while i > 0 && !s.is_char_boundary(i) {
350            i -= 1;
351        }
352        i
353    };
354
355    let mut chunks = Vec::new();
356    let mut start = 0;
357
358    while start < text.len() {
359        let end = floor_char(text, start + max_chars);
360
361        let chunk_slice = &text[start..end];
362        let break_at = if end < text.len() {
363            chunk_slice
364                .rfind("\n\n")
365                .or_else(|| chunk_slice.rfind('\n'))
366                .or_else(|| chunk_slice.rfind(". "))
367                .or_else(|| chunk_slice.rfind(' '))
368                .map(|p| start + p + 1)
369                .unwrap_or(end)
370        } else {
371            end
372        };
373
374        chunks.push(format!("{}{}", prefix, &text[start..break_at]));
375
376        if break_at >= text.len() {
377            break;
378        }
379
380        let new_start = floor_char(
381            text,
382            if break_at > overlap_chars {
383                break_at - overlap_chars
384            } else {
385                break_at
386            },
387        );
388        // Ensure forward progress — overlap must never push start backwards
389        start = if new_start <= start {
390            break_at
391        } else {
392            new_start
393        };
394    }
395
396    chunks
397}
398
399/// Convert f32 embedding to bytes for storage
400pub fn embedding_to_bytes(embedding: &[f32]) -> Vec<u8> {
401    embedding.iter().flat_map(|f| f.to_le_bytes()).collect()
402}
403
404/// Convert bytes back to f32 embedding
405pub fn bytes_to_embedding(bytes: &[u8]) -> Vec<f32> {
406    bytes
407        .chunks_exact(4)
408        .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
409        .collect()
410}
411
412/// Cosine similarity between two embeddings
413pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
414    if a.len() != b.len() || a.is_empty() {
415        return 0.0;
416    }
417
418    let mut dot = 0.0f32;
419    let mut norm_a = 0.0f32;
420    let mut norm_b = 0.0f32;
421
422    for (x, y) in a.iter().zip(b.iter()) {
423        dot += x * y;
424        norm_a += x * x;
425        norm_b += y * y;
426    }
427
428    let denom = norm_a.sqrt() * norm_b.sqrt();
429    if denom == 0.0 {
430        0.0
431    } else {
432        dot / denom
433    }
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439
440    #[test]
441    fn compacts_invalid_api_key_errors() {
442        let body = r#"{
443            "error": {
444                "code": 400,
445                "message": "API key not valid. Please pass a valid API key.",
446                "status": "INVALID_ARGUMENT"
447            }
448        }"#;
449
450        assert_eq!(
451            summarize_gemini_error(StatusCode::BAD_REQUEST, body),
452            "Invalid API key"
453        );
454    }
455
456    #[test]
457    fn compacts_rate_limit_errors() {
458        let body = r#"{
459            "error": {
460                "code": 429,
461                "message": "Quota exceeded.",
462                "status": "RESOURCE_EXHAUSTED"
463            }
464        }"#;
465
466        assert_eq!(
467            summarize_gemini_error(StatusCode::TOO_MANY_REQUESTS, body),
468            "Rate limited"
469        );
470    }
471
472    #[test]
473    fn falls_back_to_status_for_unknown_errors() {
474        assert_eq!(
475            summarize_gemini_error(StatusCode::SERVICE_UNAVAILABLE, "not-json"),
476            "Gemini unavailable"
477        );
478    }
479}
rectilinear_core/embedding/mod.rs

rectilinear_core/embedding/
mod.rs
