aimem_core/
embedder.rs

1//! Embedder — text-to-vector abstraction and implementations.
2//!
3//! Supports local (fastembed) and remote (Gemini 2.0) embedding generation.
4
5use std::sync::{Arc, Mutex};
6
7use async_trait::async_trait;
8use fastembed::{EmbeddingModel, TextEmbedding, TextInitOptions};
9use serde::{Deserialize, Serialize};
10use thiserror::Error;
11use tracing::warn;
12
13use crate::types::ContentPart;
14
15pub const LOCAL_EMBED_PROVIDER: &str = "fastembed";
16pub const LOCAL_EMBED_MODEL: &str = "all-MiniLM-L6-v2";
17pub const GEMINI_EMBED_PROVIDER: &str = "gemini";
18pub const GEMINI_EMBED_MODEL: &str = "gemini-embedding-2-preview";
19
20#[derive(Debug, Error)]
21pub enum EmbedError {
22    #[error("fastembed error: {0}")]
23    Fastembed(#[from] anyhow::Error),
24    #[error("mutex poisoned")]
25    Poison,
26    #[error("remote api error: {0}")]
27    Remote(String),
28    #[error("serialization error: {0}")]
29    Serde(#[from] serde_json::Error),
30}
31
32/// Core trait for generating multimodal embeddings.
33#[async_trait]
34pub trait Embedder: Send + Sync {
35    /// Embed a slice of multimodal content, returning one `Vec<f32>` per input.
36    async fn embed(&self, inputs: &[Vec<ContentPart>]) -> Result<Vec<Vec<f32>>, EmbedError>;
37
38    /// Convenience: embed a single string.
39    async fn embed_one(&self, text: &str) -> Result<Vec<f32>, EmbedError> {
40        let mut results = self.embed(&[vec![ContentPart::text(text)]]).await?;
41        results
42            .pop()
43            .ok_or_else(|| EmbedError::Remote("empty embedding result".to_string()))
44    }
45
46    /// Dimension of the vectors produced by this embedder.
47    fn dimension(&self) -> usize;
48
49    /// Provider identifier used for store compatibility checks.
50    fn provider_name(&self) -> &str;
51
52    /// Model identifier used for store compatibility checks.
53    fn model_name(&self) -> &str;
54}
55
56/// Local embedder using fastembed-rs.
57/// Currently only uses the text parts of multimodal input.
58#[derive(Clone)]
59pub struct LocalEmbedder {
60    inner: Arc<Mutex<TextEmbedding>>,
61}
62
63impl std::fmt::Debug for LocalEmbedder {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        f.debug_struct("LocalEmbedder")
66            .field("model", &LOCAL_EMBED_MODEL)
67            .finish()
68    }
69}
70
71impl LocalEmbedder {
72    /// Create a new local embedder with the default model (all-MiniLM-L6-v2, 384 dims).
73    pub fn new() -> Result<Self, EmbedError> {
74        let model = TextEmbedding::try_new(
75            TextInitOptions::new(EmbeddingModel::AllMiniLML6V2).with_show_download_progress(true),
76        )?;
77        Ok(Self {
78            inner: Arc::new(Mutex::new(model)),
79        })
80    }
81}
82
83#[async_trait]
84impl Embedder for LocalEmbedder {
85    async fn embed(&self, inputs: &[Vec<ContentPart>]) -> Result<Vec<Vec<f32>>, EmbedError> {
86        // Local embedder only supports text. Flatten multimodal parts into a single string.
87        let texts: Vec<String> = inputs
88            .iter()
89            .map(|parts| {
90                parts
91                    .iter()
92                    .filter_map(|p| p.as_text())
93                    .collect::<Vec<_>>()
94                    .join(" ")
95            })
96            .collect();
97
98        let text_refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
99
100        let mut guard = self.inner.lock().map_err(|_| EmbedError::Poison)?;
101        let results = guard.embed(text_refs, None)?;
102        Ok(results)
103    }
104
105    fn dimension(&self) -> usize {
106        384
107    }
108
109    fn provider_name(&self) -> &str {
110        LOCAL_EMBED_PROVIDER
111    }
112
113    fn model_name(&self) -> &str {
114        LOCAL_EMBED_MODEL
115    }
116}
117
118/// Remote embedder using Google Gemini Embedding 2 API.
119pub struct Gemini2Embedder {
120    api_key: String,
121    model: String,
122    client: reqwest::Client,
123}
124
125impl Gemini2Embedder {
126    pub fn new(api_key: String) -> Self {
127        Self {
128            api_key,
129            model: GEMINI_EMBED_MODEL.to_string(),
130            client: reqwest::Client::new(),
131        }
132    }
133}
134
135#[derive(Serialize)]
136struct GeminiEmbedRequest {
137    requests: Vec<EmbedContentRequest>,
138}
139
140#[derive(Serialize)]
141struct EmbedContentRequest {
142    model: String,
143    content: GeminiContent,
144}
145
146#[derive(Serialize)]
147struct GeminiContent {
148    parts: Vec<GeminiPart>,
149}
150
151#[derive(Serialize)]
152#[serde(rename_all = "camelCase")]
153enum GeminiPart {
154    Text(String),
155    InlineData { mime_type: String, data: String },
156}
157
158#[derive(Deserialize)]
159struct GeminiBatchEmbedResponse {
160    embeddings: Vec<GeminiEmbedding>,
161}
162
163#[derive(Deserialize)]
164struct GeminiEmbedding {
165    values: Vec<f32>,
166}
167
168#[async_trait]
169impl Embedder for Gemini2Embedder {
170    async fn embed(&self, inputs: &[Vec<ContentPart>]) -> Result<Vec<Vec<f32>>, EmbedError> {
171        let url = format!(
172            "https://generativelanguage.googleapis.com/v1beta/models/{}:batchEmbedContents?key={}",
173            self.model, self.api_key
174        );
175
176        let mut requests = Vec::new();
177        for parts in inputs {
178            let mut gemini_parts = Vec::new();
179            for part in parts {
180                match part {
181                    ContentPart::Text { text } => {
182                        gemini_parts.push(GeminiPart::Text(text.clone()));
183                    }
184                    ContentPart::Image { uri, mime, data }
185                    | ContentPart::Audio { uri, mime, data }
186                    | ContentPart::Video { uri, mime, data } => {
187                        if let Some(bytes) = data {
188                            // Explicitly provided bytes (Recommended)
189                            let encoded = base64::Engine::encode(
190                                &base64::engine::general_purpose::STANDARD,
191                                bytes,
192                            );
193                            gemini_parts.push(GeminiPart::InlineData {
194                                mime_type: mime.clone(),
195                                data: encoded,
196                            });
197                        } else if let Some(uri_str) = uri {
198                            if uri_str.starts_with("data:") {
199                                // Explicitly provided data URI
200                                if let Some(comma_pos) = uri_str.find(',') {
201                                    let data_part = uri_str[comma_pos + 1..].to_string();
202                                    gemini_parts.push(GeminiPart::InlineData {
203                                        mime_type: mime.clone(),
204                                        data: data_part,
205                                    });
206                                }
207                            } else {
208                                // We NO LONGER read local files automatically.
209                                // If the caller wants to embed a file, they must read it into bytes first.
210                                warn!(
211                                    "Skipping remote embedding for URI-only content part: {}. Provide explicit data bytes for remote embedding.",
212                                    uri_str
213                                );
214                            }
215                        }
216                    }
217                }
218            }
219
220            if gemini_parts.is_empty() {
221                return Err(EmbedError::Remote(
222                    "remote embedding requires explicit text, data URI, or raw bytes; uri-only parts are not uploaded automatically".to_string(),
223                ));
224            }
225
226            requests.push(EmbedContentRequest {
227                model: format!("models/{}", self.model),
228                content: GeminiContent {
229                    parts: gemini_parts,
230                },
231            });
232        }
233
234        let req_body = GeminiEmbedRequest { requests };
235        let resp = self
236            .client
237            .post(&url)
238            .json(&req_body)
239            .send()
240            .await
241            .map_err(|e| EmbedError::Remote(e.to_string()))?;
242
243        if !resp.status().is_success() {
244            let status = resp.status();
245            let err_text = resp
246                .text()
247                .await
248                .unwrap_or_else(|_| "Unknown error".to_string());
249            return Err(EmbedError::Remote(format!(
250                "Gemini API error ({status}): {err_text}"
251            )));
252        }
253
254        let resp_body: GeminiBatchEmbedResponse = resp
255            .json()
256            .await
257            .map_err(|e| EmbedError::Remote(e.to_string()))?;
258
259        Ok(resp_body.embeddings.into_iter().map(|e| e.values).collect())
260    }
261
262    fn dimension(&self) -> usize {
263        768
264    }
265
266    fn provider_name(&self) -> &str {
267        GEMINI_EMBED_PROVIDER
268    }
269
270    fn model_name(&self) -> &str {
271        &self.model
272    }
273}
274
275/// Helper to serialize a `Vec<f32>` to a JSON array string for Turso vector32.
276pub fn to_vector32_json(v: &[f32]) -> String {
277    let nums: Vec<String> = v.iter().map(|f| f.to_string()).collect();
278    format!("[{}]", nums.join(","))
279}
aimem_core/embedder.rs

aimem_core/
embedder.rs
