g1t/services/integrations/src/models.rs
Pick any line to see why it is the way it is: the commit, the pull request and issue it came from, and what the agent was thinking.
| Models per workspace: several providers, routed by kind of work | 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. | |
| Integrations: your own model provider, alerts that open issues, tickets agents read | 4 | |
| 5 | use g1t_contracts::integrations::{ConnectionConfig, Provider}; | |
| 6 | use worker::{Method, Result}; | |
| 7 | ||
| 8 | use crate::http; | |
| 9 | ||
| Models per workspace: several providers, routed by kind of work | 10 | /// Where requests go: without `/v1` for Anthropic's API, with the version |
| 11 | /// for OpenAI's (`…/v1`, or Gemini's `…/v1beta/openai`). | |
| Integrations: your own model provider, alerts that open issues, tickets agents read | 12 | pub fn base_url(provider: Provider, config: &ConnectionConfig) -> String { |
| Models per workspace: several providers, routed by kind of work | 13 | let given = || config.base_url.as_deref().unwrap_or_default().trim_end_matches('/').to_owned(); |
| Integrations: your own model provider, alerts that open issues, tickets agents read | 14 | match provider { |
| Models per workspace: several providers, routed by kind of work | 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(), | |
| Integrations: your own model provider, alerts that open issues, tickets agents read | 19 | _ => "https://api.anthropic.com".to_owned(), |
| 20 | } | |
| 21 | } | |
| 22 | ||
| Models per workspace: several providers, routed by kind of work | 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. | |
| Model providers: gateway tokens for endpoints, tidier rows, and the docs | 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>> { | |
| Integrations: your own model provider, alerts that open issues, tickets agents read | 49 | let base = base_url(provider, config); |
| Models per workspace: several providers, routed by kind of work | 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); | |
| Integrations: your own model provider, alerts that open issues, tickets agents read | 55 | let bearer = key.map(|key| format!("Bearer {key}")); |
| 56 | let mut headers = vec![("anthropic-version", "2023-06-01")]; | |
| Models per workspace: several providers, routed by kind of work | 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 | } | |
| Integrations: your own model provider, alerts that open issues, tickets agents read | 63 | } |
| Model providers: gateway tokens for endpoints, tidier rows, and the docs | 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 | } | |
| Models per workspace: several providers, routed by kind of work | 68 | let answer = http::send(Method::Get, &url, &headers, None).await?; |
| 69 | let system = provider.label(); | |
| Integrations: your own model provider, alerts that open issues, tickets agents read | 70 | if answer.ok() { |
| Models per workspace: several providers, routed by kind of work | 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() { | |
| Integrations: your own model provider, alerts that open issues, tickets agents read | 85 | 0 => format!("{system} accepted the key."), |
| 86 | count => format!("{system} accepted the key and offers {count} models."), | |
| Models per workspace: several providers, routed by kind of work | 87 | }; |
| 88 | return Ok(Ok((message, models))); | |
| Integrations: your own model provider, alerts that open issues, tickets agents read | 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. | |
| Models per workspace: several providers, routed by kind of work | 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 | ))); | |
| Integrations: your own model provider, alerts that open issues, tickets agents read | 97 | } |
| 98 | Ok(Err(answer.problem(system))) | |
| 99 | } | |
| Models per workspace: several providers, routed by kind of work | 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 | } |