1use 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#[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 #[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 #[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 #[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 #[must_use]
68 pub const fn credentials(&self) -> &CodexCredentials {
69 &self.credentials
70 }
71}
72
73#[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 #[serde(default)]
91 visibility: Option<String>,
92 #[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}