Skip to main content

ruve/
embedder.rs

1use reqwest::blocking::Client;
2use serde::Deserialize;
3
4pub enum EmbedderBackend {
5    OpenAI,
6    Ollama,
7}
8
9pub struct Embedder {
10    client: Client,
11    backend: EmbedderBackend,
12    api_key: Option<String>,
13}
14
15// OpenAI response shape
16#[derive(Deserialize)]
17struct OpenAIResponse {
18    data: Vec<OpenAIEmbedding>,
19}
20#[derive(Deserialize)]
21struct OpenAIEmbedding {
22    embedding: Vec<f32>,
23}
24
25// Ollama response shape
26#[derive(Deserialize)]
27struct OllamaResponse {
28    embeddings: Vec<Vec<f32>>,
29}
30
31impl Embedder {
32    pub fn openai() -> Self {
33        dotenvy::dotenv().ok();
34        let api_key = std::env::var("OPENAI_API_KEY")
35            .expect("OPENAI_API_KEY not set in environment or .env");
36        Embedder {
37            client: Client::new(),
38            backend: EmbedderBackend::OpenAI,
39            api_key: Some(api_key),
40        }
41    }
42
43    pub fn ollama() -> Self {
44        Embedder {
45            client: Client::new(),
46            backend: EmbedderBackend::Ollama,
47            api_key: None,
48        }
49    }
50
51    // keep existing new() pointing to openai so nothing else breaks
52    pub fn new() -> Self {
53        Self::openai()
54    }
55
56    pub fn embed(&self, text: &str) -> Vec<f32> {
57        match self.backend {
58            EmbedderBackend::OpenAI => {
59                let response: OpenAIResponse = self.client
60                    .post("https://api.openai.com/v1/embeddings")
61                    .bearer_auth(self.api_key.as_deref().unwrap())
62                    .json(&serde_json::json!({
63                        "input": text,
64                        "model": "text-embedding-3-large"
65                    }))
66                    .send()
67                    .expect("failed to send OpenAI embedding request")
68                    .json()
69                    .expect("failed to parse OpenAI embedding response");
70                response.data.into_iter().next().unwrap().embedding
71            }
72            EmbedderBackend::Ollama => {
73                let response: OllamaResponse = self.client
74                    .post("http://localhost:11434/api/embed")
75                    .json(&serde_json::json!({
76                        "model": "nomic-embed-text",
77                        "input": text
78                    }))
79                    .send()
80                    .expect("failed to send Ollama embedding request — is ollama running?")
81                    .json()
82                    .expect("failed to parse Ollama embedding response");
83                response.embeddings.into_iter().next().unwrap()
84            }
85        }
86    }
87}