codemem_embeddings/
openai.rs

1//! OpenAI-compatible embedding provider for Codemem.
2//!
3//! Works with OpenAI, Azure OpenAI, Together.ai, and any OpenAI-compatible API.
4//! Default model: text-embedding-3-small.
5
6use codemem_core::CodememError;
7
8/// Default OpenAI API base URL.
9pub const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1";
10
11/// Default model.
12pub const DEFAULT_MODEL: &str = "text-embedding-3-small";
13
14/// OpenAI embedding provider.
15pub struct OpenAIProvider {
16    api_key: String,
17    model: String,
18    dimensions: usize,
19    base_url: String,
20    client: reqwest::blocking::Client,
21}
22
23impl OpenAIProvider {
24    /// Create a new OpenAI provider.
25    pub fn new(api_key: &str, model: &str, dimensions: usize, base_url: Option<&str>) -> Self {
26        Self {
27            api_key: api_key.to_string(),
28            model: model.to_string(),
29            dimensions,
30            base_url: base_url.unwrap_or(DEFAULT_BASE_URL).to_string(),
31            client: reqwest::blocking::Client::new(),
32        }
33    }
34
35    /// Create with default model (text-embedding-3-small, 768 dims).
36    pub fn with_api_key(api_key: &str) -> Self {
37        Self::new(api_key, DEFAULT_MODEL, 768, None)
38    }
39}
40
41impl super::EmbeddingProvider for OpenAIProvider {
42    fn dimensions(&self) -> usize {
43        self.dimensions
44    }
45
46    fn embed(&self, text: &str) -> Result<Vec<f32>, CodememError> {
47        let url = format!("{}/embeddings", self.base_url);
48
49        let mut body = serde_json::json!({
50            "model": self.model,
51            "input": text,
52        });
53
54        // text-embedding-3-* supports custom dimensions
55        if self.model.starts_with("text-embedding-3") {
56            body["dimensions"] = serde_json::json!(self.dimensions);
57        }
58
59        let response = self
60            .client
61            .post(&url)
62            .header("Authorization", format!("Bearer {}", self.api_key))
63            .header("Content-Type", "application/json")
64            .json(&body)
65            .send()
66            .map_err(|e| CodememError::Embedding(format!("OpenAI request failed: {e}")))?;
67
68        if !response.status().is_success() {
69            let status = response.status();
70            let body = response.text().unwrap_or_default();
71            return Err(CodememError::Embedding(format!(
72                "OpenAI returned status {}: {}",
73                status, body
74            )));
75        }
76
77        let json: serde_json::Value = response
78            .json()
79            .map_err(|e| CodememError::Embedding(format!("OpenAI response parse error: {e}")))?;
80
81        let embedding = json
82            .get("data")
83            .and_then(|v| v.as_array())
84            .and_then(|arr| arr.first())
85            .and_then(|item| item.get("embedding"))
86            .and_then(|v| v.as_array())
87            .ok_or_else(|| CodememError::Embedding("Missing embedding in OpenAI response".into()))?
88            .iter()
89            .map(|v| v.as_f64().unwrap_or(0.0) as f32)
90            .collect();
91
92        Ok(embedding)
93    }
94
95    fn name(&self) -> &str {
96        "openai"
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn openai_provider_construction() {
106        let provider = OpenAIProvider::with_api_key("test-key");
107        assert_eq!(provider.model, DEFAULT_MODEL);
108        assert_eq!(provider.dimensions, 768);
109        assert_eq!(provider.base_url, DEFAULT_BASE_URL);
110    }
111
112    #[test]
113    fn openai_provider_custom_base_url() {
114        let provider = OpenAIProvider::new(
115            "test-key",
116            "custom-model",
117            1536,
118            Some("https://api.together.xyz/v1"),
119        );
120        assert_eq!(provider.base_url, "https://api.together.xyz/v1");
121        assert_eq!(provider.model, "custom-model");
122        assert_eq!(provider.dimensions, 1536);
123    }
124}
codemem_embeddings/openai.rs

codemem_embeddings/
openai.rs
