microclaw_engine/
embedding.rs

1use std::sync::Arc;
2
3use anyhow::{anyhow, Result};
4use async_trait::async_trait;
5use serde::{Deserialize, Serialize};
6
7use crate::config::Config;
8
9#[async_trait]
10pub trait EmbeddingProvider: Send + Sync {
11    async fn embed(&self, text: &str) -> Result<Vec<f32>>;
12    fn model(&self) -> &str;
13    fn dimension(&self) -> usize;
14}
15
16pub struct OpenAIEmbeddingProvider {
17    client: reqwest::Client,
18    base_url: String,
19    api_key: String,
20    model: String,
21    dim: usize,
22}
23
24pub struct OllamaEmbeddingProvider {
25    client: reqwest::Client,
26    base_url: String,
27    model: String,
28    dim: usize,
29}
30
31#[derive(Debug, Serialize)]
32struct OpenAIEmbeddingRequest<'a> {
33    model: &'a str,
34    input: &'a str,
35}
36
37#[derive(Debug, Deserialize)]
38struct OpenAIEmbeddingResponse {
39    data: Vec<OpenAIEmbeddingData>,
40}
41
42#[derive(Debug, Deserialize)]
43struct OpenAIEmbeddingData {
44    embedding: Vec<f32>,
45}
46
47#[derive(Debug, Serialize)]
48struct OllamaEmbeddingRequest<'a> {
49    model: &'a str,
50    prompt: &'a str,
51}
52
53#[derive(Debug, Deserialize)]
54struct OllamaEmbeddingResponse {
55    embedding: Vec<f32>,
56}
57
58#[cfg(feature = "sqlite-vec")]
59fn infer_default_dim(provider: &str, model: &str) -> usize {
60    match provider {
61        "openai" if model.contains("3-large") => 3072,
62        "openai" => 1536,
63        "ollama" => 1024,
64        _ => 1536,
65    }
66}
67
68#[async_trait]
69impl EmbeddingProvider for OpenAIEmbeddingProvider {
70    async fn embed(&self, text: &str) -> Result<Vec<f32>> {
71        let url = format!("{}/embeddings", self.base_url.trim_end_matches('/'));
72        let response = self
73            .client
74            .post(url)
75            .bearer_auth(&self.api_key)
76            .json(&OpenAIEmbeddingRequest {
77                model: &self.model,
78                input: text,
79            })
80            .send()
81            .await?;
82
83        if !response.status().is_success() {
84            let body = response.text().await.unwrap_or_default();
85            return Err(anyhow!("embedding request failed: {}", body));
86        }
87
88        let body: OpenAIEmbeddingResponse = response.json().await?;
89        let embedding = body
90            .data
91            .into_iter()
92            .next()
93            .ok_or_else(|| anyhow!("empty embedding response"))?
94            .embedding;
95        Ok(embedding)
96    }
97
98    fn model(&self) -> &str {
99        &self.model
100    }
101
102    fn dimension(&self) -> usize {
103        self.dim
104    }
105}
106
107#[async_trait]
108impl EmbeddingProvider for OllamaEmbeddingProvider {
109    async fn embed(&self, text: &str) -> Result<Vec<f32>> {
110        let url = format!("{}/api/embeddings", self.base_url.trim_end_matches('/'));
111        let response = self
112            .client
113            .post(url)
114            .json(&OllamaEmbeddingRequest {
115                model: &self.model,
116                prompt: text,
117            })
118            .send()
119            .await?;
120
121        if !response.status().is_success() {
122            let body = response.text().await.unwrap_or_default();
123            return Err(anyhow!("embedding request failed: {}", body));
124        }
125
126        let body: OllamaEmbeddingResponse = response.json().await?;
127        Ok(body.embedding)
128    }
129
130    fn model(&self) -> &str {
131        &self.model
132    }
133
134    fn dimension(&self) -> usize {
135        self.dim
136    }
137}
138
139pub fn create_provider(config: &Config) -> Option<Arc<dyn EmbeddingProvider>> {
140    #[cfg(not(feature = "sqlite-vec"))]
141    {
142        let _ = config;
143        None
144    }
145
146    #[cfg(feature = "sqlite-vec")]
147    {
148        let provider = config
149            .embedding_provider
150            .as_deref()
151            .unwrap_or("")
152            .trim()
153            .to_lowercase();
154        if provider.is_empty() {
155            return None;
156        }
157
158        let model = config
159            .embedding_model
160            .clone()
161            .unwrap_or_else(|| match provider.as_str() {
162                "openai" => "text-embedding-3-small".to_string(),
163                "ollama" => "nomic-embed-text".to_string(),
164                _ => "text-embedding-3-small".to_string(),
165            });
166        let dim = config
167            .embedding_dim
168            .unwrap_or_else(|| infer_default_dim(&provider, &model));
169        let client = reqwest::Client::new();
170
171        match provider.as_str() {
172            "openai" => {
173                let api_key = config.embedding_api_key.clone().unwrap_or_default();
174                if api_key.trim().is_empty() {
175                    return None;
176                }
177                let base_url = config
178                    .embedding_base_url
179                    .clone()
180                    .unwrap_or_else(|| "https://api.openai.com/v1".to_string());
181                Some(Arc::new(OpenAIEmbeddingProvider {
182                    client,
183                    base_url,
184                    api_key,
185                    model,
186                    dim,
187                }))
188            }
189            "ollama" => {
190                let base_url = config
191                    .embedding_base_url
192                    .clone()
193                    .unwrap_or_else(|| "http://127.0.0.1:11434".to_string());
194                Some(Arc::new(OllamaEmbeddingProvider {
195                    client,
196                    base_url,
197                    model,
198                    dim,
199                }))
200            }
201            _ => None,
202        }
203    }
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209    use crate::config::{Config, WorkingDirIsolation};
210
211    fn base_config() -> Config {
212        let mut cfg = Config::test_defaults();
213        cfg.working_dir_isolation = WorkingDirIsolation::Chat;
214        cfg
215    }
216
217    #[test]
218    fn test_create_provider_without_config_returns_none() {
219        let cfg = base_config();
220        assert!(create_provider(&cfg).is_none());
221    }
222
223    #[cfg(feature = "sqlite-vec")]
224    #[test]
225    fn test_create_provider_openai_when_configured() {
226        let mut cfg = base_config();
227        cfg.embedding_provider = Some("openai".into());
228        cfg.embedding_api_key = Some("sk-test".into());
229        cfg.embedding_model = Some("text-embedding-3-small".into());
230        cfg.embedding_dim = Some(1536);
231
232        let provider = create_provider(&cfg);
233        assert!(provider.is_some());
234        assert_eq!(
235            provider.as_ref().map(|p| p.model()),
236            Some("text-embedding-3-small")
237        );
238    }
239}
microclaw_engine/embedding.rs

microclaw_engine/
embedding.rs
