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#[derive(Deserialize)]
17struct OpenAIResponse {
18 data: Vec<OpenAIEmbedding>,
19}
20#[derive(Deserialize)]
21struct OpenAIEmbedding {
22 embedding: Vec<f32>,
23}
24
25#[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 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}