uni_xervo/provider/
openai.rs

1use crate::api::{ModelAliasSpec, ModelTask};
2use crate::error::{Result, RuntimeError};
3use crate::provider::remote_common::{RemoteProviderBase, check_http_status, resolve_api_key};
4use crate::traits::{
5    EmbedResult, EmbeddingModel, GenerationOptions, GenerationResult, GeneratorModel,
6    LoadedModelHandle, Message, MessageRole, ModelProvider, ProviderCapabilities, ProviderHealth,
7    TokenUsage,
8};
9use async_trait::async_trait;
10use reqwest::Client;
11use serde_json::json;
12use std::sync::Arc;
13
14/// Remote provider that calls the [OpenAI API](https://platform.openai.com/docs/api-reference)
15/// for embedding (`/embeddings`) and text generation (`/chat/completions`).
16///
17/// Requires the `OPENAI_API_KEY` environment variable (or a custom env var name
18/// via the `api_key_env` option).
19///
20/// Set the `base_url` option to target an OpenAI-compatible server (OpenRouter,
21/// vLLM, LM Studio, Ollama, internal proxies). The value should include the
22/// version path segment (e.g. `http://localhost:8000/v1`). Defaults to
23/// `https://api.openai.com/v1` when unset.
24pub struct RemoteOpenAIProvider {
25    base: RemoteProviderBase,
26}
27
28impl Default for RemoteOpenAIProvider {
29    fn default() -> Self {
30        Self {
31            base: RemoteProviderBase::new(),
32        }
33    }
34}
35
36impl RemoteOpenAIProvider {
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    #[cfg(test)]
42    fn insert_test_breaker(&self, key: crate::api::ModelRuntimeKey, age: std::time::Duration) {
43        self.base.insert_test_breaker(key, age);
44    }
45
46    #[cfg(test)]
47    fn breaker_count(&self) -> usize {
48        self.base.breaker_count()
49    }
50
51    #[cfg(test)]
52    fn force_cleanup_now_for_test(&self) {
53        self.base.force_cleanup_now_for_test();
54    }
55}
56
57#[async_trait]
58impl ModelProvider for RemoteOpenAIProvider {
59    fn provider_id(&self) -> &'static str {
60        "remote/openai"
61    }
62
63    fn capabilities(&self) -> ProviderCapabilities {
64        ProviderCapabilities {
65            supported_tasks: vec![ModelTask::Embed, ModelTask::Generate],
66        }
67    }
68
69    async fn load(&self, spec: &ModelAliasSpec) -> Result<LoadedModelHandle> {
70        match spec.task {
71            ModelTask::Embed => {
72                let api_key = resolve_api_key(&spec.options, "api_key_env", "OPENAI_API_KEY")?;
73                let base_url = resolve_base_url(&spec.options);
74                let embedding_dimensions = spec
75                    .options
76                    .get("embedding_dimensions")
77                    .and_then(|v| v.as_u64())
78                    .map(|v| v as u32);
79                let default_dims = match spec.model_id.as_str() {
80                    "text-embedding-3-large" => 3072,
81                    _ => 1536,
82                };
83                let model = OpenAIEmbeddingModel {
84                    client: self.base.client.clone(),
85                    cb: self.base.circuit_breaker_for(spec),
86                    model_id: spec.model_id.clone(),
87                    api_key,
88                    base_url,
89                    dimensions: embedding_dimensions.unwrap_or(default_dims),
90                };
91                let handle: Arc<dyn EmbeddingModel> = Arc::new(model);
92                Ok(Arc::new(handle) as LoadedModelHandle)
93            }
94            ModelTask::Generate => {
95                let api_key = resolve_api_key(&spec.options, "api_key_env", "OPENAI_API_KEY")?;
96                let base_url = resolve_base_url(&spec.options);
97                let model = OpenAIGeneratorModel {
98                    client: self.base.client.clone(),
99                    cb: self.base.circuit_breaker_for(spec),
100                    model_id: spec.model_id.clone(),
101                    api_key,
102                    base_url,
103                };
104                let handle: Arc<dyn GeneratorModel> = Arc::new(model);
105                Ok(Arc::new(handle) as LoadedModelHandle)
106            }
107            ModelTask::Raw => Err(RuntimeError::CapabilityMismatch(
108                "OpenAI provider does not support task Raw".to_string(),
109            )),
110            _ => Err(RuntimeError::CapabilityMismatch(format!(
111                "OpenAI provider does not support task {:?}",
112                spec.task
113            ))),
114        }
115    }
116
117    async fn health(&self) -> ProviderHealth {
118        ProviderHealth::Healthy
119    }
120}
121
122const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1";
123
124/// Resolve `base_url` from options, falling back to the OpenAI default and
125/// stripping a trailing `/` so callers can append `/embeddings` etc. directly.
126fn resolve_base_url(options: &serde_json::Value) -> String {
127    let raw = options
128        .get("base_url")
129        .and_then(|v| v.as_str())
130        .unwrap_or(DEFAULT_BASE_URL);
131    raw.trim_end_matches('/').to_string()
132}
133
134/// Embedding model backed by the OpenAI embeddings API.
135pub struct OpenAIEmbeddingModel {
136    client: Client,
137    cb: crate::reliability::CircuitBreakerWrapper,
138    model_id: String,
139    api_key: String,
140    base_url: String,
141    dimensions: u32,
142}
143
144#[async_trait]
145impl EmbeddingModel for OpenAIEmbeddingModel {
146    async fn embed(&self, texts: &[&str]) -> Result<EmbedResult> {
147        let texts: Vec<String> = texts.iter().map(|s| s.to_string()).collect();
148
149        self.cb
150            .call(move || async move {
151                let response = self
152                    .client
153                    .post(format!("{}/embeddings", self.base_url))
154                    .header("Authorization", format!("Bearer {}", self.api_key))
155                    .json(&json!({
156                        "model": self.model_id,
157                        "input": texts
158                    }))
159                    .send()
160                    .await
161                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
162
163                let body: serde_json::Value = check_http_status("OpenAI", response)?
164                    .json()
165                    .await
166                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
167
168                let mut vectors = Vec::new();
169                if let Some(data) = body.get("data").and_then(|d| d.as_array()) {
170                    for item in data {
171                        if let Some(embedding) = item.get("embedding").and_then(|e| e.as_array()) {
172                            let vec: Vec<f32> = embedding
173                                .iter()
174                                .filter_map(|v| v.as_f64().map(|f| f as f32))
175                                .collect();
176                            vectors.push(vec);
177                        }
178                    }
179                }
180
181                // OpenAI embeddings response carries a `usage` object with
182                // `prompt_tokens` and `total_tokens` (no completion tokens).
183                let usage = body.get("usage").map(|u| {
184                    let prompt = u["prompt_tokens"].as_u64().unwrap_or(0) as usize;
185                    let total = u["total_tokens"].as_u64().unwrap_or(prompt as u64) as usize;
186                    TokenUsage {
187                        prompt_tokens: prompt,
188                        completion_tokens: 0,
189                        total_tokens: total,
190                    }
191                });
192
193                Ok(EmbedResult { vectors, usage })
194            })
195            .await
196    }
197
198    fn dimensions(&self) -> u32 {
199        self.dimensions
200    }
201}
202
203impl crate::traits::ModelInfo for OpenAIEmbeddingModel {
204    fn model_id(&self) -> &str {
205        &self.model_id
206    }
207}
208
209// ---------------------------------------------------------------------------
210// Generator
211// ---------------------------------------------------------------------------
212
213struct OpenAIGeneratorModel {
214    client: Client,
215    cb: crate::reliability::CircuitBreakerWrapper,
216    model_id: String,
217    api_key: String,
218    base_url: String,
219}
220
221impl crate::traits::ModelInfo for OpenAIGeneratorModel {
222    fn model_id(&self) -> &str {
223        &self.model_id
224    }
225}
226
227#[async_trait]
228impl GeneratorModel for OpenAIGeneratorModel {
229    async fn generate(
230        &self,
231        messages: &[Message],
232        options: GenerationOptions,
233    ) -> Result<GenerationResult> {
234        let messages: Vec<serde_json::Value> = messages
235            .iter()
236            .map(|msg| {
237                let role = match msg.role {
238                    MessageRole::System => "system",
239                    MessageRole::User => "user",
240                    MessageRole::Assistant => "assistant",
241                };
242                json!({ "role": role, "content": msg.text() })
243            })
244            .collect();
245
246        self.cb
247            .call(move || async move {
248                let mut body = json!({
249                    "model": self.model_id,
250                    "messages": messages,
251                });
252
253                if let Some(max_tokens) = options.max_tokens {
254                    body["max_completion_tokens"] = json!(max_tokens);
255                }
256                if let Some(temperature) = options.temperature {
257                    body["temperature"] = json!(temperature);
258                }
259                if let Some(top_p) = options.top_p {
260                    body["top_p"] = json!(top_p);
261                }
262
263                let response = self
264                    .client
265                    .post(format!("{}/chat/completions", self.base_url))
266                    .header("Authorization", format!("Bearer {}", self.api_key))
267                    .json(&body)
268                    .send()
269                    .await
270                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
271
272                let body: serde_json::Value = check_http_status("OpenAI", response)?
273                    .json()
274                    .await
275                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
276
277                let text = body["choices"][0]["message"]["content"]
278                    .as_str()
279                    .unwrap_or("")
280                    .to_string();
281
282                let usage = body.get("usage").map(|u| TokenUsage {
283                    prompt_tokens: u["prompt_tokens"].as_u64().unwrap_or(0) as usize,
284                    completion_tokens: u["completion_tokens"].as_u64().unwrap_or(0) as usize,
285                    total_tokens: u["total_tokens"].as_u64().unwrap_or(0) as usize,
286                });
287
288                Ok(GenerationResult {
289                    text,
290                    usage,
291                    images: vec![],
292                    audio: None,
293                })
294            })
295            .await
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302    use crate::api::ModelRuntimeKey;
303    use crate::provider::remote_common::RemoteProviderBase;
304    use crate::traits::ModelProvider;
305    use std::time::Duration;
306
307    static ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
308
309    fn spec(alias: &str, task: ModelTask, model_id: &str) -> ModelAliasSpec {
310        ModelAliasSpec {
311            alias: alias.to_string(),
312            task,
313            provider_id: "remote/openai".to_string(),
314            model_id: model_id.to_string(),
315            revision: None,
316            warmup: crate::api::WarmupPolicy::Lazy,
317            required: false,
318            timeout: None,
319            load_timeout: None,
320            retry: None,
321            options: serde_json::Value::Null,
322        }
323    }
324
325    #[tokio::test]
326    async fn breaker_reused_for_same_runtime_key() {
327        let _lock = ENV_LOCK.lock().await;
328        // SAFETY: protected by ENV_LOCK
329        unsafe { std::env::set_var("OPENAI_API_KEY", "test-key") };
330
331        let provider = RemoteOpenAIProvider::new();
332        let s1 = spec("embed/a", ModelTask::Embed, "text-embedding-3-small");
333        let s2 = spec("embed/b", ModelTask::Embed, "text-embedding-3-small");
334
335        let _ = provider.load(&s1).await.unwrap();
336        let _ = provider.load(&s2).await.unwrap();
337
338        assert_eq!(provider.breaker_count(), 1);
339
340        // SAFETY: protected by ENV_LOCK
341        unsafe { std::env::remove_var("OPENAI_API_KEY") };
342    }
343
344    #[tokio::test]
345    async fn breaker_isolated_by_task_and_model() {
346        let _lock = ENV_LOCK.lock().await;
347        // SAFETY: protected by ENV_LOCK
348        unsafe { std::env::set_var("OPENAI_API_KEY", "test-key") };
349
350        let provider = RemoteOpenAIProvider::new();
351        let embed = spec("embed/a", ModelTask::Embed, "text-embedding-3-small");
352        let gen_spec = spec("chat/a", ModelTask::Generate, "gpt-4o-mini");
353
354        let _ = provider.load(&embed).await.unwrap();
355        let _ = provider.load(&gen_spec).await.unwrap();
356
357        assert_eq!(provider.breaker_count(), 2);
358
359        // SAFETY: protected by ENV_LOCK
360        unsafe { std::env::remove_var("OPENAI_API_KEY") };
361    }
362
363    #[tokio::test]
364    async fn breaker_cleanup_evicts_stale_entries() {
365        let _lock = ENV_LOCK.lock().await;
366        // SAFETY: protected by ENV_LOCK
367        unsafe { std::env::set_var("OPENAI_API_KEY", "test-key") };
368
369        let provider = RemoteOpenAIProvider::new();
370        let stale = spec("embed/stale", ModelTask::Embed, "text-embedding-3-small");
371        let fresh = spec("embed/fresh", ModelTask::Embed, "text-embedding-3-large");
372        provider.insert_test_breaker(
373            ModelRuntimeKey::new(&stale),
374            RemoteProviderBase::BREAKER_TTL + Duration::from_secs(5),
375        );
376        provider.insert_test_breaker(ModelRuntimeKey::new(&fresh), Duration::from_secs(1));
377        assert_eq!(provider.breaker_count(), 2);
378
379        provider.force_cleanup_now_for_test();
380        let _ = provider.load(&fresh).await.unwrap();
381
382        assert_eq!(provider.breaker_count(), 1);
383
384        // SAFETY: protected by ENV_LOCK
385        unsafe { std::env::remove_var("OPENAI_API_KEY") };
386    }
387
388    #[tokio::test]
389    async fn default_embedding_dimensions() {
390        let _lock = ENV_LOCK.lock().await;
391        // SAFETY: protected by ENV_LOCK
392        unsafe { std::env::set_var("OPENAI_API_KEY", "test-key") };
393
394        let provider = RemoteOpenAIProvider::new();
395        let s = spec("embed/dim", ModelTask::Embed, "text-embedding-3-small");
396
397        let handle = provider.load(&s).await.unwrap();
398        let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
399        assert_eq!(model.dimensions(), 1536);
400
401        // SAFETY: protected by ENV_LOCK
402        unsafe { std::env::remove_var("OPENAI_API_KEY") };
403    }
404
405    #[tokio::test]
406    async fn custom_embedding_dimensions() {
407        let _lock = ENV_LOCK.lock().await;
408        // SAFETY: protected by ENV_LOCK
409        unsafe { std::env::set_var("OPENAI_API_KEY", "test-key") };
410
411        let provider = RemoteOpenAIProvider::new();
412        let mut s = spec(
413            "embed/dim-custom",
414            ModelTask::Embed,
415            "text-embedding-3-small",
416        );
417        s.options = serde_json::json!({"embedding_dimensions": 256});
418
419        let handle = provider.load(&s).await.unwrap();
420        let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
421        assert_eq!(model.dimensions(), 256);
422
423        // SAFETY: protected by ENV_LOCK
424        unsafe { std::env::remove_var("OPENAI_API_KEY") };
425    }
426
427    #[test]
428    fn resolve_base_url_defaults_to_openai() {
429        assert_eq!(
430            resolve_base_url(&serde_json::Value::Null),
431            "https://api.openai.com/v1"
432        );
433        assert_eq!(
434            resolve_base_url(&serde_json::json!({})),
435            "https://api.openai.com/v1"
436        );
437    }
438
439    #[test]
440    fn resolve_base_url_uses_custom_value() {
441        assert_eq!(
442            resolve_base_url(&serde_json::json!({"base_url": "http://localhost:8000/v1"})),
443            "http://localhost:8000/v1"
444        );
445    }
446
447    #[test]
448    fn resolve_base_url_strips_trailing_slash() {
449        assert_eq!(
450            resolve_base_url(&serde_json::json!({"base_url": "http://localhost:8000/v1/"})),
451            "http://localhost:8000/v1"
452        );
453    }
454
455    #[tokio::test]
456    async fn load_accepts_custom_base_url() {
457        let _lock = ENV_LOCK.lock().await;
458        // SAFETY: protected by ENV_LOCK
459        unsafe { std::env::set_var("OPENAI_API_KEY", "test-key") };
460
461        let provider = RemoteOpenAIProvider::new();
462        let mut s = spec("embed/local", ModelTask::Embed, "text-embedding-3-small");
463        s.options = serde_json::json!({"base_url": "http://localhost:8000/v1"});
464
465        assert!(provider.load(&s).await.is_ok());
466
467        // SAFETY: protected by ENV_LOCK
468        unsafe { std::env::remove_var("OPENAI_API_KEY") };
469    }
470}
uni_xervo/provider/openai.rs

uni_xervo/provider/
openai.rs
