kalugin66/windmill
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ai): add google vertex ai platform support for google ai provider

Browse Source 
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
claude[bot]
2026-02-03 11:58:53 +00:00
parent 3a719cea6b
commit 2886f4a8df
4 changed files with 96 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
backend/windmill-api/src/ai.rs
View File

@@ -141,6 +141,15 @@ enum AnthropicPlatform {
GoogleVertexAi,
}
/// Platform for Google AI API
#[derive(Deserialize, Debug, Clone, Default, PartialEq)]
#[serde(rename_all = "snake_case")]
enum GoogleAIPlatform {
#[default]
Standard,
GoogleVertexAi,
}
#[derive(Deserialize, Debug)]
struct AIStandardResource {
#[serde(alias = "baseUrl", default, deserialize_with = "empty_string_as_none")]
@@ -158,6 +167,9 @@ struct AIStandardResource {
/// Platform for Anthropic API (standard or google_vertex_ai)
#[serde(default)]
platform: AnthropicPlatform,
/// Platform for Google AI API (standard or google_vertex_ai)
#[serde(default)]
google_platform: GoogleAIPlatform,
}
#[derive(Deserialize, Debug)]
@@ -186,6 +198,7 @@ struct AIRequestConfig {
#[allow(dead_code)]
pub aws_secret_access_key: Option<String>,
pub platform: AnthropicPlatform,
pub google_platform: GoogleAIPlatform,
}
impl AIRequestConfig {
@@ -205,10 +218,12 @@ impl AIRequestConfig {
aws_access_key_id,
aws_secret_access_key,
platform,
google_platform,
) = match resource {
AIResource::Standard(resource) => {
let region = resource.region.clone();
let platform = resource.platform.clone();
let google_platform = resource.google_platform.clone();
// Skip get_base_url for Bedrock - it uses SDK directly, not HTTP
let base_url = if matches!(provider, AIProvider::AWSBedrock) {
String::new()
@@ -249,6 +264,7 @@ impl AIRequestConfig {
aws_access_key_id,
aws_secret_access_key,
platform,
google_platform,
)
}
AIResource::OAuth(resource) => {
@@ -270,6 +286,7 @@ impl AIRequestConfig {
None,
None,
AnthropicPlatform::Standard,
GoogleAIPlatform::Standard,
)
}
};
@@ -284,6 +301,7 @@ impl AIRequestConfig {
aws_access_key_id,
aws_secret_access_key,
platform,
google_platform,
})
}
@@ -341,9 +359,11 @@ impl AIRequestConfig {
let is_anthropic_vertex = is_anthropic && self.platform == AnthropicPlatform::GoogleVertexAi;
let is_anthropic_sdk = headers.get("X-Anthropic-SDK").is_some();
let is_google_ai = matches!(provider, AIProvider::GoogleAI);
let is_google_ai_vertex = is_google_ai && self.google_platform == GoogleAIPlatform::GoogleVertexAi;
// GoogleAI uses OpenAI-compatible endpoint in the proxy (for the chat), but not for the ai agent
let base_url = if is_google_ai {
// For Vertex AI, the base_url is already properly configured, no need to add /openai
let base_url = if is_google_ai && !is_google_ai_vertex {
format!("{}/openai", base_url)
} else {
base_url.to_string()
@@ -388,7 +408,11 @@ impl AIRequestConfig {
if let Some(api_key) = self.api_key {
if is_azure {
request = request.header("api-key", api_key.clone())
} else if is_google_ai && !is_google_ai_vertex {
// Standard Google AI uses x-goog-api-key header
request = request.header("x-goog-api-key", api_key.clone())
} else {
// Vertex AI (both Google and Anthropic) and most other providers use Bearer token
request = request.header("authorization", format!("Bearer {}", api_key.clone()))
}
// For standard Anthropic API, also add X-API-Key header (but not for Vertex AI)

71
backend/windmill-worker/src/ai/providers/google_ai.rs
View File

@@ -6,7 +6,7 @@ use crate::ai::{
image_handler::download_and_encode_s3_image,
query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventProcessor},
sse::{GeminiSSEParser, SSEParser},
types::*,
types::{GoogleAIPlatform, *},
utils::parse_data_url,
};
@@ -214,11 +214,17 @@ pub struct GeminiPredictCandidate {
// Query Builder Implementation
// ============================================================================
pub struct GoogleAIQueryBuilder;
pub struct GoogleAIQueryBuilder {
platform: GoogleAIPlatform,
}
impl GoogleAIQueryBuilder {
pub fn new() -> Self {
Self
pub fn new(platform: GoogleAIPlatform) -> Self {
Self { platform }
}
fn is_vertex(&self) -> bool {
self.platform == GoogleAIPlatform::GoogleVertexAi
}
/// Build a text request using the native Gemini API format
@@ -661,20 +667,41 @@ impl QueryBuilder for GoogleAIQueryBuilder {
}
fn get_endpoint(&self, base_url: &str, model: &str, output_type: &OutputType) -> String {
match output_type {
OutputType::Text => {
format!(
"{}/models/{}:streamGenerateContent?alt=sse",
base_url, model
)
if self.is_vertex() {
// For Vertex AI, the base_url should be in format:
// https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/publishers/google/models
// We append the model and the appropriate action
let base_url = base_url.trim_end_matches('/');
match output_type {
OutputType::Text => {
format!("{}/{}:streamGenerateContent?alt=sse", base_url, model)
}
OutputType::Image => {
let url_suffix = if model.contains("imagen") {
"predict"
} else {
"generateContent"
};
format!("{}/{}:{}", base_url, model, url_suffix)
}
}
OutputType::Image => {
let url_suffix = if model.contains("imagen") {
"predict"
} else {
"generateContent"
};
format!("{}/models/{}:{}", base_url, model, url_suffix)
} else {
// Standard Google AI endpoint
match output_type {
OutputType::Text => {
format!(
"{}/models/{}:streamGenerateContent?alt=sse",
base_url, model
)
}
OutputType::Image => {
let url_suffix = if model.contains("imagen") {
"predict"
} else {
"generateContent"
};
format!("{}/models/{}:{}", base_url, model, url_suffix)
}
}
}
}
@@ -685,7 +712,13 @@ impl QueryBuilder for GoogleAIQueryBuilder {
_base_url: &str,
_output_type: &OutputType,
) -> Vec<(&'static str, String)> {
// Native Gemini API always uses x-goog-api-key
vec![("x-goog-api-key", api_key.to_string())]
if self.is_vertex() {
// For Vertex AI, use Bearer token authentication
// The api_key should be an OAuth2 access token (from gcloud auth print-access-token)
vec![("Authorization", format!("Bearer {}", api_key))]
} else {
// Native Gemini API uses x-goog-api-key
vec![("x-goog-api-key", api_key.to_string())]
}
}
}

6
backend/windmill-worker/src/ai/query_builder.rs
View File

@@ -113,8 +113,10 @@ pub fn create_query_builder(provider: &ProviderWithResource) -> Box<dyn QueryBui
use windmill_common::ai_providers::AIProvider;
match provider.kind {
// Google AI uses the Gemini API
AIProvider::GoogleAI => Box::new(GoogleAIQueryBuilder::new()),
// Google AI uses the Gemini API (with platform-specific handling for Vertex AI)
AIProvider::GoogleAI => Box::new(GoogleAIQueryBuilder::new(
provider.get_google_platform().clone(),
)),
// OpenAI use the Responses API
AIProvider::OpenAI => Box::new(OpenAIQueryBuilder::new(provider.kind.clone())),
// Anthropic uses its own API format (with platform-specific handling for Vertex AI)

15
backend/windmill-worker/src/ai/types.rs
View File

@@ -164,6 +164,14 @@ pub enum AnthropicPlatform {
GoogleVertexAi,
}
#[derive(Deserialize, Debug, Clone, Default, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum GoogleAIPlatform {
#[default]
Standard,
GoogleVertexAi,
}
#[derive(Deserialize, Debug)]
pub struct ProviderResource {
#[serde(alias = "apiKey", default, deserialize_with = "empty_string_as_none")]
@@ -182,6 +190,9 @@ pub struct ProviderResource {
/// Platform for Anthropic API (standard or google_vertex_ai)
#[serde(default)]
pub platform: AnthropicPlatform,
/// Platform for Google AI API (standard or google_vertex_ai)
#[serde(default)]
pub google_platform: GoogleAIPlatform,
}
#[derive(Deserialize, Debug)]
@@ -227,6 +238,10 @@ impl ProviderWithResource {
pub fn get_platform(&self) -> &AnthropicPlatform {
&self.resource.platform
}
pub fn get_google_platform(&self) -> &GoogleAIPlatform {
&self.resource.google_platform
}
}
/// Token usage information from the AI provider