1use async_trait::async_trait;
12use chrono::TimeZone;
13
14use everruns_provider::OpenAIProtocolChatDriver;
15use everruns_provider::OpenResponsesProtocolChatDriver;
16use everruns_provider::credential_schema::CredentialFormSchema;
17use everruns_provider::driver_registry::{
18 BoxedChatDriver, BoxedEmbeddingsDriver, ChatDriver, DiscoveredModel, DriverDescriptor,
19 DriverId, DriverRegistry, EmbeddingsDriverFactory, LlmCallConfig, LlmMessage,
20 LlmResponseStream, ServiceKind,
21};
22use everruns_provider::error::{AgentLoopError, Result};
23use everruns_provider::openai_protocol::{
24 apply_models_api_auth, is_azure_openai_api_url, is_openai_api_url, models_api_status_error,
25 models_url_for_api_url, normalize_api_url,
26};
27use everruns_provider::{CompactRequest, CompactResponse};
28
29use crate::types::OpenAiModelsResponse;
30
31#[derive(Clone)]
62pub struct OpenAIChatDriver {
63 inner: OpenResponsesProtocolChatDriver,
64 uses_custom_url: bool,
66}
67
68impl OpenAIChatDriver {
69 pub fn new(api_key: impl Into<String>) -> Self {
71 Self {
72 inner: OpenResponsesProtocolChatDriver::new(api_key),
73 uses_custom_url: false,
74 }
75 }
76
77 pub fn with_base_url(api_key: impl Into<String>, api_url: impl Into<String>) -> Self {
79 let api_url = normalize_api_url(&api_url.into(), "/responses");
80 Self {
81 inner: OpenResponsesProtocolChatDriver::with_base_url(api_key, api_url),
82 uses_custom_url: true,
83 }
84 }
85
86 pub fn api_url(&self) -> &str {
88 self.inner.api_url()
89 }
90
91 pub fn uses_custom_url(&self) -> bool {
93 self.uses_custom_url
94 }
95
96 pub async fn compact_conversation(&self, request: CompactRequest) -> Result<CompactResponse> {
111 self.inner.compact(request).await
112 }
113
114 pub fn can_compact(&self) -> bool {
119 self.inner.supports_compact()
120 }
121}
122
123#[async_trait]
124impl ChatDriver for OpenAIChatDriver {
125 async fn chat_completion_stream(
126 &self,
127 messages: Vec<LlmMessage>,
128 config: &LlmCallConfig,
129 ) -> Result<LlmResponseStream> {
130 self.inner.chat_completion_stream(messages, config).await
131 }
132
133 async fn list_models(&self) -> Result<Option<Vec<DiscoveredModel>>> {
134 if self.uses_custom_url && !supports_model_listing(self.api_url()) {
136 return Ok(None);
137 }
138
139 let models_url = models_url_for_api_url(self.api_url());
140 list_openai_models(self.inner.client(), self.inner.api_key(), &models_url).await
141 }
142
143 fn supports_compact(&self) -> bool {
144 self.inner.supports_compact()
145 }
146
147 fn supports_parallel_tool_calls(&self, model: &str) -> bool {
148 self.inner.supports_parallel_tool_calls(model)
149 }
150
151 async fn compact(&self, request: CompactRequest) -> Result<Option<CompactResponse>> {
152 Ok(Some(self.inner.compact(request).await?))
153 }
154}
155
156impl std::fmt::Debug for OpenAIChatDriver {
157 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
158 f.debug_struct("OpenAIChatDriver")
159 .field("api_url", &self.api_url())
160 .field("api", &"Open Responses")
161 .field("api_key", &"[REDACTED]")
162 .finish()
163 }
164}
165
166#[derive(Clone)]
193pub struct OpenAICompletionsChatDriver {
194 inner: OpenAIProtocolChatDriver,
195 uses_custom_url: bool,
197}
198
199impl OpenAICompletionsChatDriver {
200 pub fn new(api_key: impl Into<String>) -> Self {
202 Self {
203 inner: OpenAIProtocolChatDriver::new(api_key),
204 uses_custom_url: false,
205 }
206 }
207
208 pub fn with_base_url(api_key: impl Into<String>, api_url: impl Into<String>) -> Self {
210 let api_url = normalize_api_url(&api_url.into(), "/chat/completions");
211 Self {
212 inner: OpenAIProtocolChatDriver::with_base_url(api_key, api_url),
213 uses_custom_url: true,
214 }
215 }
216
217 pub fn api_url(&self) -> &str {
219 self.inner.api_url()
220 }
221
222 pub fn uses_custom_url(&self) -> bool {
224 self.uses_custom_url
225 }
226}
227
228#[async_trait]
229impl ChatDriver for OpenAICompletionsChatDriver {
230 async fn chat_completion_stream(
231 &self,
232 messages: Vec<LlmMessage>,
233 config: &LlmCallConfig,
234 ) -> Result<LlmResponseStream> {
235 self.inner.chat_completion_stream(messages, config).await
236 }
237
238 async fn list_models(&self) -> Result<Option<Vec<DiscoveredModel>>> {
239 if self.uses_custom_url && !supports_model_listing(self.api_url()) {
241 return Ok(None);
242 }
243
244 let models_url = models_url_for_api_url(self.api_url());
245 list_openai_models(self.inner.client(), self.inner.api_key(), &models_url).await
246 }
247
248 fn supports_parallel_tool_calls(&self, model: &str) -> bool {
249 self.inner.supports_parallel_tool_calls(model)
250 }
251}
252
253impl std::fmt::Debug for OpenAICompletionsChatDriver {
254 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
255 f.debug_struct("OpenAICompletionsChatDriver")
256 .field("api_url", &self.api_url())
257 .field("api", &"Chat Completions")
258 .field("api_key", &"[REDACTED]")
259 .finish()
260 }
261}
262
263async fn list_openai_models(
269 client: &reqwest::Client,
270 api_key: &str,
271 models_url: &str,
272) -> Result<Option<Vec<DiscoveredModel>>> {
273 let response = apply_models_api_auth(client.get(models_url), models_url, api_key)
274 .send()
275 .await
276 .map_err(|e| AgentLoopError::llm(format!("Failed to fetch models: {}", e)))?;
277
278 if !response.status().is_success() {
279 let status = response.status();
280 let _ = response.bytes().await; return Err(models_api_status_error(status));
282 }
283
284 let models_response: OpenAiModelsResponse = response
285 .json()
286 .await
287 .map_err(|e| AgentLoopError::llm(format!("Failed to parse models response: {}", e)))?;
288
289 let discovered: Vec<DiscoveredModel> = models_response
291 .data
292 .into_iter()
293 .filter(|m| m.is_chat_model())
294 .map(|m| DiscoveredModel {
295 model_id: m.id,
296 display_name: None, created_at: chrono::Utc.timestamp_opt(m.created, 0).single(),
298 owned_by: Some(m.owned_by),
299 discovered_profile: None,
300 })
301 .collect();
302
303 Ok(Some(discovered))
304}
305
306fn supports_model_listing(api_url: &str) -> bool {
310 is_openai_api_url(api_url) || is_azure_openai_api_url(api_url)
311}
312
313pub fn register_driver(registry: &mut DriverRegistry) {
336 let openai_embeddings_factory: EmbeddingsDriverFactory = std::sync::Arc::new(|config| {
341 let api_key = config.api_key.as_deref().unwrap_or("");
342 let driver = match config.base_url.as_deref() {
343 Some(url) => crate::embeddings::OpenAIEmbeddingsDriver::with_base_url(api_key, url),
344 None => crate::embeddings::OpenAIEmbeddingsDriver::new(api_key),
345 };
346 Box::new(driver) as BoxedEmbeddingsDriver
347 });
348 registry.register_descriptor(DriverDescriptor {
349 services: vec![ServiceKind::Chat, ServiceKind::Realtime, ServiceKind::Embeddings],
350 credential_schema: CredentialFormSchema::api_key(
351 "Create an API key at [platform.openai.com/api-keys](https://platform.openai.com/api-keys).",
352 ),
353 embeddings: Some(openai_embeddings_factory),
354 ..DriverDescriptor::chat_only(DriverId::OpenAI, |config| {
355 let api_key = config.api_key.as_deref().unwrap_or("");
356 let driver = match config.base_url.as_deref() {
357 Some(url) => OpenAIChatDriver::with_base_url(api_key, url),
358 None => OpenAIChatDriver::new(api_key),
359 };
360 Box::new(driver) as BoxedChatDriver
361 })
362 });
363
364 registry.register_descriptor(DriverDescriptor {
365 credential_schema: CredentialFormSchema::api_key(
366 "Use an API key for your Azure OpenAI resource and set the resource endpoint as the base URL.",
367 ),
368 ..DriverDescriptor::chat_only(DriverId::AzureOpenAI, |config| {
369 let api_key = config.api_key.as_deref().unwrap_or("");
370 let driver = match config.base_url.as_deref() {
371 Some(url) => OpenAIChatDriver::with_base_url(api_key, url),
372 None => OpenAIChatDriver::new(api_key),
373 };
374 Box::new(driver) as BoxedChatDriver
375 })
376 });
377
378 registry.register_descriptor(DriverDescriptor {
380 credential_schema: CredentialFormSchema::api_key(
381 "Create an API key at [platform.openai.com/api-keys](https://platform.openai.com/api-keys).",
382 ),
383 ..DriverDescriptor::chat_only(DriverId::OpenAICompletions, |config| {
384 let api_key = config.api_key.as_deref().unwrap_or("");
385 let driver = match config.base_url.as_deref() {
386 Some(url) => OpenAICompletionsChatDriver::with_base_url(api_key, url),
387 None => OpenAICompletionsChatDriver::new(api_key),
388 };
389 Box::new(driver) as BoxedChatDriver
390 })
391 });
392}
393
394#[cfg(test)]
395mod tests {
396 use super::{is_openai_api_url, supports_model_listing};
397
398 #[test]
399 fn supports_model_listing_for_openai_host_with_port() {
400 assert!(supports_model_listing(
401 "https://api.openai.com:443/v1/responses"
402 ));
403 }
404
405 #[test]
406 fn rejects_non_openai_hosts_for_model_listing() {
407 assert!(!is_openai_api_url("https://example.com/v1/responses"));
408 assert!(!supports_model_listing(
409 "https://openrouter.ai/api/v1/responses"
410 ));
411 }
412}