uni_xervo/provider/
openai.rs

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

uni_xervo/provider/
openai.rs
