From 0797e89aa00e57b4e162df32d7b0a041ad6db71e Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Thu, 22 Jan 2026 18:28:51 +0100 Subject: [PATCH] feat(ai): handle google vertex for claude models + base url overrides (#7654) * fix hardcoded gemini url * allow overriding any provider url * handle vertex * same for chat proxy * fix for chat --- backend/windmill-api/src/ai.rs | 81 ++++++++++++- backend/windmill-common/src/ai_providers.rs | 24 ++-- .../src/ai/providers/anthropic.rs | 107 ++++++++++++++---- .../src/ai/providers/google_ai.rs | 8 +- .../windmill-worker/src/ai/query_builder.rs | 7 +- backend/windmill-worker/src/ai/types.rs | 15 +++ 6 files changed, 198 insertions(+), 44 deletions(-) diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index a8d7a9e122..84f03c12b8 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -125,6 +125,15 @@ struct AIOAuthResource { user: Option, } +/// Platform for Anthropic API +#[derive(Deserialize, Debug, Clone, Default, PartialEq)] +#[serde(rename_all = "snake_case")] +enum AnthropicPlatform { + #[default] + Standard, + GoogleVertexAi, +} + #[derive(Deserialize, Debug)] struct AIStandardResource { #[serde(alias = "baseUrl")] @@ -137,6 +146,9 @@ struct AIStandardResource { aws_access_key_id: Option, #[serde(alias = "awsSecretAccessKey")] aws_secret_access_key: Option, + /// Platform for Anthropic API (standard or google_vertex_ai) + #[serde(default)] + platform: AnthropicPlatform, } #[derive(Deserialize, Debug)] @@ -161,6 +173,7 @@ struct AIRequestConfig { pub region: Option, pub aws_access_key_id: Option, pub aws_secret_access_key: Option, + pub platform: AnthropicPlatform, } impl AIRequestConfig { @@ -179,9 +192,11 @@ impl AIRequestConfig { region, aws_access_key_id, aws_secret_access_key, + platform, ) = match resource { AIResource::Standard(resource) => { let region = resource.region.clone(); + let platform = resource.platform.clone(); let base_url = provider .get_base_url(resource.base_url, resource.region, db) .await?; @@ -216,6 +231,7 @@ impl AIRequestConfig { region, aws_access_key_id, aws_secret_access_key, + platform, ) } AIResource::OAuth(resource) => { @@ -227,7 +243,17 @@ impl AIRequestConfig { let token = Self::get_token_using_oauth(resource, db, w_id).await?; let base_url = provider.get_base_url(None, None, db).await?; - (None, Some(token), None, base_url, user, None, None, None) + ( + None, + Some(token), + None, + base_url, + user, + None, + None, + None, + AnthropicPlatform::Standard, + ) } }; @@ -240,6 +266,7 @@ impl AIRequestConfig { region, aws_access_key_id, aws_secret_access_key, + platform, }) } @@ -294,8 +321,18 @@ impl AIRequestConfig { let is_azure = provider.is_azure_openai(base_url); let is_anthropic = matches!(provider, AIProvider::Anthropic); + let is_anthropic_vertex = is_anthropic && self.platform == AnthropicPlatform::GoogleVertexAi; let is_anthropic_sdk = headers.get("X-Anthropic-SDK").is_some(); let is_bedrock = matches!(provider, AIProvider::AWSBedrock); + let is_google_ai = matches!(provider, AIProvider::GoogleAI); + + // GoogleAI uses OpenAI-compatible endpoint in the proxy (for the chat), but not for the ai agent + let base_url = if is_google_ai { + format!("{}/openai", base_url) + } else { + base_url.to_string() + }; + let base_url = base_url.as_str(); // Check if using IAM credentials for Bedrock (instead of bearer token) let use_iam_auth = @@ -317,6 +354,10 @@ impl AIRequestConfig { let bedrock_base_url = base_url.replace("bedrock-runtime.", "bedrock."); let bedrock_url = format!("{}/{}", bedrock_base_url, path); (bedrock_url, body) + } else if is_anthropic_vertex && method != Method::GET { + let (model, transformed_body) = transform_anthropic_for_vertex(&body)?; + let vertex_url = format!("{}/{}:streamRawPredict", base_url, model); + (vertex_url, transformed_body) } else if is_azure { let azure_url = AIProvider::build_azure_openai_url(base_url, path); (azure_url, body) @@ -336,7 +377,12 @@ impl AIRequestConfig { .header("content-type", "application/json"); for (header_name, header_value) in headers.iter() { + // Forward anthropic-* headers, but skip anthropic-version for Vertex AI + // (Vertex AI requires anthropic_version in the request body, not as a header) if header_name.to_string().starts_with("anthropic-") { + if is_anthropic_vertex && header_name.as_str() == "anthropic-version" { + continue; + } request = request.header(header_name, header_value); } } @@ -360,13 +406,14 @@ impl AIRequestConfig { } } else { // For non-IAM auth, use bearer token or API key - if let Some(api_key) = self.api_key { + if let Some(api_key) = self.api_key.clone() { if is_azure { request = request.header("api-key", api_key.clone()) } else { request = request.header("authorization", format!("Bearer {}", api_key.clone())) } - if is_anthropic { + // For standard Anthropic API, also add X-API-Key header + if is_anthropic && !is_anthropic_vertex { request = request.header("X-API-Key", api_key); } } @@ -438,6 +485,34 @@ pub struct AIConfig { pub max_tokens_per_model: Option>, } +/// Anthropic API version for Google Vertex AI +const ANTHROPIC_VERSION_VERTEX: &str = "vertex-2023-10-16"; + +/// Transforms an Anthropic request for Google Vertex AI: +/// - Extracts the model from the body (needed for the URL) +/// - Adds anthropic_version to the body +fn transform_anthropic_for_vertex(body: &Bytes) -> Result<(String, Bytes)> { + let mut json_body: HashMap = serde_json::from_slice(body) + .map_err(|e| Error::internal_err(format!("Failed to parse Anthropic request: {}", e)))?; + + // Extract and remove model from body + let model = json_body + .remove("model") + .and_then(|v| v.as_str().map(|s| s.to_string())) + .ok_or_else(|| Error::BadRequest("Missing 'model' field in Anthropic request".to_string()))?; + + // Add anthropic_version to body (required for Vertex AI) + json_body.insert( + "anthropic_version".to_string(), + serde_json::Value::String(ANTHROPIC_VERSION_VERTEX.to_string()), + ); + + let transformed_body = serde_json::to_vec(&json_body) + .map_err(|e| Error::internal_err(format!("Failed to serialize Vertex request: {}", e)))?; + + Ok((model, Bytes::from(transformed_body))) +} + // FIM (Fill-in-the-Middle) simulation for providers that don't support native FIM #[derive(Deserialize, Debug)] struct FimRequest { diff --git a/backend/windmill-common/src/ai_providers.rs b/backend/windmill-common/src/ai_providers.rs index b8c1aaf8fe..34153195ac 100644 --- a/backend/windmill-common/src/ai_providers.rs +++ b/backend/windmill-common/src/ai_providers.rs @@ -11,6 +11,7 @@ lazy_static::lazy_static! { } pub const OPENAI_BASE_URL: &str = "https://api.openai.com/v1"; +pub const GOOGLE_AI_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta"; #[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Hash, Clone)] #[serde(rename_all = "lowercase")] @@ -38,6 +39,11 @@ impl AIProvider { region: Option, db: &DB, ) -> Result { + // If a base URL is provided in the resource, use it + if let Some(base_url) = resource_base_url { + return Ok(base_url); + } + match self { AIProvider::OpenAI => { // Check for Azure base path override @@ -62,28 +68,20 @@ impl AIProvider { Ok(azure_base_path.unwrap_or("https://api.openai.com/v1".to_string())) } AIProvider::DeepSeek => Ok("https://api.deepseek.com/v1".to_string()), - AIProvider::GoogleAI => { - Ok("https://generativelanguage.googleapis.com/v1beta/openai".to_string()) - } + AIProvider::GoogleAI => Ok(GOOGLE_AI_BASE_URL.to_string()), AIProvider::Groq => Ok("https://api.groq.com/openai/v1".to_string()), AIProvider::OpenRouter => Ok("https://openrouter.ai/api/v1".to_string()), AIProvider::TogetherAI => Ok("https://api.together.xyz/v1".to_string()), AIProvider::Anthropic => Ok("https://api.anthropic.com/v1".to_string()), AIProvider::Mistral => Ok("https://api.mistral.ai/v1".to_string()), - p @ (AIProvider::CustomAI | AIProvider::AzureOpenAI) => { - if let Some(base_url) = resource_base_url { - Ok(base_url) - } else { - Err(Error::BadRequest(format!( - "{:?} provider requires a base URL in the resource", - p - ))) - } - } AIProvider::AWSBedrock => Ok(format!( "https://bedrock-runtime.{}.amazonaws.com", region.unwrap_or_else(|| "us-east-1".to_string()) )), + AIProvider::CustomAI | AIProvider::AzureOpenAI => Err(Error::BadRequest(format!( + "{:?} provider requires a base URL in the resource", + self + ))), } } diff --git a/backend/windmill-worker/src/ai/providers/anthropic.rs b/backend/windmill-worker/src/ai/providers/anthropic.rs index a562673af9..a1fb4b8c0b 100644 --- a/backend/windmill-worker/src/ai/providers/anthropic.rs +++ b/backend/windmill-worker/src/ai/providers/anthropic.rs @@ -11,6 +11,11 @@ use crate::ai::{ utils::{extract_text_content, parse_data_url, should_use_structured_output_tool}, }; +/// Anthropic API version for standard API +const ANTHROPIC_VERSION_STANDARD: &str = "2023-06-01"; +/// Anthropic API version for Google Vertex AI +const ANTHROPIC_VERSION_VERTEX: &str = "vertex-2023-10-16"; + /// Custom tool for Anthropic native API (flat structure with type: "custom") #[derive(Serialize, Debug)] pub struct AnthropicCustomTool { @@ -96,7 +101,7 @@ pub struct AnthropicMessage { pub content: Vec, } -/// Anthropic-specific request structure +/// Anthropic-specific request structure for standard API #[derive(Serialize)] pub struct AnthropicRequest<'a> { pub model: &'a str, @@ -114,6 +119,27 @@ pub struct AnthropicRequest<'a> { pub stream: bool, } +/// Anthropic request structure for Google Vertex AI +/// Key differences from standard API: +/// - No model field (model is specified in the URL) +/// - anthropic_version is in the body instead of a header +#[derive(Serialize)] +pub struct AnthropicVertexRequest { + pub anthropic_version: &'static str, + #[serde(skip_serializing_if = "Option::is_none")] + pub system: Option>, + pub messages: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub tools: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_choice: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, + pub stream: bool, +} + /// Convert OpenAI-format messages to Anthropic native format fn convert_messages_to_anthropic(messages: &[OpenAIMessage]) -> Vec { let mut result: Vec = Vec::new(); @@ -298,11 +324,16 @@ pub struct AnthropicResponse { pub struct AnthropicQueryBuilder { #[allow(dead_code)] provider_kind: AIProvider, + platform: AnthropicPlatform, } impl AnthropicQueryBuilder { - pub fn new(provider_kind: AIProvider) -> Self { - Self { provider_kind } + pub fn new(provider_kind: AIProvider, platform: AnthropicPlatform) -> Self { + Self { provider_kind, platform } + } + + fn is_vertex(&self) -> bool { + self.platform == AnthropicPlatform::GoogleVertexAi } async fn build_text_request( @@ -362,19 +393,39 @@ impl AnthropicQueryBuilder { None }; - let request = AnthropicRequest { - model: args.model, - system, - messages: anthropic_messages, - tools: if tools.is_empty() { None } else { Some(tools) }, - tool_choice, - temperature: args.temperature, - max_tokens: Some(args.max_tokens.unwrap_or(64000)), - stream: true, - }; + let tools_option = if tools.is_empty() { None } else { Some(tools) }; + let max_tokens = Some(args.max_tokens.unwrap_or(64000)); - serde_json::to_string(&request) - .map_err(|e| Error::internal_err(format!("Failed to serialize request: {}", e))) + // Build request based on platform + if self.is_vertex() { + // For Vertex AI: no model field, anthropic_version in body + let request = AnthropicVertexRequest { + anthropic_version: ANTHROPIC_VERSION_VERTEX, + system, + messages: anthropic_messages, + tools: tools_option, + tool_choice, + temperature: args.temperature, + max_tokens, + stream: true, + }; + serde_json::to_string(&request) + .map_err(|e| Error::internal_err(format!("Failed to serialize request: {}", e))) + } else { + // For standard API: model in body, anthropic_version in header + let request = AnthropicRequest { + model: args.model, + system, + messages: anthropic_messages, + tools: tools_option, + tool_choice, + temperature: args.temperature, + max_tokens, + stream: true, + }; + serde_json::to_string(&request) + .map_err(|e| Error::internal_err(format!("Failed to serialize request: {}", e))) + } } } @@ -436,8 +487,15 @@ impl QueryBuilder for AnthropicQueryBuilder { }) } - fn get_endpoint(&self, base_url: &str, _model: &str, _output_type: &OutputType) -> String { - format!("{}/messages", base_url) + fn get_endpoint(&self, base_url: &str, model: &str, _output_type: &OutputType) -> String { + if self.is_vertex() { + // For Vertex AI, the model is specified in the URL path + // Expected base_url format: https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/publishers/anthropic/models + // We append the model and :streamRawPredict + format!("{}/{}:streamRawPredict", base_url.trim_end_matches('/'), model) + } else { + format!("{}/messages", base_url) + } } fn get_auth_headers( @@ -446,9 +504,16 @@ impl QueryBuilder for AnthropicQueryBuilder { _base_url: &str, _output_type: &OutputType, ) -> Vec<(&'static str, String)> { - vec![ - ("x-api-key", api_key.to_string()), - ("anthropic-version", "2023-06-01".to_string()), - ] + if self.is_vertex() { + // For Vertex AI, use Bearer token authentication + // The api_key should be an OAuth2 access token + vec![("Authorization", format!("Bearer {}", api_key))] + } else { + // Standard Anthropic API uses x-api-key and anthropic-version header + vec![ + ("x-api-key", api_key.to_string()), + ("anthropic-version", ANTHROPIC_VERSION_STANDARD.to_string()), + ] + } } } diff --git a/backend/windmill-worker/src/ai/providers/google_ai.rs b/backend/windmill-worker/src/ai/providers/google_ai.rs index a904fe1073..b2c31f9e35 100644 --- a/backend/windmill-worker/src/ai/providers/google_ai.rs +++ b/backend/windmill-worker/src/ai/providers/google_ai.rs @@ -212,8 +212,6 @@ pub struct GeminiPredictCandidate { // Query Builder Implementation // ============================================================================ -const GEMINI_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta"; - pub struct GoogleAIQueryBuilder; impl GoogleAIQueryBuilder { @@ -637,7 +635,7 @@ impl QueryBuilder for GoogleAIQueryBuilder { fn get_endpoint( &self, - _base_url: &str, // Ignored - always use Google's URL + base_url: &str, model: &str, output_type: &OutputType, ) -> String { @@ -645,7 +643,7 @@ impl QueryBuilder for GoogleAIQueryBuilder { OutputType::Text => { format!( "{}/models/{}:streamGenerateContent?alt=sse", - GEMINI_BASE_URL, model + base_url, model ) } OutputType::Image => { @@ -654,7 +652,7 @@ impl QueryBuilder for GoogleAIQueryBuilder { } else { "generateContent" }; - format!("{}/models/{}:{}", GEMINI_BASE_URL, model, url_suffix) + format!("{}/models/{}:{}", base_url, model, url_suffix) } } } diff --git a/backend/windmill-worker/src/ai/query_builder.rs b/backend/windmill-worker/src/ai/query_builder.rs index 37f3d3c3ba..a33177f74c 100644 --- a/backend/windmill-worker/src/ai/query_builder.rs +++ b/backend/windmill-worker/src/ai/query_builder.rs @@ -95,8 +95,11 @@ pub fn create_query_builder(provider: &ProviderWithResource) -> Box Box::new(GoogleAIQueryBuilder::new()), // OpenAI use the Responses API AIProvider::OpenAI => Box::new(OpenAIQueryBuilder::new(provider.kind.clone())), - // Anthropic uses its own API format - AIProvider::Anthropic => Box::new(AnthropicQueryBuilder::new(provider.kind.clone())), + // Anthropic uses its own API format (with platform-specific handling for Vertex AI) + AIProvider::Anthropic => Box::new(AnthropicQueryBuilder::new( + provider.kind.clone(), + provider.get_platform().clone(), + )), AIProvider::OpenRouter => Box::new(OpenRouterQueryBuilder::new()), // All other providers use the completion endpoint _ => Box::new(OtherQueryBuilder::new(provider.kind.clone())), diff --git a/backend/windmill-worker/src/ai/types.rs b/backend/windmill-worker/src/ai/types.rs index 51f5833851..1a638a5e76 100644 --- a/backend/windmill-worker/src/ai/types.rs +++ b/backend/windmill-worker/src/ai/types.rs @@ -210,6 +210,14 @@ impl From for AIAgentArgs { } } +#[derive(Deserialize, Debug, Clone, Default, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum AnthropicPlatform { + #[default] + Standard, + GoogleVertexAi, +} + #[derive(Deserialize, Debug)] pub struct ProviderResource { #[serde(alias = "apiKey")] @@ -221,6 +229,9 @@ pub struct ProviderResource { pub aws_access_key_id: Option, #[serde(alias = "awsSecretAccessKey")] pub aws_secret_access_key: Option, + /// Platform for Anthropic API (standard or google_vertex_ai) + #[serde(default)] + pub platform: AnthropicPlatform, } #[derive(Deserialize, Debug)] @@ -260,6 +271,10 @@ impl ProviderWithResource { pub fn get_aws_secret_access_key(&self) -> Option<&str> { self.resource.aws_secret_access_key.as_deref() } + + pub fn get_platform(&self) -> &AnthropicPlatform { + &self.resource.platform + } } #[derive(Serialize)]