refactor: extract google ai logic to windmill-common and use native gemini api in chat proxy (#8115)

* refactor: extract google ai logic to windmill-common and use native gemini api in chat proxy

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix: use x-goog-api-key header for google ai non-chat requests

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix: transform gemini models response to openai format and use correct auth header

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix: skip thought parts from gemini thinking models in sse stream

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Revert "fix: skip thought parts from gemini thinking models in sse stream"

This reverts commit dfa01d282c.

* fix: handle tool calls and sanitize schemas in gemini chat proxy

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* refactor: move Gemini→OpenAI response conversion to windmill-common

Extract streaming and non-streaming Gemini response conversion into
shared functions in ai_google so the API proxy and worker use the same
logic instead of duplicating format translation.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: review fixes for google ai refactor

- Remove duplicate parse_data_url from worker utils, use shared version
  from windmill_common::ai_google in both google_ai and anthropic providers
- Improve error diagnostics in google.rs by including HTTP status code
  in error messages from Gemini API responses
- Change GeminiToolCallEvent::into_extra_content to instance method
  to_extra_content using &self

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* refactor: deduplicate worker Gemini message conversion using pre-flight pattern

Replace the worker's `convert_messages_to_gemini` and
`convert_content_to_parts_with_s3` (~130 lines) with the existing
pre-flight pattern: `prepare_messages_for_api` converts S3 objects to
data URLs, then the shared `openai_messages_to_gemini` handles the rest.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: hugocasa <hugo@casademont.ch>

backend/windmill-api/src/ai.rs

@@ -29,7 +29,7 @@ const AI_TIMEOUT_MAX_SECS: u64 = 86400; // 24 hours
 const AI_TIMEOUT_DEFAULT_SECS: u64 = 3600; // 1 hour
 const HTTP_POOL_MAX_IDLE_PER_HOST: usize = 10;
 const HTTP_POOL_IDLE_TIMEOUT_SECS: u64 = 90;
-const KEEPALIVE_INTERVAL_SECS: u64 = 15;
+pub(crate) const KEEPALIVE_INTERVAL_SECS: u64 = 15;
 
 lazy_static::lazy_static! {
     /// AI request timeout in seconds.
@@ -87,7 +87,7 @@ lazy_static::lazy_static! {
         }
     };
 
-    static ref HTTP_CLIENT: Client = configure_client(reqwest::ClientBuilder::new()
+    pub(crate) static ref HTTP_CLIENT: Client = configure_client(reqwest::ClientBuilder::new()
         .timeout(std::time::Duration::from_secs(*AI_TIMEOUT_SECS))
         .pool_max_idle_per_host(HTTP_POOL_MAX_IDLE_PER_HOST)
         .pool_idle_timeout(Some(std::time::Duration::from_secs(HTTP_POOL_IDLE_TIMEOUT_SECS)))
@@ -378,12 +378,7 @@ impl AIRequestConfig {
         let is_anthropic_sdk = headers.get("X-Anthropic-SDK").is_some();
         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.to_string();
         let base_url = base_url.as_str();
 
         // Build URL based on provider
@@ -428,6 +423,9 @@ 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 {
+                // Native Gemini API uses x-goog-api-key, not Authorization: Bearer
+                request = request.header("x-goog-api-key", api_key.clone())
             } else {
                 request = request.header("authorization", format!("Bearer {}", api_key.clone()))
             }
@@ -611,7 +609,7 @@ fn is_sse_response(headers: &HeaderMap) -> bool {
         .unwrap_or(false)
 }
 
-fn inject_keepalives<S>(
+pub(crate) fn inject_keepalives<S>(
     upstream: S,
     interval: Duration,
 ) -> impl futures::Stream<Item = std::result::Result<Bytes, reqwest::Error>>
@@ -830,6 +828,36 @@ 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 mut tx = db.begin().await?;
+        audit_log(
+            &mut *tx,
+            &authed,
+            "ai.request",
+            ActionKind::Execute,
+            &w_id,
+            Some(&authed.email),
+            Some([("ai_config_path", &format!("{:?}", ai_path)[..])].into()),
+        )
+        .await?;
+        tx.commit().await?;
+
+        return match ai_path.as_str() {
+            "chat/completions" => {
+                crate::google::handle_google_ai_chat(&body, api_key, base_url).await
+            }
+            "models" => crate::google::handle_google_ai_models(api_key, base_url).await,
+            _ => Err(Error::BadRequest(format!(
+                "Unsupported Google AI path: {}",
+                ai_path
+            ))),
+        };
+    }
+
     // Handle Bedrock-specific logic when the feature is enabled
     #[cfg(feature = "bedrock")]
     {

backend/windmill-api/src/google.rs

@@ -0,0 +1,306 @@
+//! 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.
+//!
+//! 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_common::{
+    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,
+    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<OpenAIMessage>,
+    #[serde(default)]
+    stream: bool,
+    #[serde(default)]
+    temperature: Option<f32>,
+    #[serde(default)]
+    max_tokens: Option<u32>,
+    #[serde(default)]
+    tools: Option<Vec<ChatRequestTool>>,
+}
+
+#[derive(Deserialize, Debug)]
+struct ChatRequestTool {
+    function: ChatRequestToolFunction,
+}
+
+#[derive(Deserialize, Debug)]
+struct ChatRequestToolFunction {
+    name: String,
+    #[serde(default)]
+    description: Option<String>,
+    #[serde(default)]
+    parameters: Option<serde_json::Value>,
+}
+
+// ============================================================================
+// 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,
+) -> 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<GeminiFunctionDeclaration> = 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).await
+    } else {
+        handle_non_streaming(&request.model, request_body, api_key, base_url).await
+    }
+}
+
+// ============================================================================
+// Streaming path
+// ============================================================================
+
+async fn handle_streaming(
+    model: &str,
+    request_body: String,
+    api_key: &str,
+    base_url: &str,
+) -> Result<(http::StatusCode, http::HeaderMap, Body)> {
+    let endpoint = format!("{}/models/{}:streamGenerateContent?alt=sse", base_url, model);
+
+    let response = HTTP_CLIENT
+        .post(&endpoint)
+        .header("content-type", "application/json")
+        .header("x-goog-api-key", api_key)
+        .body(request_body)
+        .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, reqwest::Error>(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, reqwest::Error>(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.
+///
+/// Gemini returns `{ models: [{ name: "models/gemini-2.5-flash", displayName, ... }] }`.
+/// The frontend expects OpenAI format `{ data: [{ id: "models/gemini-2.5-flash", ... }] }`.
+pub async fn handle_google_ai_models(
+    api_key: &str,
+    base_url: &str,
+) -> 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<GeminiModel>,
+    }
+
+    let endpoint = format!("{}/models", base_url.trim_end_matches('/'));
+    let response = HTTP_CLIENT
+        .get(&endpoint)
+        .header("x-goog-api-key", api_key)
+        .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<serde_json::Value> = 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,
+) -> Result<(http::StatusCode, http::HeaderMap, Body)> {
+    let endpoint = format!("{}/models/{}:generateContent", base_url, model);
+
+    let response = HTTP_CLIENT
+        .post(&endpoint)
+        .header("content-type", "application/json")
+        .header("x-goog-api-key", api_key)
+        .body(request_body)
+        .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)))
+}

backend/windmill-api/src/lib.rs

@@ -64,6 +64,7 @@ use crate::scim_oss::has_scim_token;
 use windmill_common::error::AppError;
 
 mod ai;
+mod google;
 mod apps;
 pub mod args;
 mod audit;