aither_openai/
provider.rs

1use crate::{
2    DEEPSEEK_BASE_URL, DEFAULT_BASE_URL, OPENROUTER_BASE_URL, client::OpenAI, error::OpenAIError,
3};
4use aither_core::llm::{
5    LanguageModelProvider, model::Profile as ModelProfile, provider::Profile as ProviderProfile,
6};
7use aither_models::lookup as lookup_model_info;
8use serde::Deserialize;
9use std::{future::Future, sync::Arc};
10use zenwave::{Client, client, header};
11
12/// Provider capable of listing and instantiating `OpenAI` models.
13#[derive(Clone, Debug)]
14pub struct OpenAIProvider {
15    inner: Arc<ProviderConfig>,
16}
17
18impl OpenAIProvider {
19    /// Create a new provider with the given API key.
20    pub fn new(api_key: impl Into<String>) -> Self {
21        Self {
22            inner: Arc::new(ProviderConfig {
23                api_key: api_key.into(),
24                base_url: DEFAULT_BASE_URL.to_string(),
25                organization: None,
26            }),
27        }
28    }
29
30    /// Convenience constructor targeting [`Deepseek`](https://api-docs.deepseek.com)'s API.
31    #[must_use]
32    pub fn deepseek(api_key: impl Into<String>) -> Self {
33        Self::new(api_key).base_url(DEEPSEEK_BASE_URL)
34    }
35
36    /// Convenience constructor targeting [`OpenRouter`](https://openrouter.ai)'s API.
37    #[must_use]
38    pub fn openrouter(api_key: impl Into<String>) -> Self {
39        Self::new(api_key).base_url(OPENROUTER_BASE_URL)
40    }
41
42    /// Override the REST base URL (useful for Azure or sandboxes).
43    #[must_use]
44    pub fn base_url(mut self, url: impl Into<String>) -> Self {
45        Arc::make_mut(&mut self.inner).base_url = url.into();
46        self
47    }
48
49    /// Attach an organization header for model management calls.
50    #[must_use]
51    pub fn organization(mut self, organization: impl Into<String>) -> Self {
52        Arc::make_mut(&mut self.inner).organization = Some(organization.into());
53        self
54    }
55
56    fn client_for(&self, model: impl Into<String>) -> OpenAI {
57        let mut builder = OpenAI::builder(self.inner.api_key.clone())
58            .base_url(self.inner.base_url.clone())
59            .model(model);
60        if let Some(org) = &self.inner.organization {
61            builder = builder.organization(org.clone());
62        }
63        builder.build()
64    }
65}
66
67impl LanguageModelProvider for OpenAIProvider {
68    type Model = OpenAI;
69    type Error = OpenAIError;
70
71    fn list_models(&self) -> impl Future<Output = Result<Vec<ModelProfile>, Self::Error>> + Send {
72        let cfg = self.inner.clone();
73        async move {
74            let endpoint = format!("{}/models", cfg.base_url.trim_end_matches('/'));
75            let mut backend = client();
76            let mut builder = backend
77                .get(endpoint)
78                .map_err(OpenAIError::Http)?
79                .header(
80                    header::AUTHORIZATION.as_str(),
81                    format!("Bearer {}", cfg.api_key),
82                )
83                .map_err(OpenAIError::Http)?;
84            if let Some(org) = &cfg.organization {
85                builder = builder
86                    .header("OpenAI-Organization", org.clone())
87                    .map_err(OpenAIError::Http)?;
88            }
89            let response: ModelListResponse = builder.json().await.map_err(OpenAIError::Http)?;
90            Ok(response
91                .data
92                .into_iter()
93                .filter_map(|model| {
94                    let context_length = model
95                        .context_length
96                        .or(model.max_tokens)
97                        .or(model.input_token_limit)
98                        .or_else(|| lookup_model_info(&model.id).and_then(aither_models::ModelEntry::max_input_tokens));
99                    let Some(context_length) = context_length else {
100                        tracing::warn!(
101                            model = %model.id,
102                            "skip model from OpenAI list: missing context_length and no aither-models fallback"
103                        );
104                        return None;
105                    };
106                    Some(ModelProfile::new(
107                        model.id.clone(),
108                        model.owned_by.clone().unwrap_or_else(|| "openai".into()),
109                        model.id,
110                        "OpenAI model",
111                        context_length,
112                    ))
113                })
114                .collect())
115        }
116    }
117
118    fn get_model(
119        &self,
120        name: &str,
121    ) -> impl Future<Output = Result<Self::Model, Self::Error>> + Send {
122        let client = self.client_for(name.to_string());
123        async move { Ok(client) }
124    }
125
126    fn profile() -> ProviderProfile {
127        ProviderProfile::new("openai", "Official OpenAI API provider")
128    }
129}
130
131#[derive(Debug, Clone)]
132struct ProviderConfig {
133    api_key: String,
134    base_url: String,
135    organization: Option<String>,
136}
137
138#[derive(Debug, Deserialize)]
139struct ModelListResponse {
140    data: Vec<ModelDescriptor>,
141}
142
143#[derive(Debug, Deserialize)]
144struct ModelDescriptor {
145    id: String,
146    #[serde(default)]
147    owned_by: Option<String>,
148    #[serde(default)]
149    context_length: Option<u32>,
150    #[serde(default)]
151    max_tokens: Option<u32>,
152    #[serde(default, alias = "inputTokenLimit", alias = "input_token_limit")]
153    input_token_limit: Option<u32>,
154}
aither_openai/provider.rs

aither_openai/
provider.rs
