Skip to main content

codex_helper_core/
basellm_metadata.rs

1use std::collections::BTreeMap;
2use std::path::{Path, PathBuf};
3use std::time::Duration;
4
5use anyhow::{Context, Result};
6use reqwest::header::{
7    ETAG, HeaderMap, HeaderValue, IF_MODIFIED_SINCE, IF_NONE_MATCH, LAST_MODIFIED,
8};
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11use tokio::fs;
12use tracing::{debug, info, warn};
13
14use crate::config::proxy_home_dir;
15use crate::pricing::basellm_all_json_url;
16
17const BASELLM_RESPONSE_BODY_LIMIT: usize = 16 * 1024 * 1024;
18const BASELLM_CACHE_MAX_BYTES: u64 = 2 * 1024 * 1024;
19const BASELLM_SYNC_TIMEOUT_SECS: u64 = 30;
20
21#[derive(Debug, Clone, Default, Serialize, Deserialize)]
22#[serde(default)]
23pub struct BasellmMetadataCache {
24    pub source_url: String,
25    pub fetched_at_unix: i64,
26    pub etag: Option<String>,
27    pub last_modified: Option<String>,
28    pub openai_models: BTreeMap<String, BasellmOpenAiModelMetadata>,
29}
30
31#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
32#[serde(default)]
33pub struct BasellmOpenAiModelMetadata {
34    pub model_id: String,
35    pub display_name: Option<String>,
36    pub description: Option<String>,
37    pub context_window: Option<i64>,
38    pub max_context_window: Option<i64>,
39    pub input_modalities: Vec<String>,
40    pub reasoning: Option<bool>,
41    pub tool_call: Option<bool>,
42    pub structured_output: Option<bool>,
43    pub supports_fast_priority: bool,
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum BasellmMetadataSyncStatus {
48    Updated,
49    NotModified,
50    Unavailable,
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct BasellmMetadataSyncReport {
55    pub status: BasellmMetadataSyncStatus,
56    pub model_count: usize,
57}
58
59pub fn basellm_metadata_cache_path() -> PathBuf {
60    proxy_home_dir()
61        .join("model-metadata")
62        .join("basellm-openai-cache.json")
63}
64
65pub async fn sync_basellm_metadata_cache_background() {
66    if std::env::var("CODEX_HELPER_BASELLM_METADATA_SYNC")
67        .ok()
68        .is_some_and(|value| {
69            matches!(
70                value.trim().to_ascii_lowercase().as_str(),
71                "0" | "false" | "off" | "no"
72            )
73        })
74    {
75        debug!("BaseLLM model metadata background sync disabled by env");
76        return;
77    }
78
79    tokio::spawn(async {
80        match sync_basellm_metadata_cache(false).await {
81            Ok(report) => {
82                debug!(
83                    status = ?report.status,
84                    model_count = report.model_count,
85                    "BaseLLM model metadata background sync completed"
86                );
87            }
88            Err(err) => {
89                warn!("BaseLLM model metadata background sync failed: {err}");
90            }
91        }
92    });
93}
94
95pub async fn sync_basellm_metadata_cache(force: bool) -> Result<BasellmMetadataSyncReport> {
96    let cache_path = basellm_metadata_cache_path();
97    let existing = if force {
98        BasellmMetadataCache::default()
99    } else {
100        read_basellm_metadata_cache_from_path(&cache_path)
101            .await
102            .unwrap_or_default()
103    };
104
105    let client = reqwest::Client::builder()
106        .timeout(Duration::from_secs(BASELLM_SYNC_TIMEOUT_SECS))
107        .build()
108        .context("build BaseLLM metadata HTTP client")?;
109
110    let mut request = client
111        .get(basellm_all_json_url())
112        .header(reqwest::header::ACCEPT, "application/json");
113    if !force {
114        let headers = cache_headers(&existing);
115        if !headers.is_empty() {
116            request = request.headers(headers);
117        }
118    }
119
120    let response = request
121        .send()
122        .await
123        .context("request BaseLLM model metadata")?;
124
125    if response.status() == reqwest::StatusCode::NOT_MODIFIED && !force {
126        return Ok(BasellmMetadataSyncReport {
127            status: BasellmMetadataSyncStatus::NotModified,
128            model_count: existing.openai_models.len(),
129        });
130    }
131
132    if !response.status().is_success() {
133        return Ok(BasellmMetadataSyncReport {
134            status: BasellmMetadataSyncStatus::Unavailable,
135            model_count: existing.openai_models.len(),
136        });
137    }
138
139    let etag = header_to_string(response.headers(), ETAG);
140    let last_modified = header_to_string(response.headers(), LAST_MODIFIED);
141    let body = read_limited_text(response, BASELLM_RESPONSE_BODY_LIMIT)
142        .await
143        .context("read BaseLLM model metadata response")?;
144    let mut cache =
145        parse_basellm_openai_metadata_json(&body).context("parse BaseLLM OpenAI model metadata")?;
146    cache.source_url = basellm_all_json_url().to_string();
147    cache.fetched_at_unix = unix_now();
148    cache.etag = etag;
149    cache.last_modified = last_modified;
150
151    write_basellm_metadata_cache_to_path(&cache_path, &cache)
152        .await
153        .with_context(|| format!("write BaseLLM model metadata cache {:?}", cache_path))?;
154
155    info!(
156        model_count = cache.openai_models.len(),
157        cache_path = %cache_path.display(),
158        "BaseLLM OpenAI model metadata cache updated"
159    );
160
161    Ok(BasellmMetadataSyncReport {
162        status: BasellmMetadataSyncStatus::Updated,
163        model_count: cache.openai_models.len(),
164    })
165}
166
167pub fn load_cached_openai_model_metadata(model_id: &str) -> Option<BasellmOpenAiModelMetadata> {
168    let cache = load_cached_basellm_metadata_cache()?;
169    cache
170        .openai_models
171        .get(&normalize_model_id(model_id))
172        .cloned()
173}
174
175pub fn load_cached_basellm_metadata_cache() -> Option<BasellmMetadataCache> {
176    read_basellm_metadata_cache_from_path_blocking(&basellm_metadata_cache_path()).ok()
177}
178
179pub fn parse_basellm_openai_metadata_json(text: &str) -> Result<BasellmMetadataCache> {
180    let root: Value = serde_json::from_str(text).context("invalid BaseLLM metadata JSON")?;
181    let Some(openai_models) = root
182        .get("openai")
183        .and_then(|provider| provider.get("models"))
184        .and_then(Value::as_object)
185    else {
186        return Ok(BasellmMetadataCache::default());
187    };
188
189    let mut models = BTreeMap::new();
190    for (model_id, model_value) in openai_models {
191        let normalized = normalize_model_id(model_id);
192        if normalized.is_empty() {
193            continue;
194        }
195
196        let display_name = model_value
197            .get("name")
198            .and_then(json_scalar_to_string)
199            .or_else(|| {
200                model_value
201                    .get("display_name")
202                    .and_then(json_scalar_to_string)
203            })
204            .filter(|value| !value.eq_ignore_ascii_case(model_id));
205        let description = model_value
206            .get("description")
207            .and_then(json_scalar_to_string);
208        let input_modalities = parse_input_modalities(model_value);
209        let context_window = model_value
210            .get("limit")
211            .and_then(|limit| limit.get("input").or_else(|| limit.get("context")))
212            .and_then(json_i64);
213        let max_context_window = model_value
214            .get("limit")
215            .and_then(|limit| limit.get("context").or_else(|| limit.get("input")))
216            .and_then(json_i64);
217
218        models.insert(
219            normalized.clone(),
220            BasellmOpenAiModelMetadata {
221                model_id: model_id.to_string(),
222                display_name,
223                description,
224                context_window,
225                max_context_window,
226                input_modalities,
227                reasoning: model_value.get("reasoning").and_then(Value::as_bool),
228                tool_call: model_value.get("tool_call").and_then(Value::as_bool),
229                structured_output: model_value
230                    .get("structured_output")
231                    .and_then(Value::as_bool),
232                supports_fast_priority: basellm_model_supports_fast_priority(model_value),
233            },
234        );
235    }
236
237    Ok(BasellmMetadataCache {
238        source_url: basellm_all_json_url().to_string(),
239        openai_models: models,
240        ..BasellmMetadataCache::default()
241    })
242}
243
244async fn read_basellm_metadata_cache_from_path(path: &Path) -> Result<BasellmMetadataCache> {
245    let metadata = fs::metadata(path).await?;
246    if metadata.len() > BASELLM_CACHE_MAX_BYTES {
247        anyhow::bail!("BaseLLM metadata cache is too large");
248    }
249    let bytes = fs::read(path).await?;
250    Ok(serde_json::from_slice(&bytes)?)
251}
252
253fn read_basellm_metadata_cache_from_path_blocking(path: &Path) -> Result<BasellmMetadataCache> {
254    let metadata = std::fs::metadata(path)?;
255    if metadata.len() > BASELLM_CACHE_MAX_BYTES {
256        anyhow::bail!("BaseLLM metadata cache is too large");
257    }
258    let bytes = std::fs::read(path)?;
259    Ok(serde_json::from_slice(&bytes)?)
260}
261
262async fn write_basellm_metadata_cache_to_path(
263    path: &Path,
264    cache: &BasellmMetadataCache,
265) -> Result<()> {
266    let bytes = serde_json::to_vec_pretty(cache)?;
267    if bytes.len() as u64 > BASELLM_CACHE_MAX_BYTES {
268        anyhow::bail!("BaseLLM metadata cache payload is too large");
269    }
270    if let Some(parent) = path.parent() {
271        fs::create_dir_all(parent).await?;
272    }
273    let tmp_path = path.with_extension("json.tmp");
274    fs::write(&tmp_path, bytes).await?;
275    if fs::try_exists(path).await.unwrap_or(false) {
276        let _ = fs::remove_file(path).await;
277    }
278    fs::rename(&tmp_path, path).await?;
279    Ok(())
280}
281
282async fn read_limited_text(response: reqwest::Response, limit: usize) -> Result<String> {
283    let bytes = response.bytes().await?;
284    if bytes.len() > limit {
285        anyhow::bail!("response body exceeds {limit} bytes");
286    }
287    Ok(String::from_utf8(bytes.to_vec())?)
288}
289
290fn cache_headers(cache: &BasellmMetadataCache) -> HeaderMap {
291    let mut headers = HeaderMap::new();
292    if let Some(etag) = cache.etag.as_deref().and_then(header_value) {
293        headers.insert(IF_NONE_MATCH, etag);
294    }
295    if let Some(last_modified) = cache.last_modified.as_deref().and_then(header_value) {
296        headers.insert(IF_MODIFIED_SINCE, last_modified);
297    }
298    headers
299}
300
301fn header_value(value: &str) -> Option<HeaderValue> {
302    HeaderValue::from_str(value).ok()
303}
304
305fn header_to_string(headers: &HeaderMap, name: reqwest::header::HeaderName) -> Option<String> {
306    headers
307        .get(name)
308        .and_then(|value| value.to_str().ok())
309        .map(ToOwned::to_owned)
310}
311
312fn parse_input_modalities(model_value: &Value) -> Vec<String> {
313    let mut out = Vec::new();
314    if let Some(items) = model_value
315        .get("modalities")
316        .and_then(|modalities| modalities.get("input"))
317        .and_then(Value::as_array)
318    {
319        for item in items {
320            let Some(modality) = item
321                .as_str()
322                .map(str::trim)
323                .filter(|value| !value.is_empty())
324            else {
325                continue;
326            };
327            let modality = match modality.to_ascii_lowercase().as_str() {
328                "text" => "text",
329                "image" => "image",
330                // Codex catalog currently only models text/image input. Treat PDFs as attachments
331                // rather than adding a modality unknown to older Codex clients.
332                "pdf" => continue,
333                _ => continue,
334            };
335            if !out.iter().any(|existing| existing == modality) {
336                out.push(modality.to_string());
337            }
338        }
339    }
340    out
341}
342
343fn basellm_model_supports_fast_priority(model_value: &Value) -> bool {
344    model_value
345        .get("experimental")
346        .and_then(|experimental| experimental.get("modes"))
347        .and_then(|modes| modes.get("fast"))
348        .and_then(|fast| fast.get("provider"))
349        .and_then(|provider| provider.get("body"))
350        .and_then(|body| body.get("service_tier"))
351        .and_then(Value::as_str)
352        .is_some_and(|service_tier| service_tier.eq_ignore_ascii_case("priority"))
353}
354
355fn json_scalar_to_string(value: &Value) -> Option<String> {
356    match value {
357        Value::Number(number) => Some(number.to_string()),
358        Value::String(text) => {
359            let text = text.trim();
360            (!text.is_empty()).then(|| text.to_string())
361        }
362        _ => None,
363    }
364}
365
366fn json_i64(value: &Value) -> Option<i64> {
367    value
368        .as_i64()
369        .or_else(|| value.as_str()?.trim().parse::<i64>().ok())
370}
371
372fn normalize_model_id(value: &str) -> String {
373    value.trim().to_ascii_lowercase()
374}
375
376fn unix_now() -> i64 {
377    std::time::SystemTime::now()
378        .duration_since(std::time::UNIX_EPOCH)
379        .map(|duration| duration.as_secs() as i64)
380        .unwrap_or(0)
381}
382
383#[cfg(test)]
384mod tests {
385    use super::*;
386
387    #[test]
388    fn parses_fast_priority_and_capabilities_from_basellm_openai_models() {
389        let text = r#"
390{
391  "openai": {
392    "models": {
393      "gpt-test": {
394        "name": "GPT Test",
395        "description": "Test model",
396        "limit": { "context": 1050000, "input": 922000, "output": 128000 },
397        "modalities": { "input": ["text", "image", "pdf"], "output": ["text"] },
398        "reasoning": true,
399        "tool_call": true,
400        "structured_output": true,
401        "experimental": {
402          "modes": {
403            "fast": {
404              "cost": { "input": 10, "output": 20 },
405              "provider": { "body": { "service_tier": "priority" } }
406            }
407          }
408        }
409      }
410    }
411  }
412}
413"#;
414
415        let cache = parse_basellm_openai_metadata_json(text).expect("parse");
416        let model = cache.openai_models.get("gpt-test").expect("gpt-test");
417
418        assert_eq!(model.display_name.as_deref(), Some("GPT Test"));
419        assert_eq!(model.description.as_deref(), Some("Test model"));
420        assert_eq!(model.context_window, Some(922_000));
421        assert_eq!(model.max_context_window, Some(1_050_000));
422        assert_eq!(model.input_modalities, vec!["text", "image"]);
423        assert_eq!(model.reasoning, Some(true));
424        assert_eq!(model.tool_call, Some(true));
425        assert_eq!(model.structured_output, Some(true));
426        assert!(model.supports_fast_priority);
427    }
428}