Compare commits

...

2 Commits

Author SHA1 Message Date
claude[bot]
6518ed184a refactor(ai): unify platform enum for Anthropic and Google AI providers
- Merge AnthropicPlatform and GoogleAIPlatform into a single AIPlatform enum
- Remove duplicate google_platform field from AIStandardResource, AIRequestConfig, and ProviderResource
- Both Anthropic and Google AI providers now use the same platform field to indicate Vertex AI usage
- Simplify ProviderWithResource to have a single get_platform() method

This reduces code duplication and makes the resource configuration cleaner -
the same platform field works for both providers since both support the same
values (standard and google_vertex_ai).

Co-authored-by: centdix <centdix@users.noreply.github.com>
2026-02-03 13:02:47 +00:00
claude[bot]
2886f4a8df feat(ai): add google vertex ai platform support for google ai provider
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2026-02-03 11:58:53 +00:00
5 changed files with 84 additions and 39 deletions

View File

@@ -132,10 +132,11 @@ struct AIOAuthResource {
user: Option<String>,
}
/// Platform for Anthropic API
/// Platform for AI providers (Anthropic, Google AI)
/// Both Anthropic and Google AI can run on standard endpoints or Google Vertex AI
#[derive(Deserialize, Debug, Clone, Default, PartialEq)]
#[serde(rename_all = "snake_case")]
enum AnthropicPlatform {
enum AIPlatform {
#[default]
Standard,
GoogleVertexAi,
@@ -155,9 +156,10 @@ struct AIStandardResource {
aws_access_key_id: Option<String>,
#[serde(alias = "awsSecretAccessKey", default, deserialize_with = "empty_string_as_none")]
aws_secret_access_key: Option<String>,
/// Platform for Anthropic API (standard or google_vertex_ai)
/// Platform for AI providers (standard or google_vertex_ai)
/// Used by both Anthropic and Google AI providers to indicate Vertex AI usage
#[serde(default)]
platform: AnthropicPlatform,
platform: AIPlatform,
}
#[derive(Deserialize, Debug)]
@@ -185,7 +187,8 @@ struct AIRequestConfig {
pub aws_access_key_id: Option<String>,
#[allow(dead_code)]
pub aws_secret_access_key: Option<String>,
pub platform: AnthropicPlatform,
/// Platform for AI providers (standard or google_vertex_ai)
pub platform: AIPlatform,
}
impl AIRequestConfig {
@@ -213,9 +216,7 @@ impl AIRequestConfig {
let base_url = if matches!(provider, AIProvider::AWSBedrock) {
String::new()
} else {
provider
.get_base_url(resource.base_url, db)
.await?
provider.get_base_url(resource.base_url, db).await?
};
let api_key = if let Some(api_key) = resource.api_key {
Some(get_variable_or_self(api_key, db, w_id).await?)
@@ -269,7 +270,7 @@ impl AIRequestConfig {
None,
None,
None,
AnthropicPlatform::Standard,
AIPlatform::Standard,
)
}
};
@@ -338,12 +339,14 @@ 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_vertex = is_anthropic && self.platform == AIPlatform::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.platform == AIPlatform::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 +391,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)

View File

@@ -324,16 +324,16 @@ pub struct AnthropicResponse {
pub struct AnthropicQueryBuilder {
#[allow(dead_code)]
provider_kind: AIProvider,
platform: AnthropicPlatform,
platform: AIPlatform,
}
impl AnthropicQueryBuilder {
pub fn new(provider_kind: AIProvider, platform: AnthropicPlatform) -> Self {
pub fn new(provider_kind: AIProvider, platform: AIPlatform) -> Self {
Self { provider_kind, platform }
}
fn is_vertex(&self) -> bool {
self.platform == AnthropicPlatform::GoogleVertexAi
self.platform == AIPlatform::GoogleVertexAi
}
async fn build_text_request(

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::{AIPlatform, *},
utils::parse_data_url,
};
@@ -214,11 +214,17 @@ pub struct GeminiPredictCandidate {
// Query Builder Implementation
// ============================================================================
pub struct GoogleAIQueryBuilder;
pub struct GoogleAIQueryBuilder {
platform: AIPlatform,
}
impl GoogleAIQueryBuilder {
pub fn new() -> Self {
Self
pub fn new(platform: AIPlatform) -> Self {
Self { platform }
}
fn is_vertex(&self) -> bool {
self.platform == AIPlatform::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())]
}
}
}

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_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)

View File

@@ -156,9 +156,11 @@ impl From<AIAgentArgsRaw> for AIAgentArgs {
}
}
/// Platform for AI providers (Anthropic, Google AI)
/// Both Anthropic and Google AI can run on standard endpoints or Google Vertex AI
#[derive(Deserialize, Debug, Clone, Default, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum AnthropicPlatform {
pub enum AIPlatform {
#[default]
Standard,
GoogleVertexAi,
@@ -179,9 +181,10 @@ pub struct ProviderResource {
#[allow(dead_code)]
#[serde(alias = "awsSecretAccessKey", default, deserialize_with = "empty_string_as_none")]
pub aws_secret_access_key: Option<String>,
/// Platform for Anthropic API (standard or google_vertex_ai)
/// Platform for AI providers (standard or google_vertex_ai)
/// Used by both Anthropic and Google AI providers to indicate Vertex AI usage
#[serde(default)]
pub platform: AnthropicPlatform,
pub platform: AIPlatform,
}
#[derive(Deserialize, Debug)]
@@ -224,7 +227,7 @@ impl ProviderWithResource {
self.resource.aws_secret_access_key.as_deref()
}
pub fn get_platform(&self) -> &AnthropicPlatform {
pub fn get_platform(&self) -> &AIPlatform {
&self.resource.platform
}
}