From 289549a048e96019e16d84f40ebbb5f0f58ff5eb Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Wed, 20 May 2026 17:24:07 +0200 Subject: [PATCH] refactor: move google ai proxy handling to windmill-ai (#9260) * refactor: add ai proxy execution mode * refactor: move google ai proxy handling * refactor: share google ai request building --- backend/Cargo.lock | 1 + backend/windmill-ai/Cargo.toml | 1 + .../windmill-ai/src/providers/google_ai.rs | 540 +++++++++++++++++- backend/windmill-ai/src/proxy.rs | 66 ++- backend/windmill-api/src/ai.rs | 189 +++--- backend/windmill-api/src/google.rs | 354 ------------ backend/windmill-api/src/lib.rs | 1 - docs/windmill-ai-refactor-plan.md | 37 +- 8 files changed, 678 insertions(+), 511 deletions(-) delete mode 100644 backend/windmill-api/src/google.rs diff --git a/backend/Cargo.lock b/backend/Cargo.lock index b8b0ab8625..8e590776cd 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -13871,6 +13871,7 @@ dependencies = [ name = "windmill-ai" version = "1.704.1" dependencies = [ + "async-stream", "async-trait", "aws-config", "aws-credential-types", diff --git a/backend/windmill-ai/Cargo.toml b/backend/windmill-ai/Cargo.toml index b419101f25..8f2f679c7e 100644 --- a/backend/windmill-ai/Cargo.toml +++ b/backend/windmill-ai/Cargo.toml @@ -20,6 +20,7 @@ windmill-parser.workspace = true windmill-mcp = { workspace = true, optional = true } async-trait.workspace = true +async-stream.workspace = true base64.workspace = true bytes.workspace = true eventsource-stream.workspace = true diff --git a/backend/windmill-ai/src/providers/google_ai.rs b/backend/windmill-ai/src/providers/google_ai.rs index 81f34e7acb..b6a6295d88 100644 --- a/backend/windmill-ai/src/providers/google_ai.rs +++ b/backend/windmill-ai/src/providers/google_ai.rs @@ -1,15 +1,24 @@ use crate::{ ai_google::{ - openai_messages_to_gemini, openai_tools_to_gemini, GeminiGenerationConfig, + gemini_event_to_openai_sse_chunks, gemini_response_to_openai, openai_messages_to_gemini, + openai_tools_to_gemini, parse_gemini_response, parse_gemini_sse_event, + sanitize_schema_for_google, GeminiFunctionDeclaration, GeminiGenerationConfig, GeminiImageContent, GeminiImageRequest, GeminiImageResponse, GeminiInlineData, GeminiPart, GeminiPredictContent, GeminiTextRequest, GeminiTool, }, image_handler::{download_and_encode_s3_image, prepare_messages_for_api}, + proxy::{ProxyBuildArgs, ProxyRequest}, query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink}, sse::{GeminiSSEParser, SSEParser}, types::*, }; use async_trait::async_trait; +use bytes::Bytes; +use eventsource_stream::Eventsource; +use futures::{stream::BoxStream, StreamExt}; +use http::{header, HeaderMap, HeaderValue, Method, StatusCode}; +use serde::Deserialize; +use serde_json::json; use windmill_common::{client::AuthedClient, error::Error}; // ============================================================================ @@ -37,22 +46,11 @@ impl GoogleAIQueryBuilder { ) -> Result { let prepared_messages = prepare_messages_for_api(args.messages, client, workspace_id).await?; - let (contents, system_instruction) = openai_messages_to_gemini(&prepared_messages); - - let tools = self.convert_tools_to_gemini(args.tools, args.has_websearch); - - let generation_config = self.build_generation_config(args); - - let request = GeminiTextRequest { - contents, - tools, - tool_config: None, - system_instruction, - generation_config, - }; - - serde_json::to_string(&request) - .map_err(|e| Error::internal_err(format!("Failed to serialize request: {}", e))) + build_gemini_text_request_body( + &prepared_messages, + self.convert_tools_to_gemini(args.tools, args.has_websearch), + self.build_generation_config(args), + ) } async fn build_image_request( @@ -155,19 +153,378 @@ impl GoogleAIQueryBuilder { (None, None) }; - if args.temperature.is_some() || args.max_tokens.is_some() || response_mime_type.is_some() { - Some(GeminiGenerationConfig { - temperature: args.temperature, - max_output_tokens: args.max_tokens, - response_mime_type, - response_schema, - }) - } else { - None - } + build_gemini_generation_config( + args.temperature, + args.max_tokens, + response_mime_type, + response_schema, + ) } } +fn build_gemini_text_request_body( + messages: &[OpenAIMessage], + tools: Option>, + generation_config: Option, +) -> Result { + let request = build_gemini_text_request(messages, tools, generation_config); + serde_json::to_string(&request) + .map_err(|e| Error::internal_err(format!("Failed to serialize Gemini request: {}", e))) +} + +fn build_gemini_text_request( + messages: &[OpenAIMessage], + tools: Option>, + generation_config: Option, +) -> GeminiTextRequest { + let (contents, system_instruction) = openai_messages_to_gemini(messages); + + GeminiTextRequest { contents, tools, tool_config: None, system_instruction, generation_config } +} + +fn build_gemini_generation_config( + temperature: Option, + max_tokens: Option, + response_mime_type: Option, + response_schema: Option, +) -> Option { + if temperature.is_some() + || max_tokens.is_some() + || response_mime_type.is_some() + || response_schema.is_some() + { + Some(GeminiGenerationConfig { + temperature, + max_output_tokens: max_tokens, + response_mime_type, + response_schema, + }) + } else { + None + } +} + +#[derive(Deserialize, Debug)] +struct GoogleAIProxyChatRequest { + model: String, + messages: Vec, + #[serde(default)] + stream: bool, + #[serde(default)] + temperature: Option, + #[serde(default)] + max_tokens: Option, + #[serde(default)] + tools: Option>, +} + +#[derive(Deserialize, Debug)] +struct GoogleAIProxyChatTool { + function: GoogleAIProxyChatToolFunction, +} + +#[derive(Deserialize, Debug)] +struct GoogleAIProxyChatToolFunction { + name: String, + #[serde(default)] + description: Option, + #[serde(default)] + parameters: Option, +} + +#[derive(Deserialize)] +struct GeminiModel { + name: String, + #[serde(rename = "displayName", default)] + display_name: String, +} + +#[derive(Deserialize)] +struct GeminiModelsResponse { + #[serde(default)] + models: Vec, +} + +struct GoogleAIProxyRequest { + request: ProxyRequest, + model: String, + stream: bool, +} + +pub enum GoogleAIProxyResponseBody { + Fixed(Bytes), + Stream(BoxStream<'static, std::result::Result>), +} + +pub struct GoogleAIProxyResponse { + pub status_code: StatusCode, + pub headers: HeaderMap, + pub body: GoogleAIProxyResponseBody, +} + +/// Handle a workspace Google AI chat proxy request. +/// +/// The API still owns credential resolution, auditing, and keepalive injection. +/// Callers must verify the user can use the supplied credentials before calling. +/// This helper owns the provider-specific OpenAI <-> Gemini transformations. +pub async fn handle_google_ai_chat_proxy( + client: &reqwest::Client, + args: &ProxyBuildArgs<'_>, +) -> Result { + let GoogleAIProxyRequest { request, model, stream } = build_google_ai_chat_proxy_request(args)?; + + let response = + send_google_ai_proxy_request(client, request, "Failed to send request to Gemini API") + .await?; + + if stream { + Ok(convert_streaming_response(response, &model)) + } else { + convert_non_streaming_response(response, &model).await + } +} + +/// Handle a workspace Google AI model-list proxy request. +/// +/// The API still owns credential resolution and auditing. Callers must verify +/// the user can use the supplied credentials before calling. +pub async fn handle_google_ai_models_proxy( + client: &reqwest::Client, + args: &ProxyBuildArgs<'_>, +) -> Result { + let request = build_google_ai_models_proxy_request(args); + let response = + send_google_ai_proxy_request(client, request, "Failed to fetch Gemini models").await?; + + let gemini_resp: GeminiModelsResponse = response.json().await.map_err(|e| { + Error::internal_err(format!("Failed to parse Gemini models response: {}", e)) + })?; + + let data: Vec = gemini_resp + .models + .into_iter() + .map(|m| { + json!({ + "id": m.name, + "object": "model", + "display_name": m.display_name, + }) + }) + .collect(); + + let body = serde_json::to_vec(&json!({ "data": data })) + .map_err(|e| Error::internal_err(format!("Failed to serialize models: {}", e)))?; + + Ok(GoogleAIProxyResponse { + status_code: StatusCode::OK, + headers: json_response_headers(), + body: GoogleAIProxyResponseBody::Fixed(Bytes::from(body)), + }) +} + +fn build_google_ai_chat_proxy_request( + args: &ProxyBuildArgs<'_>, +) -> Result { + let request: GoogleAIProxyChatRequest = serde_json::from_slice(args.body) + .map_err(|e| Error::BadRequest(format!("Failed to parse request body: {}", e)))?; + + let gemini_tools = request.tools.as_ref().map(|tools| { + let declarations: Vec = tools + .iter() + .map(|t| { + let mut params = t.function.parameters.clone().unwrap_or(json!({})); + sanitize_schema_for_google(&mut params); + GeminiFunctionDeclaration { + name: t.function.name.clone(), + description: t.function.description.clone(), + parameters: params, + } + }) + .collect(); + vec![GeminiTool { function_declarations: Some(declarations), google_search: None }] + }); + + let body = build_gemini_text_request_body( + &request.messages, + gemini_tools, + build_gemini_generation_config(request.temperature, request.max_tokens, None, None), + )? + .into_bytes(); + + let credentials = args.credentials; + let base_url = credentials.base_url.trim_end_matches('/'); + let is_vertex = credentials.platform == AIPlatform::GoogleVertexAi; + let endpoint = if request.stream { + format!( + "{}?alt=sse", + build_google_ai_model_endpoint( + base_url, + &request.model, + "streamGenerateContent", + is_vertex, + ) + ) + } else { + build_google_ai_model_endpoint(base_url, &request.model, "generateContent", is_vertex) + }; + + let mut headers = vec![("content-type".to_string(), "application/json".to_string())]; + add_google_ai_auth_header( + &mut headers, + credentials.api_key.as_deref().unwrap_or(""), + is_vertex, + ); + + Ok(GoogleAIProxyRequest { + request: ProxyRequest { method: Method::POST, url: endpoint, headers, body }, + model: request.model, + stream: request.stream, + }) +} + +fn build_google_ai_models_proxy_request(args: &ProxyBuildArgs<'_>) -> ProxyRequest { + let credentials = args.credentials; + let base_url = credentials.base_url.trim_end_matches('/'); + let is_vertex = credentials.platform == AIPlatform::GoogleVertexAi; + let url = if is_vertex { + base_url.to_string() + } else { + format!("{}/models", base_url) + }; + + let mut headers = Vec::new(); + add_google_ai_auth_header( + &mut headers, + credentials.api_key.as_deref().unwrap_or(""), + is_vertex, + ); + + ProxyRequest { method: Method::GET, url, headers, body: Vec::new() } +} + +fn build_google_ai_model_endpoint( + base_url: &str, + model: &str, + action: &str, + is_vertex: bool, +) -> String { + if is_vertex { + format!("{}/{}:{}", base_url, model, action) + } else { + format!("{}/models/{}:{}", base_url, model, action) + } +} + +fn add_google_ai_auth_header(headers: &mut Vec<(String, String)>, api_key: &str, is_vertex: bool) { + if is_vertex { + headers.push(("Authorization".to_string(), format!("Bearer {}", api_key))); + } else { + headers.push(("x-goog-api-key".to_string(), api_key.to_string())); + } +} + +async fn send_google_ai_proxy_request( + client: &reqwest::Client, + proxy_request: ProxyRequest, + send_error_message: &str, +) -> Result { + let mut request = client.request(proxy_request.method.clone(), &proxy_request.url); + for (header_name, header_value) in &proxy_request.headers { + request = request.header(header_name.as_str(), header_value.as_str()); + } + + let response = request + .body(proxy_request.body) + .send() + .await + .map_err(|e| Error::internal_err(format!("{}: {}", send_error_message, e)))?; + + if let Err(e) = response.error_for_status_ref() { + let status = e.status().map(|s| s.to_string()).unwrap_or_default(); + let body = response.text().await.unwrap_or_default(); + return Err(Error::AIError(format!("{}: {}", status, body))); + } + + Ok(response) +} + +fn convert_streaming_response(response: reqwest::Response, model: &str) -> GoogleAIProxyResponse { + let id = format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()); + let model = model.to_string(); + + let gemini_sse_stream = response.bytes_stream().eventsource(); + let openai_sse_stream = async_stream::stream! { + tokio::pin!(gemini_sse_stream); + let mut tool_call_index: usize = 0; + while let Some(event) = gemini_sse_stream.next().await { + match event { + Ok(event) => match parse_gemini_sse_event(&event.data) { + Ok(Some(parsed)) => { + for chunk in gemini_event_to_openai_sse_chunks( + &parsed, &id, &model, &mut tool_call_index, + ) { + yield Ok::(Bytes::from(chunk)); + } + } + Ok(None) => {} + Err(e) => tracing::error!("Error parsing Gemini SSE event: {}", e), + }, + Err(e) => tracing::error!("Error reading Gemini SSE stream: {}", e), + } + } + yield Ok::(Bytes::from("data: [DONE]\n\n")); + } + .boxed(); + + GoogleAIProxyResponse { + status_code: StatusCode::OK, + headers: event_stream_response_headers(), + body: GoogleAIProxyResponseBody::Stream(openai_sse_stream), + } +} + +async fn convert_non_streaming_response( + response: reqwest::Response, + model: &str, +) -> Result { + let body = response + .bytes() + .await + .map_err(|e| Error::internal_err(format!("Failed to read Gemini response body: {}", e)))?; + + let parsed = parse_gemini_response(&body)?; + let openai_response = gemini_response_to_openai(&parsed, model); + + let body = serde_json::to_vec(&openai_response) + .map_err(|e| Error::internal_err(format!("Failed to serialize response: {}", e)))?; + + Ok(GoogleAIProxyResponse { + status_code: StatusCode::OK, + headers: json_response_headers(), + body: GoogleAIProxyResponseBody::Fixed(Bytes::from(body)), + }) +} + +fn json_response_headers() -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert( + header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ); + headers +} + +fn event_stream_response_headers() -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert( + header::CONTENT_TYPE, + HeaderValue::from_static("text/event-stream"), + ); + headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-cache")); + headers.insert(header::CONNECTION, HeaderValue::from_static("keep-alive")); + headers +} + #[async_trait] impl QueryBuilder for GoogleAIQueryBuilder { fn supports_tools_with_output_type(&self, output_type: &OutputType) -> bool { @@ -324,3 +681,132 @@ impl QueryBuilder for GoogleAIQueryBuilder { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ai_providers::AIProvider, proxy::ProviderCredentials}; + use std::collections::HashMap; + + fn credentials(base_url: &str, platform: AIPlatform) -> ProviderCredentials { + ProviderCredentials { + provider: AIProvider::GoogleAI, + base_url: base_url.to_string(), + api_key: Some("api-key".to_string()), + access_token: None, + organization_id: None, + user: None, + region: None, + aws_access_key_id: None, + aws_secret_access_key: None, + aws_session_token: None, + platform, + enable_1m_context: false, + custom_headers: HashMap::new(), + } + } + + #[test] + fn builds_standard_google_ai_chat_proxy_request() { + let credentials = credentials( + "https://generativelanguage.googleapis.com/v1beta/", + AIPlatform::Standard, + ); + let method = Method::POST; + let headers = HeaderMap::new(); + let body = br#"{ + "model": "gemini-2.0-flash", + "messages": [{"role": "user", "content": "hello"}], + "temperature": 0.2, + "max_tokens": 123, + "stream": false + }"#; + + let request = build_google_ai_chat_proxy_request(&ProxyBuildArgs { + method: &method, + path: "chat/completions", + headers: &headers, + body, + credentials: &credentials, + }) + .unwrap(); + + assert_eq!(request.request.method, Method::POST); + assert_eq!( + request.request.url, + "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent" + ); + assert!(!request.stream); + assert_eq!(request.model, "gemini-2.0-flash"); + assert!(request + .request + .headers + .contains(&("x-goog-api-key".to_string(), "api-key".to_string()))); + + let body: serde_json::Value = serde_json::from_slice(&request.request.body).unwrap(); + assert_eq!(body["generationConfig"]["maxOutputTokens"], 123); + assert_eq!(body["generationConfig"]["temperature"], 0.2); + assert!(body["contents"].is_array()); + } + + #[test] + fn builds_vertex_google_ai_streaming_proxy_request() { + let credentials = credentials( + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/publishers/google/models/", + AIPlatform::GoogleVertexAi, + ); + let method = Method::POST; + let headers = HeaderMap::new(); + let body = br#"{ + "model": "gemini-2.0-flash", + "messages": [{"role": "user", "content": "hello"}], + "stream": true + }"#; + + let request = build_google_ai_chat_proxy_request(&ProxyBuildArgs { + method: &method, + path: "chat/completions", + headers: &headers, + body, + credentials: &credentials, + }) + .unwrap(); + + assert_eq!( + request.request.url, + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:streamGenerateContent?alt=sse" + ); + assert!(request.stream); + assert!(request + .request + .headers + .contains(&("Authorization".to_string(), "Bearer api-key".to_string()))); + } + + #[test] + fn builds_google_ai_models_proxy_request() { + let credentials = credentials( + "https://generativelanguage.googleapis.com/v1beta/", + AIPlatform::Standard, + ); + let method = Method::GET; + let headers = HeaderMap::new(); + + let request = build_google_ai_models_proxy_request(&ProxyBuildArgs { + method: &method, + path: "models", + headers: &headers, + body: &[], + credentials: &credentials, + }); + + assert_eq!(request.method, Method::GET); + assert_eq!( + request.url, + "https://generativelanguage.googleapis.com/v1beta/models" + ); + assert!(request + .headers + .contains(&("x-goog-api-key".to_string(), "api-key".to_string()))); + } +} diff --git a/backend/windmill-ai/src/proxy.rs b/backend/windmill-ai/src/proxy.rs index 21b92705db..bbc570a18a 100644 --- a/backend/windmill-ai/src/proxy.rs +++ b/backend/windmill-ai/src/proxy.rs @@ -46,6 +46,24 @@ pub struct ProxyRequest { pub body: Vec, } +/// How the API proxy should execute a request for a provider. +/// +/// Most providers can be represented as a transformed HTTP request. Google AI +/// and Bedrock need native execution because their proxy paths also transform +/// responses or call an SDK rather than forwarding an HTTP request directly. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ProxyExecutionMode { + HttpForward, + NativeGoogleAi, + NativeAwsBedrock, +} + +impl ProxyExecutionMode { + pub fn uses_query_builder_proxy(self) -> bool { + matches!(self, Self::HttpForward) + } +} + pub fn supports_openai_compatible_proxy(provider: &AIProvider) -> bool { matches!( provider, @@ -60,8 +78,24 @@ pub fn supports_openai_compatible_proxy(provider: &AIProvider) -> bool { ) } +pub fn proxy_execution_mode(provider: &AIProvider) -> ProxyExecutionMode { + match provider { + AIProvider::OpenAI + | AIProvider::AzureOpenAI + | AIProvider::Anthropic + | AIProvider::Mistral + | AIProvider::DeepSeek + | AIProvider::Groq + | AIProvider::OpenRouter + | AIProvider::TogetherAI + | AIProvider::CustomAI => ProxyExecutionMode::HttpForward, + AIProvider::GoogleAI => ProxyExecutionMode::NativeGoogleAi, + AIProvider::AWSBedrock => ProxyExecutionMode::NativeAwsBedrock, + } +} + pub fn supports_query_builder_proxy(provider: &AIProvider) -> bool { - supports_openai_compatible_proxy(provider) || matches!(provider, AIProvider::Anthropic) + proxy_execution_mode(provider).uses_query_builder_proxy() } pub fn build_openai_compatible_proxy_request(args: &ProxyBuildArgs<'_>) -> Result { @@ -180,10 +214,32 @@ mod tests { #[test] fn query_builder_proxy_support_includes_anthropic() { - assert!(supports_query_builder_proxy(&AIProvider::OpenAI)); - assert!(supports_query_builder_proxy(&AIProvider::Anthropic)); - assert!(!supports_query_builder_proxy(&AIProvider::GoogleAI)); - assert!(!supports_query_builder_proxy(&AIProvider::AWSBedrock)); + let cases = [ + (AIProvider::OpenAI, ProxyExecutionMode::HttpForward), + (AIProvider::AzureOpenAI, ProxyExecutionMode::HttpForward), + (AIProvider::Anthropic, ProxyExecutionMode::HttpForward), + (AIProvider::Mistral, ProxyExecutionMode::HttpForward), + (AIProvider::DeepSeek, ProxyExecutionMode::HttpForward), + (AIProvider::Groq, ProxyExecutionMode::HttpForward), + (AIProvider::OpenRouter, ProxyExecutionMode::HttpForward), + (AIProvider::TogetherAI, ProxyExecutionMode::HttpForward), + (AIProvider::CustomAI, ProxyExecutionMode::HttpForward), + (AIProvider::GoogleAI, ProxyExecutionMode::NativeGoogleAi), + (AIProvider::AWSBedrock, ProxyExecutionMode::NativeAwsBedrock), + ]; + + for (provider, expected_mode) in cases { + let mode = proxy_execution_mode(&provider); + assert_eq!( + mode, expected_mode, + "unexpected proxy mode for {provider:?}" + ); + assert_eq!( + supports_query_builder_proxy(&provider), + mode.uses_query_builder_proxy(), + "query-builder support drifted for {provider:?}" + ); + } } #[test] diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index eb337e8125..05bf995700 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -19,9 +19,16 @@ use windmill_ai::ai_cache::current_instance_ai_config_revision; use windmill_ai::ai_providers::{ empty_string_as_none, AIPlatform, AIProvider, ProviderConfig, ProviderModel, }; -use windmill_ai::providers::create_proxy_query_builder; +use windmill_ai::providers::{ + create_proxy_query_builder, + google_ai::{ + handle_google_ai_chat_proxy, handle_google_ai_models_proxy, GoogleAIProxyResponse, + GoogleAIProxyResponseBody, + }, +}; use windmill_ai::proxy::{ - supports_query_builder_proxy, ProviderCredentials, ProxyBuildArgs, ProxyRequest, + proxy_execution_mode, supports_query_builder_proxy, ProviderCredentials, ProxyBuildArgs, + ProxyExecutionMode, ProxyRequest, }; use windmill_ai::utils::AI_HTTP_HEADERS; use windmill_audit::{audit_oss::audit_log, ActionKind}; @@ -368,82 +375,6 @@ impl AIRequestConfig { Ok(response.access_token) } - pub fn prepare_request( - self, - provider: &AIProvider, - path: &str, - method: Method, - _headers: HeaderMap, - body: Bytes, - ) -> Result { - let credentials = self.into_provider_credentials(provider.clone()); - - let body = if let Some(user) = credentials.user.as_ref() { - Self::add_user_to_body(body, user.clone())? - } else { - body - }; - - let base_url = credentials.base_url.trim_end_matches('/'); - - let is_azure = credentials.provider.is_azure_openai(base_url); - let is_google_ai = credentials.provider == AIProvider::GoogleAI; - - let base_url = base_url.to_string(); - let base_url = base_url.as_str(); - - // Build URL based on provider - let url = if is_azure { - let azure_url = AIProvider::build_azure_openai_url(base_url, path); - azure_url - } else { - let default_url = format!("{}/{}", base_url, path); - default_url - }; - - tracing::debug!("AI request URL: {}", url); - - let mut request = HTTP_CLIENT - .request(method.clone(), &url) - .header("content-type", "application/json"); - - // Add authentication headers - if let Some(api_key) = credentials.api_key { - if is_azure { - request = request.header("api-key", api_key.clone()) - } else if is_google_ai { - // Note: GoogleAI requests are intercepted earlier (see the GoogleAI - // handler block above) and never reach this code path. This branch - // is kept as a safety net for the standard Gemini API auth format. - request = request.header("x-goog-api-key", api_key.clone()) - } else { - request = request.header("authorization", format!("Bearer {}", api_key.clone())) - } - } - - if let Some(access_token) = credentials.access_token { - request = request.header("authorization", format!("Bearer {}", access_token)) - } - - request = request.body(body); - - if let Some(org_id) = credentials.organization_id { - request = request.header("OpenAI-Organization", org_id); - } - - // Apply custom headers from AI_HTTP_HEADERS environment variable - for (header_name, header_value) in AI_HTTP_HEADERS.iter() { - request = request.header(header_name.as_str(), header_value.as_str()); - } - - // Apply custom headers from the resource - for (header_name, header_value) in &credentials.custom_headers { - request = request.header(header_name.as_str(), header_value.as_str()); - } - - Ok(request) - } - fn into_provider_credentials(self, provider: AIProvider) -> ProviderCredentials { ProviderCredentials { provider, @@ -461,24 +392,6 @@ impl AIRequestConfig { custom_headers: self.custom_headers, } } - - fn add_user_to_body(body: Bytes, user: String) -> Result { - tracing::debug!("Adding user to request body"); - let mut json_body: HashMap> = serde_json::from_slice(&body) - .map_err(|e| Error::internal_err(format!("Failed to parse request body: {}", e)))?; - - let user_json_string = serde_json::Value::String(user).to_string(); // makes sure to escape characters - - json_body.insert( - "user".to_string(), - RawValue::from_string(user_json_string) - .map_err(|e| Error::internal_err(format!("Failed to parse user: {}", e)))?, - ); - - Ok(serde_json::to_vec(&json_body) - .map_err(|e| Error::internal_err(format!("Failed to reserialize request body: {}", e)))? - .into()) - } } #[derive(Clone, Debug)] @@ -613,6 +526,19 @@ fn proxy_request_to_request_builder(proxy_request: ProxyRequest) -> RequestBuild request.body(proxy_request.body) } +fn google_ai_proxy_response_to_body( + response: GoogleAIProxyResponse, +) -> (http::StatusCode, HeaderMap, axum::body::Body) { + let body = match response.body { + GoogleAIProxyResponseBody::Fixed(body) => axum::body::Body::from(body), + GoogleAIProxyResponseBody::Stream(stream) => axum::body::Body::from_stream( + inject_keepalives(stream, Duration::from_secs(KEEPALIVE_INTERVAL_SECS)), + ), + }; + + (response.status_code, response.headers, body) +} + pub(crate) fn inject_keepalives( upstream: S, interval: Duration, @@ -888,12 +814,10 @@ async fn proxy( ai_path = chat_path; } - // Handle GoogleAI (Gemini) using the native Gemini API - if matches!(provider, AIProvider::GoogleAI) { - let api_key = request_config.api_key.as_deref().unwrap_or(""); - let base_url = request_config.base_url.trim_end_matches('/'); - let is_vertex = request_config.platform == AIPlatform::GoogleVertexAi; + let proxy_mode = proxy_execution_mode(&provider); + // Handle GoogleAI (Gemini) using the native Gemini API + if matches!(proxy_mode, ProxyExecutionMode::NativeGoogleAi) { let mut tx = db.begin().await?; audit_log( &mut *tx, @@ -907,23 +831,32 @@ async fn proxy( .await?; tx.commit().await?; - return match ai_path.as_str() { - "chat/completions" => { - crate::google::handle_google_ai_chat(&body, api_key, base_url, is_vertex).await - } - "models" => crate::google::handle_google_ai_models(api_key, base_url, is_vertex).await, + let credentials = request_config.into_provider_credentials(provider.clone()); + let proxy_args = ProxyBuildArgs { + method: &method, + path: &ai_path, + headers: &headers, + body: &body, + credentials: &credentials, + }; + + let response = match ai_path.as_str() { + "chat/completions" => handle_google_ai_chat_proxy(&HTTP_CLIENT, &proxy_args).await, + "models" => handle_google_ai_models_proxy(&HTTP_CLIENT, &proxy_args).await, _ => Err(Error::BadRequest(format!( "Unsupported Google AI path: {}", ai_path ))), - }; + }?; + + return Ok(google_ai_proxy_response_to_body(response)); } // Handle Bedrock-specific logic when the feature is enabled #[cfg(feature = "bedrock")] { // Extract model and streaming flag for Bedrock transformation (only for POST requests) - let (model, is_streaming) = if matches!(provider, AIProvider::AWSBedrock) + let (model, is_streaming) = if matches!(proxy_mode, ProxyExecutionMode::NativeAwsBedrock) && method == Method::POST { #[derive(Deserialize, Debug)] @@ -940,7 +873,7 @@ async fn proxy( }; // For Bedrock requests, use the SDK-based approach - if matches!(provider, AIProvider::AWSBedrock) { + if matches!(proxy_mode, ProxyExecutionMode::NativeAwsBedrock) { let region = request_config .region .as_deref() @@ -1014,25 +947,35 @@ async fn proxy( // When bedrock feature is disabled, return error for Bedrock provider #[cfg(not(feature = "bedrock"))] - if matches!(provider, AIProvider::AWSBedrock) { + if matches!(proxy_mode, ProxyExecutionMode::NativeAwsBedrock) { return Err(Error::BadRequest( "AWS Bedrock support is not enabled. Build with 'bedrock' feature.".to_string(), )); } - let request = if supports_query_builder_proxy(&provider) { - let credentials = request_config.into_provider_credentials(provider.clone()); - let query_builder = create_proxy_query_builder(&credentials); - let proxy_request = query_builder.build_proxy_request(&ProxyBuildArgs { - method: &method, - path: &ai_path, - headers: &headers, - body: &body, - credentials: &credentials, - })?; - proxy_request_to_request_builder(proxy_request) - } else { - request_config.prepare_request(&provider, &ai_path, method, headers, body)? + let request = match proxy_mode { + ProxyExecutionMode::HttpForward => { + let credentials = request_config.into_provider_credentials(provider.clone()); + let query_builder = create_proxy_query_builder(&credentials); + let proxy_request = query_builder.build_proxy_request(&ProxyBuildArgs { + method: &method, + path: &ai_path, + headers: &headers, + body: &body, + credentials: &credentials, + })?; + proxy_request_to_request_builder(proxy_request) + } + ProxyExecutionMode::NativeGoogleAi => { + return Err(Error::internal_err( + "Google AI proxy route was not handled".to_string(), + )) + } + ProxyExecutionMode::NativeAwsBedrock => { + return Err(Error::BadRequest( + "Unsupported AWS Bedrock proxy request".to_string(), + )) + } }; let response = request.send().await.map_err(to_anyhow)?; diff --git a/backend/windmill-api/src/google.rs b/backend/windmill-api/src/google.rs deleted file mode 100644 index 95987ab789..0000000000 --- a/backend/windmill-api/src/google.rs +++ /dev/null @@ -1,354 +0,0 @@ -//! Google AI (Gemini API) handler for the AI chat proxy. -//! -//! Handles POST `chat/completions` requests using the native Gemini API, -//! converting from/to OpenAI format so the existing frontend parsers continue to work. -//! -//! Supports both standard Google AI (generativelanguage.googleapis.com) and -//! Google Vertex AI ({region}-aiplatform.googleapis.com) endpoints. -//! -//! Used by `windmill-api/src/ai.rs` when the provider is `GoogleAI`. -//! Shared conversion logic lives in `windmill_common::ai_google`. - -use axum::body::Body; -use bytes::Bytes; -use eventsource_stream::Eventsource; -use futures::StreamExt; -use serde::Deserialize; -use serde_json::json; -use windmill_ai::{ - ai_google::{ - gemini_event_to_openai_sse_chunks, gemini_response_to_openai, openai_messages_to_gemini, - parse_gemini_response, parse_gemini_sse_event, sanitize_schema_for_google, - GeminiFunctionDeclaration, GeminiGenerationConfig, GeminiTextRequest, GeminiTool, - }, - ai_types::OpenAIMessage, -}; -use windmill_common::error::{Error, Result}; - -use crate::ai::{inject_keepalives, HTTP_CLIENT, KEEPALIVE_INTERVAL_SECS}; - -// ============================================================================ -// Request type (OpenAI format received from the frontend) -// ============================================================================ - -#[derive(Deserialize, Debug)] -struct ChatRequest { - model: String, - messages: Vec, - #[serde(default)] - stream: bool, - #[serde(default)] - temperature: Option, - #[serde(default)] - max_tokens: Option, - #[serde(default)] - tools: Option>, -} - -#[derive(Deserialize, Debug)] -struct ChatRequestTool { - function: ChatRequestToolFunction, -} - -#[derive(Deserialize, Debug)] -struct ChatRequestToolFunction { - name: String, - #[serde(default)] - description: Option, - #[serde(default)] - parameters: Option, -} - -// ============================================================================ -// Helpers for Vertex AI vs standard Google AI URL/auth -// ============================================================================ - -/// Build the endpoint URL for a model action (streamGenerateContent, generateContent, predict). -/// -/// - Standard: `{base_url}/models/{model}:{action}` -/// - Vertex AI: `{base_url}/{model}:{action}` (base_url already contains .../publishers/google/models) -fn build_model_endpoint(base_url: &str, model: &str, action: &str, is_vertex: bool) -> String { - if is_vertex { - format!("{}/{}:{}", base_url, model, action) - } else { - format!("{}/models/{}:{}", base_url, model, action) - } -} - -/// Set the appropriate auth header on a request builder. -/// -/// - Standard: `x-goog-api-key` header -/// - Vertex AI: `Authorization: Bearer` header -fn set_auth( - request: reqwest::RequestBuilder, - api_key: &str, - is_vertex: bool, -) -> reqwest::RequestBuilder { - if is_vertex { - request.header("Authorization", format!("Bearer {}", api_key)) - } else { - request.header("x-goog-api-key", api_key) - } -} - -// ============================================================================ -// Public handler -// ============================================================================ - -/// Handle a `chat/completions` POST request using the native Gemini API. -/// -/// Converts the incoming OpenAI-format body to a `GeminiTextRequest`, sends it -/// to the appropriate Gemini endpoint, and converts the response back to the -/// OpenAI SSE or JSON format that the frontend expects. -pub async fn handle_google_ai_chat( - body: &Bytes, - api_key: &str, - base_url: &str, - is_vertex: bool, -) -> Result<(http::StatusCode, http::HeaderMap, Body)> { - let request: ChatRequest = serde_json::from_slice(body) - .map_err(|e| Error::BadRequest(format!("Failed to parse request body: {}", e)))?; - - let (contents, system_instruction) = openai_messages_to_gemini(&request.messages); - - let generation_config = if request.temperature.is_some() || request.max_tokens.is_some() { - Some(GeminiGenerationConfig { - temperature: request.temperature, - max_output_tokens: request.max_tokens, - response_mime_type: None, - response_schema: None, - }) - } else { - None - }; - - let gemini_tools = request.tools.as_ref().map(|tools| { - let declarations: Vec = tools - .iter() - .map(|t| { - let mut params = t.function.parameters.clone().unwrap_or(json!({})); - sanitize_schema_for_google(&mut params); - GeminiFunctionDeclaration { - name: t.function.name.clone(), - description: t.function.description.clone(), - parameters: params, - } - }) - .collect(); - vec![GeminiTool { function_declarations: Some(declarations), google_search: None }] - }); - - let gemini_request = GeminiTextRequest { - contents, - tools: gemini_tools, - tool_config: None, - system_instruction, - generation_config, - }; - - let request_body = serde_json::to_string(&gemini_request) - .map_err(|e| Error::internal_err(format!("Failed to serialize Gemini request: {}", e)))?; - - let base_url = base_url.trim_end_matches('/'); - - if request.stream { - handle_streaming(&request.model, request_body, api_key, base_url, is_vertex).await - } else { - handle_non_streaming(&request.model, request_body, api_key, base_url, is_vertex).await - } -} - -// ============================================================================ -// Streaming path -// ============================================================================ - -async fn handle_streaming( - model: &str, - request_body: String, - api_key: &str, - base_url: &str, - is_vertex: bool, -) -> Result<(http::StatusCode, http::HeaderMap, Body)> { - let endpoint = format!( - "{}?alt=sse", - build_model_endpoint(base_url, model, "streamGenerateContent", is_vertex) - ); - - let request = HTTP_CLIENT - .post(&endpoint) - .header("content-type", "application/json") - .body(request_body); - let request = set_auth(request, api_key, is_vertex); - - let response = request - .send() - .await - .map_err(|e| Error::internal_err(format!("Failed to send request to Gemini API: {}", e)))?; - - if let Err(e) = response.error_for_status_ref() { - let status = e.status().map(|s| s.to_string()).unwrap_or_default(); - let body = response.text().await.unwrap_or_default(); - return Err(Error::AIError(format!("{}: {}", status, body))); - } - - let id = format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()); - let model_str = model.to_string(); - - let gemini_sse_stream = response.bytes_stream().eventsource(); - let openai_sse_stream = async_stream::stream! { - tokio::pin!(gemini_sse_stream); - let mut tool_call_index: usize = 0; - while let Some(event) = gemini_sse_stream.next().await { - match event { - Ok(event) => match parse_gemini_sse_event(&event.data) { - Ok(Some(parsed)) => { - for chunk in gemini_event_to_openai_sse_chunks( - &parsed, &id, &model_str, &mut tool_call_index, - ) { - yield Ok::(Bytes::from(chunk)); - } - } - Ok(None) => {} - Err(e) => tracing::error!("Error parsing Gemini SSE event: {}", e), - }, - Err(e) => tracing::error!("Error reading Gemini SSE stream: {}", e), - } - } - yield Ok::(Bytes::from("data: [DONE]\n\n")); - }; - - let mut headers = http::HeaderMap::new(); - headers.insert("content-type", "text/event-stream".parse().unwrap()); - headers.insert("cache-control", "no-cache".parse().unwrap()); - headers.insert("connection", "keep-alive".parse().unwrap()); - - Ok(( - http::StatusCode::OK, - headers, - Body::from_stream(inject_keepalives( - Box::pin(openai_sse_stream), - std::time::Duration::from_secs(KEEPALIVE_INTERVAL_SECS), - )), - )) -} - -// ============================================================================ -// Model listing -// ============================================================================ - -/// List available Gemini models and convert to OpenAI format. -/// -/// - Standard: `GET {base_url}/models` — returns `{ models: [...] }` -/// - Vertex AI: `GET {base_url}` — returns `{ models: [...] }` (base_url already ends with .../models) -pub async fn handle_google_ai_models( - api_key: &str, - base_url: &str, - is_vertex: bool, -) -> Result<(http::StatusCode, http::HeaderMap, Body)> { - #[derive(Deserialize)] - struct GeminiModel { - name: String, - #[serde(rename = "displayName", default)] - display_name: String, - } - - #[derive(Deserialize)] - struct GeminiModelsResponse { - #[serde(default)] - models: Vec, - } - - let base_url = base_url.trim_end_matches('/'); - let endpoint = if is_vertex { - // Vertex AI: base_url is .../publishers/google/models - base_url.to_string() - } else { - // Standard: append /models - format!("{}/models", base_url) - }; - - let request = HTTP_CLIENT.get(&endpoint); - let request = set_auth(request, api_key, is_vertex); - - let response = request - .send() - .await - .map_err(|e| Error::internal_err(format!("Failed to fetch Gemini models: {}", e)))?; - - if let Err(e) = response.error_for_status_ref() { - let status = e.status().map(|s| s.to_string()).unwrap_or_default(); - let body = response.text().await.unwrap_or_default(); - return Err(Error::AIError(format!("{}: {}", status, body))); - } - - let gemini_resp: GeminiModelsResponse = response.json().await.map_err(|e| { - Error::internal_err(format!("Failed to parse Gemini models response: {}", e)) - })?; - - let data: Vec = gemini_resp - .models - .into_iter() - .map(|m| { - json!({ - "id": m.name, - "object": "model", - "display_name": m.display_name, - }) - }) - .collect(); - - let body_bytes = serde_json::to_vec(&json!({ "data": data })) - .map_err(|e| Error::internal_err(format!("Failed to serialize models: {}", e)))?; - - let mut headers = http::HeaderMap::new(); - headers.insert("content-type", "application/json".parse().unwrap()); - - Ok((http::StatusCode::OK, headers, Body::from(body_bytes))) -} - -// ============================================================================ -// Non-streaming path -// ============================================================================ - -async fn handle_non_streaming( - model: &str, - request_body: String, - api_key: &str, - base_url: &str, - is_vertex: bool, -) -> Result<(http::StatusCode, http::HeaderMap, Body)> { - let endpoint = build_model_endpoint(base_url, model, "generateContent", is_vertex); - - let request = HTTP_CLIENT - .post(&endpoint) - .header("content-type", "application/json") - .body(request_body); - let request = set_auth(request, api_key, is_vertex); - - let response = request - .send() - .await - .map_err(|e| Error::internal_err(format!("Failed to send request to Gemini API: {}", e)))?; - - if let Err(e) = response.error_for_status_ref() { - let status = e.status().map(|s| s.to_string()).unwrap_or_default(); - let body = response.text().await.unwrap_or_default(); - return Err(Error::AIError(format!("{}: {}", status, body))); - } - - let body = response - .bytes() - .await - .map_err(|e| Error::internal_err(format!("Failed to read Gemini response body: {}", e)))?; - - let parsed = parse_gemini_response(&body)?; - let openai_response = gemini_response_to_openai(&parsed, model); - - let body_bytes = serde_json::to_vec(&openai_response) - .map_err(|e| Error::internal_err(format!("Failed to serialize response: {}", e)))?; - - let mut headers = http::HeaderMap::new(); - headers.insert("content-type", "application/json".parse().unwrap()); - - Ok((http::StatusCode::OK, headers, Body::from(body_bytes))) -} diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index 75b3469c41..e7e4508971 100644 --- a/backend/windmill-api/src/lib.rs +++ b/backend/windmill-api/src/lib.rs @@ -79,7 +79,6 @@ mod capture; mod concurrency_groups; mod db; mod db_health; -mod google; mod drafts; #[cfg(feature = "private")] diff --git a/docs/windmill-ai-refactor-plan.md b/docs/windmill-ai-refactor-plan.md index 4c0e109735..ea2b9f451e 100644 --- a/docs/windmill-ai-refactor-plan.md +++ b/docs/windmill-ai-refactor-plan.md @@ -37,7 +37,7 @@ Avoid adding modules whose only purpose is to re-export moved code. Direct impor Also do not make `build_proxy_request(raw_body, path)` too narrow. The proxy path needs method, incoming headers, resolved credentials, base URL/platform, organization/user fields, custom headers, and Bedrock/Azure/Vertex-specific context. Introduce a structured `ProxyBuildArgs`/`ProviderCredentials` shape before deleting `AIRequestConfig::prepare_request`, `google.rs`, or `bedrock.rs`. -## Current Phase PR: Proxy Contract + OpenAI-Compatible Proxy +## Completed Phase: Proxy Contract + OpenAI-Compatible Proxy ✅ Goal: introduce the shared API proxy contract in `windmill-ai` and move the OpenAI-compatible proxy request builder there without changing provider behavior. @@ -66,6 +66,41 @@ Validation: - `cargo check -p windmill-ai -p windmill-api` - `cargo check -p windmill-ai -p windmill-api --features bedrock` +Follow-up status: Anthropic/Vertex proxy handling has since moved into +`windmill-ai`, and the dead `AIRequestConfig::prepare_request` fallback has +been removed. + +## Current Phase PR: Proxy Execution Mode + Google AI Proxy Migration + +Goal: introduce a shared provider execution classifier before moving Google AI +and Bedrock. `ProxyRequest` is a good contract for HTTP-forwarding providers +such as OpenAI-compatible providers and Anthropic, but Google AI also converts +responses back to OpenAI shape and Bedrock uses SDK execution. Model that split +explicitly before moving those providers, then move the Google AI proxy +transformation into `windmill-ai` as the first native-provider migration. + +Suggested PR title: `refactor(ai): add provider proxy execution mode`. + +Scope: +- Add `ProxyExecutionMode` in `windmill-ai::proxy`. +- Classify providers as HTTP-forwarding, native Google AI, or native Bedrock. +- Make `supports_query_builder_proxy` derive from the shared execution mode. +- Use the shared execution mode in `windmill-api/src/ai.rs` for workspace proxy routing. +- Move Google AI workspace proxy request conversion, streaming/non-streaming response conversion, and model-list normalization into `windmill-ai::providers::google_ai`. +- Share Google AI `GeminiTextRequest` and generation-config construction between worker agent requests and API proxy requests. +- Delete the API-local `windmill-api/src/google.rs` module. +- Keep global proxy behavior, Bedrock native handling, credential resolution, audit logging, caching, and SSE keepalive behavior unchanged. + +Out of scope: +- Do not move `windmill-api/src/bedrock.rs`. +- Do not unify `AIRequestConfig` and `ProviderWithResource`. + +Validation: +- `cargo test -p windmill-ai google_ai` +- `cargo test -p windmill-ai proxy` +- `cargo test -p windmill-api maps_request_config_to_provider_credentials` +- `cargo test -p windmill-ai anthropic` + ## Step-by-Step Plan Each step produces a compiling, working backend.