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( |
| 44 | provider: Provider, |
| 45 | config: &ConnectionConfig, |
| 46 | key: Option<&str>, |
| 47 | gateway_token: Option<&str>, |
| 48 | ) -> Result<std::result::Result<(String, Vec<String>), String>> { |
| 49 | let base = base_url(provider, config); |
| 50 | let url = match provider.api() { |
| 51 | "anthropic" => format!("{base}/v1/models?limit=100"), |
| 52 | _ => format!("{base}/models"), |
| 53 | }; |
| 54 | let header = auth_header(provider, config); |
| 55 | let bearer = key.map(|key| format!("Bearer {key}")); |
| 56 | let mut headers = vec![("anthropic-version", "2023-06-01")]; |
| 57 | if let Some(key) = key { |
| 58 | if header == "authorization" { |
| 59 | headers.push(("authorization", bearer.as_deref().unwrap_or_default())); |
| 60 | } else { |
| 61 | headers.push(("x-api-key", key)); |
| 62 | } |
| 63 | } |
| 64 | let gateway = gateway_token.map(|token| format!("Bearer {token}")); |
| 65 | if let Some(gateway) = gateway.as_deref() { |
| 66 | headers.push(("cf-aig-authorization", gateway)); |
| 67 | } |
| 68 | let answer = http::send(Method::Get, &url, &headers, None).await?; |
| 69 | let system = provider.label(); |
| 70 | if answer.ok() { |
| 71 | let mut models: Vec<String> = answer.json()["data"] |
| 72 | .as_array() |
| 73 | .map(|data| { |
| 74 | data.iter() |
| 75 | .filter_map(|model| model["id"].as_str()) |
| 76 | // Gemini names models `models/gemini-…`. |
| 77 | .map(|id| id.trim_start_matches("models/").to_owned()) |
| 78 | .filter(|id| for_chat(id)) |
| 79 | .collect() |
| 80 | }) |
| 81 | .unwrap_or_default(); |
| 82 | models.sort(); |
| 83 | models.truncate(200); |
| 84 | let message = match models.len() { |
| 85 | 0 => format!("{system} accepted the key."), |
| 86 | count => format!("{system} accepted the key and offers {count} models."), |
| 87 | }; |
| 88 | return Ok(Ok((message, models))); |
| 89 | } |
| 90 | // A proxy may answer messages but not list models: it was reached, and |
| 91 | // whether the key works shows on the first run. |
| 92 | if answer.status == 404 && matches!(provider, Provider::AnthropicEndpoint | Provider::OpenaiEndpoint) { |
| 93 | return Ok(Ok(( |
| 94 | "Reached the endpoint. It does not list models, so the key will be checked on the first run.".to_owned(), |
| 95 | Vec::new(), |
| 96 | ))); |
| 97 | } |
| 98 | Ok(Err(answer.problem(system))) |
| 99 | } |
| 100 | |
| 101 | #[cfg(test)] |
| 102 | mod tests { |
| 103 | use super::*; |
| 104 | |
| 105 | #[test] |
| 106 | fn each_provider_has_its_address_and_header() { |
| 107 | let config = ConnectionConfig { |
| 108 | base_url: Some("https://llm.acme.dev/v1/".to_owned()), |
| 109 | ..ConnectionConfig::default() |
| 110 | }; |
| 111 | assert_eq!(base_url(Provider::Openai, &config), "https://api.openai.com/v1"); |
| 112 | assert_eq!(base_url(Provider::OpenaiEndpoint, &config), "https://llm.acme.dev/v1"); |
| 113 | assert_eq!(base_url(Provider::AnthropicEndpoint, &config), "https://llm.acme.dev"); |
| 114 | assert_eq!(auth_header(Provider::Gemini, &config), "authorization"); |
| 115 | assert_eq!(auth_header(Provider::Anthropic, &config), "x-api-key"); |
| 116 | } |
| 117 | |
| 118 | #[test] |
| 119 | fn only_models_that_can_chat_are_offered() { |
| 120 | assert!(for_chat("gpt-5")); |
| 121 | assert!(for_chat("gemini-2.5-pro")); |
| 122 | assert!(!for_chat("text-embedding-3-large")); |
| 123 | assert!(!for_chat("gpt-image-1")); |
| 124 | } |
| 125 | } |