llm_kernel/embedding/
openai.rs

1//! OpenAI text-embedding provider (sync, via ureq).
2
3use serde::Deserialize;
4
5use crate::embedding::types::{EmbeddingProvider, EmbeddingResult};
6use crate::error::{KernelError, Result};
7
8#[derive(Deserialize)]
9struct EmbeddingData {
10    embedding: Vec<f32>,
11    index: usize,
12}
13
14#[derive(Deserialize)]
15struct EmbeddingResponse {
16    data: Vec<EmbeddingData>,
17}
18
19/// OpenAI embedding provider.
20///
21/// Uses `text-embedding-3-small` (1536-dim) by default.
22/// Swap model to `text-embedding-3-large` (3072-dim) for higher accuracy.
23///
24/// `api_key` is not exposed via `Debug` to prevent accidental logging.
25pub struct OpenAIEmbeddingClient {
26    api_key: String,
27    model: String,
28    dim: usize,
29}
30
31impl OpenAIEmbeddingClient {
32    /// Create with `text-embedding-3-small` (1536 dimensions).
33    pub fn new_small(api_key: impl Into<String>) -> Self {
34        Self {
35            api_key: api_key.into(),
36            model: "text-embedding-3-small".into(),
37            dim: 1536,
38        }
39    }
40
41    /// Create with `text-embedding-3-large` (3072 dimensions).
42    pub fn new_large(api_key: impl Into<String>) -> Self {
43        Self {
44            api_key: api_key.into(),
45            model: "text-embedding-3-large".into(),
46            dim: 3072,
47        }
48    }
49
50    /// Create with an explicit model name and embedding dimension.
51    ///
52    /// Use this for legacy models (`text-embedding-ada-002`), reduced-dimension
53    /// variants, or any future model not covered by [`new_small`](Self::new_small)
54    /// and [`new_large`](Self::new_large).
55    ///
56    /// # Panics
57    ///
58    /// Panics if `dim` is zero.
59    pub fn new_with_model(
60        api_key: impl Into<String>,
61        model: impl Into<String>,
62        dim: usize,
63    ) -> Self {
64        assert!(dim > 0, "dim must be non-zero");
65        Self {
66            api_key: api_key.into(),
67            model: model.into(),
68            dim,
69        }
70    }
71
72    /// Create from environment variable `OPENAI_API_KEY`.
73    ///
74    /// Always uses `text-embedding-3-small` (1536-dim). For a different model
75    /// use [`new_with_model`](Self::new_with_model) after reading the key manually.
76    pub fn from_env() -> Result<Self> {
77        let key = std::env::var("OPENAI_API_KEY")
78            .map_err(|_| KernelError::Embedding("OPENAI_API_KEY not set".into()))?;
79        Ok(Self::new_small(key))
80    }
81}
82
83use super::types::text_preview;
84
85impl EmbeddingProvider for OpenAIEmbeddingClient {
86    fn dim(&self) -> usize {
87        self.dim
88    }
89
90    fn name(&self) -> &str {
91        &self.model
92    }
93
94    fn embed(&self, text: &str) -> Result<EmbeddingResult> {
95        let config = ureq::config::Config::builder()
96            .timeout_global(Some(std::time::Duration::from_secs(30)))
97            .build();
98        let agent = ureq::Agent::new_with_config(config);
99
100        let body = serde_json::json!({
101            "model": self.model,
102            "input": text,
103        });
104
105        let mut resp = agent
106            .post("https://api.openai.com/v1/embeddings")
107            .header("Authorization", format!("Bearer {}", self.api_key))
108            .header("Content-Type", "application/json")
109            .send_json(body)
110            .map_err(KernelError::embedding)?;
111
112        let payload: EmbeddingResponse = resp
113            .body_mut()
114            .read_json()
115            .map_err(KernelError::embedding)?;
116
117        let vector = payload
118            .data
119            .into_iter()
120            .next()
121            .ok_or_else(|| KernelError::Embedding("empty embedding response".into()))?
122            .embedding;
123
124        Ok(EmbeddingResult {
125            vector,
126            text_preview: text_preview(text),
127        })
128    }
129
130    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<EmbeddingResult>> {
131        if texts.is_empty() {
132            return Ok(vec![]);
133        }
134
135        let config = ureq::config::Config::builder()
136            .timeout_global(Some(std::time::Duration::from_secs(60)))
137            .build();
138        let agent = ureq::Agent::new_with_config(config);
139
140        let body = serde_json::json!({
141            "model": self.model,
142            "input": texts,
143        });
144
145        let mut resp = agent
146            .post("https://api.openai.com/v1/embeddings")
147            .header("Authorization", format!("Bearer {}", self.api_key))
148            .header("Content-Type", "application/json")
149            .send_json(body)
150            .map_err(KernelError::embedding)?;
151
152        let payload: EmbeddingResponse = resp
153            .body_mut()
154            .read_json()
155            .map_err(KernelError::embedding)?;
156
157        // The OpenAI API does not guarantee that `data` is returned in input
158        // order; sort by `index` before zipping with `texts`.
159        let mut data = payload.data;
160        data.sort_unstable_by_key(|d| d.index);
161
162        if data.len() != texts.len() {
163            return Err(KernelError::Embedding(format!(
164                "API returned {} embeddings for {} inputs",
165                data.len(),
166                texts.len()
167            )));
168        }
169
170        let results = data
171            .into_iter()
172            .zip(texts.iter())
173            .map(|(item, &text)| EmbeddingResult {
174                vector: item.embedding,
175                text_preview: text_preview(text),
176            })
177            .collect();
178
179        Ok(results)
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186
187    #[test]
188    fn small_client_has_correct_dim() {
189        let client = OpenAIEmbeddingClient::new_small("test-key");
190        assert_eq!(client.dim(), 1536);
191        assert_eq!(client.name(), "text-embedding-3-small");
192    }
193
194    #[test]
195    fn large_client_has_correct_dim() {
196        let client = OpenAIEmbeddingClient::new_large("test-key");
197        assert_eq!(client.dim(), 3072);
198        assert_eq!(client.name(), "text-embedding-3-large");
199    }
200
201    #[test]
202    fn from_env_fails_without_key() {
203        // SAFETY: `remove_var` is unsafe in Rust 2024 because env mutation is
204        // not thread-safe. Tests run in their own process and this binary does
205        // not spawn threads that read OPENAI_API_KEY concurrently.
206        unsafe { std::env::remove_var("OPENAI_API_KEY") };
207        assert!(OpenAIEmbeddingClient::from_env().is_err());
208    }
209
210    #[test]
211    fn parse_embedding_response() {
212        let raw = r#"{"object":"list","data":[{"object":"embedding","embedding":[0.1,-0.2,0.3],"index":0}],"model":"text-embedding-3-small","usage":{"prompt_tokens":5,"total_tokens":5}}"#;
213        let payload: EmbeddingResponse = serde_json::from_str(raw).unwrap();
214        assert_eq!(payload.data.len(), 1);
215        assert_eq!(payload.data[0].embedding, vec![0.1f32, -0.2, 0.3]);
216        assert_eq!(payload.data[0].index, 0);
217    }
218
219    #[test]
220    fn embed_batch_reorders_by_index() {
221        // Simulate an out-of-order API response (index 1 before index 0).
222        let raw = r#"{"object":"list","data":[{"object":"embedding","embedding":[0.2],"index":1},{"object":"embedding","embedding":[0.1],"index":0}],"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}"#;
223        let mut payload: EmbeddingResponse = serde_json::from_str(raw).unwrap();
224        payload.data.sort_unstable_by_key(|d| d.index);
225        assert_eq!(payload.data[0].embedding, vec![0.1f32]);
226        assert_eq!(payload.data[1].embedding, vec![0.2f32]);
227    }
228
229    #[test]
230    fn new_with_model_sets_name_and_dim() {
231        let client = OpenAIEmbeddingClient::new_with_model("key", "text-embedding-ada-002", 1536);
232        assert_eq!(client.dim(), 1536);
233        assert_eq!(client.name(), "text-embedding-ada-002");
234    }
235
236    #[test]
237    fn new_with_model_custom_dim() {
238        let client = OpenAIEmbeddingClient::new_with_model("key", "text-embedding-3-small", 512);
239        assert_eq!(client.dim(), 512);
240        assert_eq!(client.name(), "text-embedding-3-small");
241    }
242
243    #[test]
244    #[should_panic(expected = "dim must be non-zero")]
245    fn new_with_model_zero_dim_panics() {
246        OpenAIEmbeddingClient::new_with_model("key", "text-embedding-ada-002", 0);
247    }
248
249    #[test]
250    fn preview_ascii_truncated() {
251        let long = "a".repeat(100);
252        let preview = text_preview(&long);
253        assert!(preview.ends_with('…'));
254        assert_eq!(preview.chars().filter(|&c| c != '…').count(), 64);
255    }
256
257    #[test]
258    fn preview_short_not_truncated() {
259        assert_eq!(text_preview("hello"), "hello");
260    }
261
262    #[test]
263    fn preview_multibyte_no_panic() {
264        // Each Korean char is 3 bytes; byte-slicing at 64 would panic.
265        let korean = "안녕하세요".repeat(20);
266        let preview = text_preview(&korean);
267        assert!(preview.ends_with('…'));
268        assert_eq!(preview.chars().filter(|&c| c != '…').count(), 64);
269    }
270}
llm_kernel/embedding/openai.rs

llm_kernel/embedding/
openai.rs
