|
1 | | -use codex_app_server_protocol::AuthMode; |
| 1 | +use codex_api::ModelsClient; |
| 2 | +use codex_api::ReqwestTransport; |
| 3 | +use codex_protocol::openai_models::ModelInfo; |
2 | 4 | use codex_protocol::openai_models::ModelPreset; |
| 5 | +use http::HeaderMap; |
| 6 | +use std::sync::Arc; |
3 | 7 | use tokio::sync::RwLock; |
4 | 8 |
|
| 9 | +use crate::api_bridge::auth_provider_from_auth; |
| 10 | +use crate::api_bridge::map_api_error; |
| 11 | +use crate::auth::AuthManager; |
5 | 12 | use crate::config::Config; |
| 13 | +use crate::default_client::build_reqwest_client; |
| 14 | +use crate::error::Result as CoreResult; |
| 15 | +use crate::model_provider_info::ModelProviderInfo; |
6 | 16 | use crate::openai_models::model_family::ModelFamily; |
7 | 17 | use crate::openai_models::model_family::find_family_for_model; |
8 | 18 | use crate::openai_models::model_presets::builtin_model_presets; |
9 | 19 |
|
10 | 20 | #[derive(Debug)] |
11 | 21 | pub struct ModelsManager { |
| 22 | + // todo(aibrahim) merge available_models and model family creation into one struct |
12 | 23 | pub available_models: RwLock<Vec<ModelPreset>>, |
| 24 | + pub remote_models: RwLock<Vec<ModelInfo>>, |
13 | 25 | pub etag: String, |
14 | | - pub auth_mode: Option<AuthMode>, |
| 26 | + pub auth_manager: Arc<AuthManager>, |
15 | 27 | } |
16 | 28 |
|
17 | 29 | impl ModelsManager { |
18 | | - pub fn new(auth_mode: Option<AuthMode>) -> Self { |
| 30 | + pub fn new(auth_manager: Arc<AuthManager>) -> Self { |
19 | 31 | Self { |
20 | | - available_models: RwLock::new(builtin_model_presets(auth_mode)), |
| 32 | + available_models: RwLock::new(builtin_model_presets(auth_manager.get_auth_mode())), |
| 33 | + remote_models: RwLock::new(Vec::new()), |
21 | 34 | etag: String::new(), |
22 | | - auth_mode, |
| 35 | + auth_manager, |
23 | 36 | } |
24 | 37 | } |
25 | 38 |
|
26 | | - pub async fn refresh_available_models(&self) { |
27 | | - let models = builtin_model_presets(self.auth_mode); |
28 | | - *self.available_models.write().await = models; |
| 39 | + // do not use this function yet. It's work in progress. |
| 40 | + pub async fn refresh_available_models( |
| 41 | + &self, |
| 42 | + provider: &ModelProviderInfo, |
| 43 | + ) -> CoreResult<Vec<ModelInfo>> { |
| 44 | + let auth = self.auth_manager.auth(); |
| 45 | + let api_provider = provider.to_api_provider(auth.as_ref().map(|auth| auth.mode))?; |
| 46 | + let api_auth = auth_provider_from_auth(auth.clone(), provider).await?; |
| 47 | + let transport = ReqwestTransport::new(build_reqwest_client()); |
| 48 | + let client = ModelsClient::new(transport, api_provider, api_auth); |
| 49 | + |
| 50 | + let response = client |
| 51 | + .list_models(env!("CARGO_PKG_VERSION"), HeaderMap::new()) |
| 52 | + .await |
| 53 | + .map_err(map_api_error)?; |
| 54 | + |
| 55 | + let models = response.models; |
| 56 | + *self.remote_models.write().await = models.clone(); |
| 57 | + { |
| 58 | + let mut available_models_guard = self.available_models.write().await; |
| 59 | + *available_models_guard = self.build_available_models().await; |
| 60 | + } |
| 61 | + Ok(models) |
29 | 62 | } |
30 | 63 |
|
31 | 64 | pub fn construct_model_family(&self, model: &str, config: &Config) -> ModelFamily { |
32 | 65 | find_family_for_model(model).with_config_overrides(config) |
33 | 66 | } |
| 67 | + |
| 68 | + async fn build_available_models(&self) -> Vec<ModelPreset> { |
| 69 | + let mut available_models = self.remote_models.read().await.clone(); |
| 70 | + available_models.sort_by(|a, b| b.priority.cmp(&a.priority)); |
| 71 | + let mut model_presets: Vec<ModelPreset> = |
| 72 | + available_models.into_iter().map(Into::into).collect(); |
| 73 | + if let Some(default) = model_presets.first_mut() { |
| 74 | + default.is_default = true; |
| 75 | + } |
| 76 | + model_presets |
| 77 | + } |
| 78 | +} |
| 79 | + |
| 80 | +#[cfg(test)] |
| 81 | +mod tests { |
| 82 | + use super::*; |
| 83 | + use crate::CodexAuth; |
| 84 | + use crate::model_provider_info::WireApi; |
| 85 | + use codex_protocol::openai_models::ModelsResponse; |
| 86 | + use serde_json::json; |
| 87 | + use wiremock::Mock; |
| 88 | + use wiremock::MockServer; |
| 89 | + use wiremock::ResponseTemplate; |
| 90 | + use wiremock::matchers::method; |
| 91 | + use wiremock::matchers::path; |
| 92 | + |
| 93 | + fn remote_model(slug: &str, display: &str, priority: i32) -> ModelInfo { |
| 94 | + serde_json::from_value(json!({ |
| 95 | + "slug": slug, |
| 96 | + "display_name": display, |
| 97 | + "description": format!("{display} desc"), |
| 98 | + "default_reasoning_level": "medium", |
| 99 | + "supported_reasoning_levels": ["low", "medium"], |
| 100 | + "shell_type": "shell_command", |
| 101 | + "visibility": "list", |
| 102 | + "minimal_client_version": [0, 1, 0], |
| 103 | + "supported_in_api": true, |
| 104 | + "priority": priority |
| 105 | + })) |
| 106 | + .expect("valid model") |
| 107 | + } |
| 108 | + |
| 109 | + fn provider_for(base_url: String) -> ModelProviderInfo { |
| 110 | + ModelProviderInfo { |
| 111 | + name: "mock".into(), |
| 112 | + base_url: Some(base_url), |
| 113 | + env_key: None, |
| 114 | + env_key_instructions: None, |
| 115 | + experimental_bearer_token: None, |
| 116 | + wire_api: WireApi::Responses, |
| 117 | + query_params: None, |
| 118 | + http_headers: None, |
| 119 | + env_http_headers: None, |
| 120 | + request_max_retries: Some(0), |
| 121 | + stream_max_retries: Some(0), |
| 122 | + stream_idle_timeout_ms: Some(5_000), |
| 123 | + requires_openai_auth: false, |
| 124 | + } |
| 125 | + } |
| 126 | + |
| 127 | + #[tokio::test] |
| 128 | + async fn refresh_available_models_sorts_and_marks_default() { |
| 129 | + let server = MockServer::start().await; |
| 130 | + let remote_models = vec![ |
| 131 | + remote_model("priority-low", "Low", 1), |
| 132 | + remote_model("priority-high", "High", 10), |
| 133 | + ]; |
| 134 | + let response = ModelsResponse { |
| 135 | + models: remote_models.clone(), |
| 136 | + }; |
| 137 | + Mock::given(method("GET")) |
| 138 | + .and(path("/models")) |
| 139 | + .respond_with( |
| 140 | + ResponseTemplate::new(200) |
| 141 | + .insert_header("content-type", "application/json") |
| 142 | + .set_body_json(&response), |
| 143 | + ) |
| 144 | + .expect(1) |
| 145 | + .mount(&server) |
| 146 | + .await; |
| 147 | + |
| 148 | + let auth_manager = |
| 149 | + AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key")); |
| 150 | + let manager = ModelsManager::new(auth_manager); |
| 151 | + let provider = provider_for(server.uri()); |
| 152 | + |
| 153 | + let returned = manager |
| 154 | + .refresh_available_models(&provider) |
| 155 | + .await |
| 156 | + .expect("refresh succeeds"); |
| 157 | + |
| 158 | + assert_eq!(returned, remote_models); |
| 159 | + let cached_remote = manager.remote_models.read().await.clone(); |
| 160 | + assert_eq!(cached_remote, remote_models); |
| 161 | + |
| 162 | + let available = manager.available_models.read().await.clone(); |
| 163 | + assert_eq!(available.len(), 2); |
| 164 | + assert_eq!(available[0].model, "priority-high"); |
| 165 | + assert!( |
| 166 | + available[0].is_default, |
| 167 | + "highest priority should be default" |
| 168 | + ); |
| 169 | + assert_eq!(available[1].model, "priority-low"); |
| 170 | + assert!(!available[1].is_default); |
| 171 | + } |
34 | 172 | } |
0 commit comments