pr_01m47d15m3e54sn21z27rpy5n9/services/integrations/src/models.rs
| 1 | //! A workspace's own model provider: checking that its key works. The |
| 2 | //! requests themselves go through the model proxy, which never lets a |
| 3 | //! sandbox see the key. |
| 4 | |
| 5 | use g1t_contracts::integrations::{ConnectionConfig, Provider}; |
| 6 | use worker::{Method, Result}; |
| 7 | |
| 8 | use crate::http; |
| 9 | |
| 10 | /// Where requests go, without `/v1`. |
| 11 | pub fn base_url(provider: Provider, config: &ConnectionConfig) -> String { |
| 12 | match provider { |
| 13 | Provider::AnthropicEndpoint => config.base_url.as_deref().unwrap_or_default().trim_end_matches('/').trim_end_matches("/v1").to_owned(), |
| 14 | _ => "https://api.anthropic.com".to_owned(), |
| 15 | } |
| 16 | } |
| 17 | |
| 18 | pub async fn test(provider: Provider, config: &ConnectionConfig, key: Option<&str>) -> Result<std::result::Result<String, String>> { |
| 19 | let base = base_url(provider, config); |
| 20 | let bearer = key.map(|key| format!("Bearer {key}")); |
| 21 | let mut headers = vec![("anthropic-version", "2023-06-01")]; |
| 22 | match (key, config.auth_header.as_deref()) { |
| 23 | (Some(_), Some("authorization")) => headers.push(("authorization", bearer.as_deref().unwrap_or_default())), |
| 24 | (Some(key), _) => headers.push(("x-api-key", key)), |
| 25 | (None, _) => {} |
| 26 | } |
| 27 | let answer = http::send(Method::Get, &format!("{base}/v1/models"), &headers, None).await?; |
| 28 | let system = if provider == Provider::Anthropic { "Anthropic" } else { "The endpoint" }; |
| 29 | if answer.ok() { |
| 30 | let models = answer.json()["data"].as_array().map_or(0, Vec::len); |
| 31 | return Ok(Ok(match models { |
| 32 | 0 => format!("{system} accepted the key."), |
| 33 | count => format!("{system} accepted the key and offers {count} models."), |
| 34 | })); |
| 35 | } |
| 36 | // A proxy may answer messages but not list models: it was reached, and |
| 37 | // whether the key works shows on the first run. |
| 38 | if answer.status == 404 && provider == Provider::AnthropicEndpoint { |
| 39 | return Ok(Ok("Reached the endpoint. It does not list models, so the key will be checked on the first run.".to_owned())); |
| 40 | } |
| 41 | Ok(Err(answer.problem(system))) |
| 42 | } |