Skip to main content

aither_codex/
provider.rs

1//! [`LanguageModelProvider`] for the Codex subscription backend.
2//!
3//! The backend exposes its own model catalog at
4//! `GET {base_url}/models?client_version=…`, which is authoritative for what a
5//! subscription can run — [`CodexProvider::list_models`] surfaces it directly.
6
7use serde::Deserialize;
8use zenwave::{Client as _, client};
9
10use aither_core::llm::model::Profile as ModelProfile;
11use aither_core::llm::provider::{LanguageModelProvider, Profile};
12
13use crate::auth::CodexCredentials;
14use crate::client::{Codex, OnRefresh};
15use crate::constant::{ACCOUNT_ID_HEADER, CLIENT_VERSION, CODEX_BASE_URL};
16use crate::error::CodexError;
17
18/// Provider over the `ChatGPT` Codex backend.
19///
20/// Holds the credential set used for both catalog listing and model clients.
21/// Token lifecycle (expiry refresh, rotation persistence) is handled inside
22/// each [`Codex`] model client created by [`get_model`].
23#[derive(Clone)]
24pub struct CodexProvider {
25    credentials: CodexCredentials,
26    base_url: String,
27    on_refresh: Option<OnRefresh>,
28}
29
30impl std::fmt::Debug for CodexProvider {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        f.debug_struct("CodexProvider")
33            .field("base_url", &self.base_url)
34            .finish_non_exhaustive()
35    }
36}
37
38impl CodexProvider {
39    /// Create a provider from OAuth credentials.
40    #[must_use]
41    pub fn new(credentials: CodexCredentials) -> Self {
42        Self {
43            credentials,
44            base_url: CODEX_BASE_URL.into(),
45            on_refresh: None,
46        }
47    }
48
49    /// Override the backend base URL.
50    #[must_use]
51    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
52        self.base_url = base_url.into();
53        self
54    }
55
56    /// Propagate a token-rotation persistence callback to model clients.
57    #[must_use]
58    pub fn on_refresh(
59        mut self,
60        callback: impl Fn(&CodexCredentials) + Send + Sync + 'static,
61    ) -> Self {
62        self.on_refresh = Some(std::sync::Arc::new(callback));
63        self
64    }
65
66    /// The credentials backing this provider.
67    #[must_use]
68    pub const fn credentials(&self) -> &CodexCredentials {
69        &self.credentials
70    }
71}
72
73/// Response of `GET /codex/models`.
74#[derive(Debug, Deserialize)]
75struct ModelsResponse {
76    #[serde(default)]
77    models: Vec<ModelDescriptor>,
78}
79
80#[derive(Debug, Deserialize)]
81struct ModelDescriptor {
82    slug: String,
83    #[serde(default)]
84    display_name: Option<String>,
85    #[serde(default)]
86    description: Option<String>,
87    #[serde(default)]
88    context_window: Option<u32>,
89    /// Catalog visibility; `"hide"` entries are filtered out.
90    #[serde(default)]
91    visibility: Option<String>,
92    /// Sort priority (higher first); missing sorts last.
93    #[serde(default)]
94    priority: Option<i64>,
95}
96
97impl ModelDescriptor {
98    fn visible(&self) -> bool {
99        self.visibility.as_deref() != Some("hide")
100    }
101
102    fn into_profile(self) -> ModelProfile {
103        ModelProfile::new(
104            self.display_name.unwrap_or_else(|| self.slug.clone()),
105            "openai",
106            self.slug,
107            self.description.unwrap_or_default(),
108            self.context_window.unwrap_or_default(),
109        )
110    }
111}
112
113impl LanguageModelProvider for CodexProvider {
114    type Model = Codex;
115    type Error = CodexError;
116
117    async fn list_models(&self) -> Result<Vec<ModelProfile>, Self::Error> {
118        let url = format!("{}/models?client_version={CLIENT_VERSION}", self.base_url);
119        let mut backend = client();
120        let mut request = backend.get(&url)?.header(
121            "authorization",
122            format!("Bearer {}", self.credentials.access_token),
123        )?;
124        if let Some(account_id) = &self.credentials.account_id {
125            request = request.header(ACCOUNT_ID_HEADER, account_id)?;
126        }
127        let mut response: ModelsResponse = request.json().await?;
128        response
129            .models
130            .sort_by_key(|m| std::cmp::Reverse(m.priority.unwrap_or(0)));
131        Ok(response
132            .models
133            .into_iter()
134            .filter(ModelDescriptor::visible)
135            .map(ModelDescriptor::into_profile)
136            .collect())
137    }
138
139    fn get_model(
140        &self,
141        name: &str,
142    ) -> impl core::future::Future<Output = Result<Self::Model, Self::Error>> + Send {
143        let mut model = Codex::new(self.credentials.clone(), name);
144        if let Some(callback) = &self.on_refresh {
145            model = model.on_refresh({
146                let callback = callback.clone();
147                move |creds| callback(creds)
148            });
149        }
150        std::future::ready(Ok(model))
151    }
152
153    fn profile() -> Profile {
154        Profile::new("codex", "OpenAI models through a ChatGPT subscription")
155    }
156}