artifactize_core/llm/
models.rs

1use std::{path::Path, time::Duration};
2
3use rig_core::providers::{anthropic, openai};
4use serde::Serialize;
5
6use crate::config::Backend;
7
8/// Model discovery is interactive metadata, not a review; stop a stalled listing
9/// after thirty seconds for both API-key backends and the Codex account endpoint.
10pub(super) const LIST_TIMEOUT: Duration = Duration::from_secs(30);
11
12#[derive(Serialize)]
13pub struct ListedModel {
14    pub slug: crate::config::ModelId,
15    pub display_name: String,
16}
17
18#[derive(Serialize)]
19pub struct Listing {
20    pub backend: Backend,
21    pub models: Vec<ListedModel>,
22}
23
24/// `state` and `repo` locate Codex credentials; the API-key backends ignore them.
25pub async fn list(
26    backend: Backend,
27    state: Option<&Path>,
28    repo: Option<&Path>,
29) -> Result<Listing, String> {
30    let base = super::base_url(backend)?;
31    let models = if backend == Backend::Codex {
32        super::codex::models(&base, state, repo).await?
33    } else {
34        list_at(backend, &super::api_key(backend, "to list models")?, &base).await?
35    };
36    Ok(Listing { backend, models })
37}
38
39async fn list_at(backend: Backend, key: &str, base: &str) -> Result<Vec<ListedModel>, String> {
40    let http = rig_reqwest::ReqwestClient::from(super::http_client()?);
41    let request = async {
42        match backend {
43            Backend::Openai => {
44                openai::OpenAIConfig::new(key.to_owned())
45                    .with_base_url(base)
46                    .connect(http)
47                    .list_models()
48                    .await
49            }
50            Backend::Anthropic => {
51                anthropic::AnthropicConfig::new(key.to_owned())
52                    .with_base_url(base)
53                    .connect(http)
54                    .list_models()
55                    .await
56            }
57            Backend::Codex => unreachable!("Codex lists its own models"),
58        }
59    };
60    let models = tokio::time::timeout(LIST_TIMEOUT, request)
61        .await
62        .map_err(|_| "Model listing timed out.".to_owned())?
63        .map_err(|error| {
64            let text = super::diagnostic_hiding(&error, &[key]);
65            if text == crate::secret::WITHHELD {
66                format!("Model listing failed: {text}")
67            } else {
68                text
69            }
70        })?;
71    // A listed model without a usable id cannot be selected, so it is left out.
72    Ok(models
73        .into_iter()
74        .filter_map(|model| {
75            Some(ListedModel {
76                slug: model.id.parse().ok()?,
77                display_name: model.name.unwrap_or(model.id),
78            })
79        })
80        .collect())
81}
82
83#[cfg(test)]
84mod tests {
85    use rig_core::test_utils::MockHttpResponse;
86    use serde_json::json;
87
88    use super::*;
89    use crate::llm::tests::Server;
90
91    #[tokio::test]
92    async fn api_key_models_use_provider_headers_and_follow_anthropic_pages() {
93        let openai = Server::new(vec![MockHttpResponse::success(
94            json!({"data":[
95                {"id":"z-model"}, {"id":"a-model", "name":"A model"}
96            ]})
97            .to_string(),
98        )])
99        .await;
100        let models = list_at(Backend::Openai, "openai-key", &openai.base)
101            .await
102            .unwrap();
103        assert_eq!(
104            serde_json::to_value(models).unwrap(),
105            json!([
106                {"slug":"z-model","display_name":"z-model"},
107                {"slug":"a-model","display_name":"A model"}
108            ])
109        );
110        let requests = openai.requests();
111        assert_eq!(requests[0].line, "GET /v1/models HTTP/1.1");
112        assert_eq!(requests[0].headers["authorization"], "Bearer openai-key");
113        assert!(!requests[0].headers.contains_key("x-api-key"));
114        assert!(requests[0].body.is_null());
115
116        let anthropic = Server::new(vec![
117            MockHttpResponse::success(
118                json!({
119                    "data":[{"id":"first", "display_name":"First"}],
120                    "has_more":true,
121                    "last_id":"first",
122                })
123                .to_string(),
124            ),
125            MockHttpResponse::success(
126                json!({"data":[{"id":"second", "display_name":"Second"}],"has_more":false})
127                    .to_string(),
128            ),
129        ])
130        .await;
131        let models = list_at(Backend::Anthropic, "anthropic-key", &anthropic.base)
132            .await
133            .unwrap();
134        assert_eq!(
135            serde_json::to_value(models).unwrap(),
136            json!([
137                {"slug":"first","display_name":"First"},
138                {"slug":"second","display_name":"Second"}
139            ])
140        );
141        let requests = anthropic.requests();
142        assert_eq!(requests[0].line, "GET /v1/models HTTP/1.1");
143        assert_eq!(requests[1].line, "GET /v1/models?after_id=first HTTP/1.1");
144        assert_eq!(requests[0].headers["x-api-key"], "anthropic-key");
145        assert!(requests[0].headers.contains_key("anthropic-version"));
146        assert!(!requests[0].headers.contains_key("authorization"));
147    }
148
149    #[tokio::test]
150    async fn model_listing_errors_do_not_expose_api_keys() {
151        for backend in [Backend::Openai, Backend::Anthropic] {
152            let server = Server::new(vec![MockHttpResponse::error(
153                rig_core::http_client::StatusCode::UNAUTHORIZED,
154                json!({"error":{"message":"Invalid secret-key", "type":"authentication_error"}})
155                    .to_string(),
156            )])
157            .await;
158            let error = list_at(backend, "secret-key", &server.base)
159                .await
160                .err()
161                .unwrap();
162            // The provider's text held the key: withheld whole, as every server text is.
163            assert_eq!(
164                error,
165                format!("Model listing failed: {}", crate::secret::WITHHELD)
166            );
167        }
168        // Text without the key still explains the failure.
169        let server = Server::new(vec![MockHttpResponse::error(
170            rig_core::http_client::StatusCode::UNAUTHORIZED,
171            json!({"error":{"message":"Invalid API key\n", "type":"authentication_error"}})
172                .to_string(),
173        )])
174        .await;
175        let error = list_at(Backend::Openai, "secret-key-value", &server.base)
176            .await
177            .err()
178            .unwrap();
179        assert_eq!(error, "Invalid API key");
180    }
181}
artifactize_core/llm/models.rs

artifactize_core/llm/
models.rs
