g1t/services/integrations/src/models.rs
| 1 | //! A workspace's own model providers: where each one's API is, how it takes |
| 2 | //! its key, and checking that the key works. The requests themselves go |
| 3 | //! through the model proxy, which never lets a sandbox see a 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` for Anthropic's API, with the version |
| 11 | /// for OpenAI's (`…/v1`, or Gemini's `…/v1beta/openai`). |
| 12 | pub fn base_url(provider: Provider, config: &ConnectionConfig) -> String { |
| 13 | let given = || config.base_url.as_deref().unwrap_or_default().trim_end_matches('/').to_owned(); |
| 14 | match provider { |
| 15 | Provider::AnthropicEndpoint => given().trim_end_matches("/v1").to_owned(), |
| 16 | Provider::Openai => "https://api.openai.com/v1".to_owned(), |
| 17 | Provider::Gemini => "https://generativelanguage.googleapis.com/v1beta/openai".to_owned(), |
| 18 | Provider::OpenaiEndpoint => given(), |
| 19 | _ => "https://api.anthropic.com".to_owned(), |
| 20 | } |
| 21 | } |
| 22 | |
| 23 | /// The header the key goes in. |
| 24 | pub fn auth_header(provider: Provider, config: &ConnectionConfig) -> String { |
| 25 | match provider { |
| 26 | Provider::Anthropic => "x-api-key".to_owned(), |
| 27 | Provider::AnthropicEndpoint => config.auth_header.clone().unwrap_or_else(|| "x-api-key".to_owned()), |
| 28 | _ => config.auth_header.clone().unwrap_or_else(|| "authorization".to_owned()), |
| 29 | } |
| 30 | } |
| 31 | |
| 32 | /// Whether a model id is one an agent could use: not embeddings, images, |
| 33 | /// speech or moderation. |
| 34 | fn for_chat(id: &str) -> bool { |
| 35 | let id = id.to_ascii_lowercase(); |
| 36 | !["embed", "tts", "whisper", "dall-e", "image", "moderation", "audio", "transcribe", "realtime", "search", "aqa", "imagen", "veo"] |
| 37 | .iter() |
| 38 | .any(|word| id.contains(word)) |
| 39 | } |
| 40 | |
| 41 | /// Asks the provider for its models with the key. `Ok` with what to say and |
| 42 | /// the models it offers; `Err` with what went wrong. |
| 43 | pub async fn test(provider: Provider, config: &ConnectionConfig, key: Option<&str>) -> Result<std::result::Result<(String, Vec<String>), String>> { |
| 44 | let base = base_url(provider, config); |
| 45 | let url = match provider.api() { |
| 46 | "anthropic" => format!("{base}/v1/models?limit=100"), |
| 47 | _ => format!("{base}/models"), |
| 48 | }; |
| 49 | let header = auth_header(provider, config); |
| 50 | let bearer = key.map(|key| format!("Bearer {key}")); |
| 51 | let mut headers = vec![("anthropic-version", "2023-06-01")]; |
| 52 | if let Some(key) = key { |
| 53 | if header == "authorization" { |
| 54 | headers.push(("authorization", bearer.as_deref().unwrap_or_default())); |
| 55 | } else { |
| 56 | headers.push(("x-api-key", key)); |
| 57 | } |
| 58 | } |
| 59 | let answer = http::send(Method::Get, &url, &headers, None).await?; |
| 60 | let system = provider.label(); |
| 61 | if answer.ok() { |
| 62 | let mut models: Vec<String> = answer.json()["data"] |
| 63 | .as_array() |
| 64 | .map(|data| { |
| 65 | data.iter() |
| 66 | .filter_map(|model| model["id"].as_str()) |
| 67 | // Gemini names models `models/gemini-…`. |
| 68 | .map(|id| id.trim_start_matches("models/").to_owned()) |
| 69 | .filter(|id| for_chat(id)) |
| 70 | .collect() |
| 71 | }) |
| 72 | .unwrap_or_default(); |
| 73 | models.sort(); |
| 74 | models.truncate(200); |
| 75 | let message = match models.len() { |
| 76 | 0 => format!("{system} accepted the key."), |
| 77 | count => format!("{system} accepted the key and offers {count} models."), |
| 78 | }; |
| 79 | return Ok(Ok((message, models))); |
| 80 | } |
| 81 | // A proxy may answer messages but not list models: it was reached, and |
| 82 | // whether the key works shows on the first run. |
| 83 | if answer.status == 404 && matches!(provider, Provider::AnthropicEndpoint | Provider::OpenaiEndpoint) { |
| 84 | return Ok(Ok(( |
| 85 | "Reached the endpoint. It does not list models, so the key will be checked on the first run.".to_owned(), |
| 86 | Vec::new(), |
| 87 | ))); |
| 88 | } |
| 89 | Ok(Err(answer.problem(system))) |
| 90 | } |
| 91 | |
| 92 | #[cfg(test)] |
| 93 | mod tests { |
| 94 | use super::*; |
| 95 | |
| 96 | #[test] |
| 97 | fn each_provider_has_its_address_and_header() { |
| 98 | let config = ConnectionConfig { |
| 99 | base_url: Some("https://llm.acme.dev/v1/".to_owned()), |
| 100 | ..ConnectionConfig::default() |
| 101 | }; |
| 102 | assert_eq!(base_url(Provider::Openai, &config), "https://api.openai.com/v1"); |
| 103 | assert_eq!(base_url(Provider::OpenaiEndpoint, &config), "https://llm.acme.dev/v1"); |
| 104 | assert_eq!(base_url(Provider::AnthropicEndpoint, &config), "https://llm.acme.dev"); |
| 105 | assert_eq!(auth_header(Provider::Gemini, &config), "authorization"); |
| 106 | assert_eq!(auth_header(Provider::Anthropic, &config), "x-api-key"); |
| 107 | } |
| 108 | |
| 109 | #[test] |
| 110 | fn only_models_that_can_chat_are_offered() { |
| 111 | assert!(for_chat("gpt-5")); |
| 112 | assert!(for_chat("gemini-2.5-pro")); |
| 113 | assert!(!for_chat("text-embedding-3-large")); |
| 114 | assert!(!for_chat("gpt-image-1")); |
| 115 | } |
| 116 | } |