diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 1ce936a245..3f0ee056f9 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -15849,6 +15849,7 @@ dependencies = [ "dashmap 6.1.0", "datafusion", "ed25519-dalek", + "eventsource-stream", "flate2", "futures", "git-version", diff --git a/backend/windmill-api/Cargo.toml b/backend/windmill-api/Cargo.toml index de8935d097..14069590eb 100644 --- a/backend/windmill-api/Cargo.toml +++ b/backend/windmill-api/Cargo.toml @@ -174,6 +174,7 @@ aws-sdk-bedrock = { workspace = true, optional = true } aws-sdk-bedrockruntime = { workspace = true, optional = true } aws-smithy-types = { workspace = true, optional = true } async-trait.workspace = true +eventsource-stream.workspace = true windmill-jseval.workspace = true tar.workspace = true flate2.workspace = true diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index 730df2a875..92ff49f4a4 100644 --- a/backend/windmill-api/src/ai.rs +++ b/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( +pub(crate) fn inject_keepalives( upstream: S, interval: Duration, ) -> impl futures::Stream> @@ -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")] { diff --git a/backend/windmill-api/src/google.rs b/backend/windmill-api/src/google.rs new file mode 100644 index 0000000000..bc32c20ba4 --- /dev/null +++ b/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, + #[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, +} + +// ============================================================================ +// 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 = 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::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. +/// +/// 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, + } + + 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 = 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))) +} diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index 073c1a8aa1..a25a431dba 100644 --- a/backend/windmill-api/src/lib.rs +++ b/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; diff --git a/backend/windmill-common/src/ai_google.rs b/backend/windmill-common/src/ai_google.rs new file mode 100644 index 0000000000..ccf34685e5 --- /dev/null +++ b/backend/windmill-common/src/ai_google.rs @@ -0,0 +1,726 @@ +//! Shared Google AI (Gemini API) types and conversion utilities. +//! +//! This module provides: +//! - Gemini request/response types +//! - OpenAI → Gemini message conversion +//! - Gemini SSE event parsing +//! +//! Used by both windmill-api (chat proxy) and windmill-worker (AI agent). + +use serde::{Deserialize, Serialize}; + +use crate::ai_types::{ContentPart, ExtraContent, GoogleExtraContent, OpenAIContent, OpenAIMessage, ToolDef, UrlCitation}; +use crate::error::Error; + +// ============================================================================ +// Request / Content Types +// ============================================================================ + +/// Inline data for binary content (images). +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct GeminiInlineData { + #[serde(rename = "mimeType")] + pub mime_type: String, + pub data: String, +} + +/// A part of content — text, inline data, function call, or function response. +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(untagged)] +pub enum GeminiPart { + Text { + text: String, + }, + InlineData { + #[serde(rename = "inlineData")] + inline_data: GeminiInlineData, + }, + FunctionCall { + #[serde(rename = "functionCall")] + function_call: GeminiFunctionCall, + /// Thought signature for Gemini 3+ models — required when replaying function calls. + #[serde(rename = "thoughtSignature", skip_serializing_if = "Option::is_none")] + thought_signature: Option, + }, + FunctionResponse { + #[serde(rename = "functionResponse")] + function_response: GeminiFunctionResponse, + }, +} + +/// A function call from the model. +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct GeminiFunctionCall { + pub name: String, + pub args: serde_json::Value, +} + +/// A function response sent back to the model. +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct GeminiFunctionResponse { + pub name: String, + pub response: serde_json::Value, +} + +/// Content message with an optional role and a list of parts. +#[derive(Serialize, Clone, Debug)] +pub struct GeminiContentMessage { + #[serde(skip_serializing_if = "Option::is_none")] + pub role: Option, + pub parts: Vec, +} + +/// Main request body for `generateContent` / `streamGenerateContent`. +#[derive(Serialize)] +pub struct GeminiTextRequest { + pub contents: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub tools: Option>, + #[serde(rename = "toolConfig", skip_serializing_if = "Option::is_none")] + pub tool_config: Option, + #[serde(rename = "systemInstruction", skip_serializing_if = "Option::is_none")] + pub system_instruction: Option, + #[serde(rename = "generationConfig", skip_serializing_if = "Option::is_none")] + pub generation_config: Option, +} + +/// Tool definition — function declarations and/or Google Search grounding. +#[derive(Serialize)] +pub struct GeminiTool { + #[serde(rename = "functionDeclarations", skip_serializing_if = "Option::is_none")] + pub function_declarations: Option>, + #[serde(rename = "googleSearch", skip_serializing_if = "Option::is_none")] + pub google_search: Option, +} + +/// A single function declaration. +/// +/// `parameters` holds a pre-serialized (and, for the worker, pre-sanitized) JSON Schema. +#[derive(Serialize)] +pub struct GeminiFunctionDeclaration { + pub name: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + pub parameters: serde_json::Value, +} + +/// Tool configuration controlling when and how functions are called. +#[derive(Serialize)] +pub struct GeminiToolConfig { + #[serde(rename = "functionCallingConfig")] + pub function_calling_config: GeminiFunctionCallingConfig, +} + +/// Function calling mode and optional allow-list. +#[derive(Serialize)] +pub struct GeminiFunctionCallingConfig { + pub mode: String, + #[serde(rename = "allowedFunctionNames", skip_serializing_if = "Option::is_none")] + pub allowed_function_names: Option>, +} + +/// Generation parameters (temperature, token limits, structured output). +#[derive(Serialize)] +pub struct GeminiGenerationConfig { + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + #[serde(rename = "maxOutputTokens", skip_serializing_if = "Option::is_none")] + pub max_output_tokens: Option, + #[serde(rename = "responseMimeType", skip_serializing_if = "Option::is_none")] + pub response_mime_type: Option, + #[serde(rename = "responseSchema", skip_serializing_if = "Option::is_none")] + pub response_schema: Option, +} + +// ============================================================================ +// Image Generation Types +// ============================================================================ + +/// Request body for Imagen / Gemini image generation. +#[derive(Serialize)] +pub struct GeminiImageRequest { + #[serde(skip_serializing_if = "Option::is_none")] + pub contents: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub instances: Option>, +} + +/// Content wrapper used in `generateContent` image requests. +#[derive(Serialize)] +pub struct GeminiImageContent { + pub parts: Vec, +} + +/// Prompt wrapper for Imagen `predict` endpoint. +#[derive(Serialize)] +pub struct GeminiPredictContent { + pub prompt: String, +} + +/// Top-level response from Gemini/Imagen image generation. +#[derive(Deserialize)] +pub struct GeminiImageResponse { + pub candidates: Option>, + pub predictions: Option>, +} + +#[derive(Deserialize)] +pub struct GeminiImageCandidate { + pub content: GeminiImageCandidateContent, +} + +#[derive(Deserialize)] +pub struct GeminiImageCandidateContent { + pub parts: Vec, +} + +#[derive(Deserialize)] +pub struct GeminiImageCandidatePart { + #[serde(rename = "inlineData")] + pub inline_data: Option, +} + +#[derive(Deserialize)] +pub struct GeminiPredictCandidate { + #[serde(rename = "bytesBase64Encoded")] + pub bytes_base64_encoded: String, +} + +// ============================================================================ +// SSE Response Types +// ============================================================================ + +/// One part inside a streaming candidate — text, function call, or thought signature. +#[derive(Deserialize, Debug)] +pub struct GeminiSSEPart { + #[serde(default)] + pub text: Option, + #[serde(rename = "functionCall")] + pub function_call: Option, + /// Thought signature for Gemini 3+ models. + #[serde(rename = "thoughtSignature")] + pub thought_signature: Option, +} + +/// Function call contained in a streaming part. +#[derive(Deserialize, Debug)] +pub struct GeminiSSEFunctionCall { + pub name: String, + pub args: serde_json::Value, +} + +/// Content block inside a streaming candidate. +#[derive(Deserialize, Debug)] +pub struct GeminiSSEContent { + pub parts: Option>, +} + +/// Web source from a Gemini grounding chunk. +#[derive(Deserialize, Debug)] +pub struct GeminiGroundingChunkWeb { + pub uri: String, + #[serde(default)] + pub title: Option, +} + +/// One grounding chunk (search result) from Gemini web search. +#[derive(Deserialize, Debug)] +pub struct GeminiGroundingChunk { + pub web: Option, +} + +/// Grounding metadata attached to a streaming candidate. +#[derive(Deserialize, Debug)] +pub struct GeminiGroundingMetadata { + #[serde(rename = "groundingChunks", default)] + pub grounding_chunks: Vec, + #[serde(rename = "webSearchQueries", default)] + pub web_search_queries: Vec, +} + +/// One candidate inside a streaming Gemini response. +#[derive(Deserialize, Debug)] +pub struct GeminiSSECandidate { + pub content: Option, + #[serde(rename = "finishReason")] + pub finish_reason: Option, + #[serde(rename = "groundingMetadata")] + pub grounding_metadata: Option, +} + +/// Token usage from the `usageMetadata` field of a Gemini SSE event. +#[derive(Deserialize, Debug, Clone)] +pub struct GeminiUsageMetadata { + #[serde(rename = "promptTokenCount", default)] + pub prompt_token_count: Option, + #[serde(rename = "candidatesTokenCount", default)] + pub candidates_token_count: Option, + #[serde(rename = "totalTokenCount", default)] + pub total_token_count: Option, +} + +/// Top-level structure of one Gemini SSE event. +#[derive(Deserialize, Debug)] +pub struct GeminiSSEEvent { + pub candidates: Option>, + #[serde(rename = "usageMetadata")] + pub usage_metadata: Option, +} + +// ============================================================================ +// Parsed Event Result +// ============================================================================ + +/// A single function call extracted from a Gemini SSE event. +#[derive(Debug)] +pub struct GeminiToolCallEvent { + pub name: String, + pub args: serde_json::Value, + pub thought_signature: Option, +} + +impl GeminiToolCallEvent { + /// Convert the thought signature (if present) into an [`ExtraContent`]. + pub fn to_extra_content(&self) -> Option { + self.thought_signature.as_ref().map(|sig| ExtraContent { + google: Some(GoogleExtraContent { thought_signature: Some(sig.clone()) }), + }) + } +} + +/// Structured result of parsing a Gemini response (streaming SSE event or non-streaming body). +#[derive(Debug, Default)] +pub struct GeminiParsedEvent { + pub text: Option, + pub tool_calls: Vec, + pub annotations: Vec, + pub used_websearch: bool, + pub usage: Option, + pub finish_reason: Option, +} + +// ============================================================================ +// Helper Functions +// ============================================================================ + +/// Parse a data URL into `(mime_type, base64_data)`. +/// +/// Expected format: `data:;base64,`. +pub fn parse_data_url(url: &str) -> Option<(String, String)> { + let rest = url.strip_prefix("data:")?; + let (header, data) = rest.split_once(',')?; + let media_type = header.strip_suffix(";base64")?; + Some((media_type.to_string(), data.to_string())) +} + +/// Find the function name associated with a `tool_call_id` by scanning prior messages. +pub fn find_gemini_function_name(messages: &[OpenAIMessage], tool_call_id: &str) -> String { + messages + .iter() + .filter_map(|msg| msg.tool_calls.as_ref()) + .flatten() + .find(|tc| tc.id == tool_call_id) + .map(|tc| tc.function.name.clone()) + .unwrap_or_else(|| "unknown_function".to_string()) +} + +/// Convert an [`OpenAIContent`] value to a list of [`GeminiPart`]s. +/// +/// Handles text and `image_url` (data URLs). `S3Object` variants are skipped here; +/// the worker handles them by downloading and injecting inline data beforehand. +pub fn convert_content_to_gemini_parts(content: &OpenAIContent) -> Vec { + match content { + OpenAIContent::Text(text) if !text.is_empty() => { + vec![GeminiPart::Text { text: text.clone() }] + } + OpenAIContent::Text(_) => vec![], + OpenAIContent::Parts(parts) => parts + .iter() + .filter_map(|part| match part { + ContentPart::Text { text } if !text.is_empty() => { + Some(GeminiPart::Text { text: text.clone() }) + } + ContentPart::ImageUrl { image_url } => { + parse_data_url(&image_url.url).map(|(mime_type, data)| { + GeminiPart::InlineData { + inline_data: GeminiInlineData { mime_type, data }, + } + }) + } + // S3Objects are handled by the worker + _ => None, + }) + .collect(), + } +} + +/// Convert OpenAI-format messages to Gemini `contents` and an optional `systemInstruction`. +/// +/// Returns `(contents, system_instruction)`. +/// +/// `S3Object` images in content parts are skipped (the worker pre-converts them). +/// Tool call history is preserved correctly for multi-turn agent conversations. +pub fn openai_messages_to_gemini( + messages: &[OpenAIMessage], +) -> (Vec, Option) { + let mut contents: Vec = Vec::new(); + let mut system_instruction: Option = None; + + for msg in messages { + match msg.role.as_str() { + "system" => { + if let Some(content) = &msg.content { + let parts = convert_content_to_gemini_parts(content); + if !parts.is_empty() { + system_instruction = + Some(GeminiContentMessage { role: None, parts }); + } + } + } + "tool" => { + if let (Some(tool_call_id), Some(content)) = + (&msg.tool_call_id, &msg.content) + { + let func_name = find_gemini_function_name(messages, tool_call_id); + let response_text = match content { + OpenAIContent::Text(text) => text.clone(), + OpenAIContent::Parts(parts) => parts + .iter() + .filter_map(|p| { + if let ContentPart::Text { text } = p { + Some(text.as_str()) + } else { + None + } + }) + .collect::>() + .join(" "), + }; + contents.push(GeminiContentMessage { + role: Some("user".to_string()), + parts: vec![GeminiPart::FunctionResponse { + function_response: GeminiFunctionResponse { + name: func_name, + response: serde_json::json!({ "result": response_text }), + }, + }], + }); + } + } + role => { + let gemini_role = if role == "assistant" { "model" } else { "user" }; + let mut parts: Vec = Vec::new(); + + if let Some(content) = &msg.content { + parts.extend(convert_content_to_gemini_parts(content)); + } + + if let Some(tool_calls) = &msg.tool_calls { + for tc in tool_calls { + let args: serde_json::Value = + serde_json::from_str(&tc.function.arguments).unwrap_or_default(); + let thought_signature = tc + .extra_content + .as_ref() + .and_then(|ec| ec.google.as_ref()) + .and_then(|g| g.thought_signature.clone()); + parts.push(GeminiPart::FunctionCall { + function_call: GeminiFunctionCall { + name: tc.function.name.clone(), + args, + }, + thought_signature, + }); + } + } + + if !parts.is_empty() { + contents.push(GeminiContentMessage { + role: Some(gemini_role.to_string()), + parts, + }); + } + } + } + } + + (contents, system_instruction) +} + +/// Convert OpenAI tool definitions to Gemini format. +/// +/// `tool_params` must be pre-serialized (and, for the worker, pre-sanitized for Google) +/// JSON schema values, one per entry in `tools` in the same order. +pub fn openai_tools_to_gemini( + tools: &[ToolDef], + tool_params: &[serde_json::Value], + has_websearch: bool, +) -> Option> { + let mut gemini_tools: Vec = Vec::new(); + + let declarations: Vec = tools + .iter() + .zip(tool_params.iter()) + .map(|(t, params)| GeminiFunctionDeclaration { + name: t.function.name.clone(), + description: t.function.description.clone(), + parameters: params.clone(), + }) + .collect(); + + if !declarations.is_empty() { + gemini_tools.push(GeminiTool { + function_declarations: Some(declarations), + google_search: None, + }); + } + + if has_websearch { + gemini_tools.push(GeminiTool { + function_declarations: None, + google_search: Some(serde_json::json!({})), + }); + } + + if gemini_tools.is_empty() { + None + } else { + Some(gemini_tools) + } +} + +/// Parse one Gemini SSE data line into a [`GeminiParsedEvent`]. +/// +/// Returns `Ok(None)` for empty data or unrecognised payloads (e.g. `"[DONE]"`). +/// Logs a warning and returns `Ok(None)` on JSON parse errors rather than propagating. +pub fn parse_gemini_sse_event(data: &str) -> Result, Error> { + if data.is_empty() || data == "[DONE]" { + return Ok(None); + } + + let event: GeminiSSEEvent = match serde_json::from_str(data) { + Ok(e) => e, + Err(e) => { + tracing::error!("Failed to parse Gemini SSE event {}: {}", data, e); + return Ok(None); + } + }; + + let mut parsed = GeminiParsedEvent { usage: event.usage_metadata, ..Default::default() }; + + let Some(candidates) = event.candidates else { + return Ok(Some(parsed)); + }; + + extract_candidates_into(&candidates, &mut parsed); + + Ok(Some(parsed)) +} + +/// Parse a non-streaming Gemini `generateContent` response body. +pub fn parse_gemini_response(data: &[u8]) -> Result { + let event: GeminiSSEEvent = serde_json::from_slice(data) + .map_err(|e| Error::internal_err(format!("Failed to parse Gemini response: {}", e)))?; + + let mut parsed = GeminiParsedEvent { usage: event.usage_metadata, ..Default::default() }; + + if let Some(candidates) = event.candidates { + extract_candidates_into(&candidates, &mut parsed); + } + + Ok(parsed) +} + +// ============================================================================ +// Gemini → OpenAI Format Conversion +// ============================================================================ + +/// Convert a `GeminiParsedEvent` from a non-streaming response to an OpenAI chat completion JSON. +pub fn gemini_response_to_openai(parsed: &GeminiParsedEvent, model: &str) -> serde_json::Value { + let content = parsed.text.as_deref().unwrap_or_default(); + + let tool_calls: Vec = parsed + .tool_calls + .iter() + .enumerate() + .map(|(i, tc)| { + serde_json::json!({ + "index": i, + "id": format!("call_{}", uuid::Uuid::new_v4().simple()), + "type": "function", + "function": { + "name": tc.name, + "arguments": serde_json::to_string(&tc.args).unwrap_or_default() + } + }) + }) + .collect(); + + let finish_reason = parsed + .finish_reason + .as_deref() + .map(|r| r.to_lowercase()) + .unwrap_or_else(|| "stop".to_string()); + + let usage = parsed.usage.as_ref().map(|u| { + serde_json::json!({ + "prompt_tokens": u.prompt_token_count.unwrap_or(0), + "completion_tokens": u.candidates_token_count.unwrap_or(0), + "total_tokens": u.total_token_count.unwrap_or(0), + }) + }); + + let mut message = serde_json::json!({ + "role": "assistant", + "content": content, + }); + if !tool_calls.is_empty() { + message["tool_calls"] = serde_json::json!(tool_calls); + } + + serde_json::json!({ + "id": format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()), + "object": "chat.completion", + "model": model, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": finish_reason, + }], + "usage": usage, + }) +} + +/// Convert a `GeminiParsedEvent` from a streaming SSE event into OpenAI-format SSE lines. +/// +/// Returns the serialized `"data: {...}\n\n"` lines ready to be written to the response stream. +/// `tool_call_index` is mutated to track the running index across multiple SSE events. +pub fn gemini_event_to_openai_sse_chunks( + parsed: &GeminiParsedEvent, + id: &str, + model: &str, + tool_call_index: &mut usize, +) -> Vec { + let mut chunks = Vec::new(); + + if let Some(text) = &parsed.text { + let chunk = serde_json::json!({ + "id": id, + "object": "chat.completion.chunk", + "model": model, + "choices": [{ + "index": 0, + "delta": { "content": text }, + "finish_reason": null, + }] + }); + chunks.push(format!("data: {}\n\n", chunk)); + } + + for tc in &parsed.tool_calls { + let args_str = serde_json::to_string(&tc.args).unwrap_or_default(); + let call_id = format!("call_{}", uuid::Uuid::new_v4().simple()); + let chunk = serde_json::json!({ + "id": id, + "object": "chat.completion.chunk", + "model": model, + "choices": [{ + "index": 0, + "delta": { + "tool_calls": [{ + "index": *tool_call_index, + "id": call_id, + "type": "function", + "function": { + "name": tc.name, + "arguments": args_str, + } + }] + }, + "finish_reason": null, + }] + }); + chunks.push(format!("data: {}\n\n", chunk)); + *tool_call_index += 1; + } + + chunks +} + +/// Recursively remove JSON Schema fields unsupported by the Gemini API. +pub fn sanitize_schema_for_google(value: &mut serde_json::Value) { + const UNSUPPORTED: &[&str] = &[ + "additionalProperties", + "strict", + "$schema", + "default", + "exclusiveMinimum", + "exclusiveMaximum", + "const", + "multipleOf", + ]; + + if let Some(obj) = value.as_object_mut() { + for field in UNSUPPORTED { + obj.remove(*field); + } + for v in obj.values_mut() { + sanitize_schema_for_google(v); + } + } else if let Some(arr) = value.as_array_mut() { + for v in arr.iter_mut() { + sanitize_schema_for_google(v); + } + } +} + +// ============================================================================ +// Internal Helpers +// ============================================================================ + +fn extract_candidates_into(candidates: &[GeminiSSECandidate], parsed: &mut GeminiParsedEvent) { + for candidate in candidates { + if let Some(content) = &candidate.content { + if let Some(parts) = &content.parts { + for part in parts { + if let Some(text) = &part.text { + if !text.is_empty() { + match parsed.text.as_mut() { + Some(existing) => existing.push_str(text), + None => parsed.text = Some(text.clone()), + } + } + } + + if let Some(function_call) = &part.function_call { + parsed.tool_calls.push(GeminiToolCallEvent { + name: function_call.name.clone(), + args: function_call.args.clone(), + thought_signature: part.thought_signature.clone(), + }); + } + } + } + } + + if candidate.finish_reason.is_some() { + parsed.finish_reason = candidate.finish_reason.clone(); + } + + if let Some(grounding) = &candidate.grounding_metadata { + if !grounding.web_search_queries.is_empty() || !grounding.grounding_chunks.is_empty() { + parsed.used_websearch = true; + } + for chunk in &grounding.grounding_chunks { + if let Some(web) = &chunk.web { + parsed.annotations.push(UrlCitation { + start_index: 0, + end_index: 0, + url: web.uri.clone(), + title: web.title.clone(), + }); + } + } + } + } +} diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index c4e0a47cd8..85f5563419 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -29,6 +29,7 @@ use sqlx::{Acquire, Postgres}; pub mod agent_workers; #[cfg(feature = "bedrock")] pub mod ai_bedrock; +pub mod ai_google; pub mod ai_providers; pub mod ai_types; pub mod apps; diff --git a/backend/windmill-worker/src/ai/providers/anthropic.rs b/backend/windmill-worker/src/ai/providers/anthropic.rs index e42e1dc49b..f9f5edf452 100644 --- a/backend/windmill-worker/src/ai/providers/anthropic.rs +++ b/backend/windmill-worker/src/ai/providers/anthropic.rs @@ -1,14 +1,16 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json::value::RawValue; -use windmill_common::{ai_providers::AIProvider, client::AuthedClient, error::Error}; +use windmill_common::{ + ai_google::parse_data_url, ai_providers::AIProvider, client::AuthedClient, error::Error, +}; use crate::ai::{ image_handler::prepare_messages_for_api, query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventProcessor}, sse::{AnthropicSSEParser, SSEParser}, types::*, - utils::{extract_text_content, parse_data_url, should_use_structured_output_tool}, + utils::{extract_text_content, should_use_structured_output_tool}, }; /// Anthropic API version for standard API diff --git a/backend/windmill-worker/src/ai/providers/google_ai.rs b/backend/windmill-worker/src/ai/providers/google_ai.rs index 62e6afff75..e2030f3269 100644 --- a/backend/windmill-worker/src/ai/providers/google_ai.rs +++ b/backend/windmill-worker/src/ai/providers/google_ai.rs @@ -1,215 +1,21 @@ use async_trait::async_trait; -use serde::{Deserialize, Serialize}; -use windmill_common::{client::AuthedClient, error::Error}; +use windmill_common::{ + ai_google::{ + openai_messages_to_gemini, openai_tools_to_gemini, GeminiGenerationConfig, + GeminiImageContent, GeminiImageRequest, GeminiImageResponse, GeminiInlineData, GeminiPart, + GeminiPredictContent, GeminiTextRequest, GeminiTool, + }, + client::AuthedClient, + error::Error, +}; use crate::ai::{ - image_handler::download_and_encode_s3_image, + image_handler::{download_and_encode_s3_image, prepare_messages_for_api}, query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventProcessor}, sse::{GeminiSSEParser, SSEParser}, types::*, - utils::parse_data_url, }; -// ============================================================================ -// Gemini API Types - Shared between text and image -// ============================================================================ - -/// Inline data for binary content (images) -#[derive(Serialize, Deserialize, Clone, Debug)] -pub struct GeminiInlineData { - #[serde(rename = "mimeType")] - pub mime_type: String, - pub data: String, -} - -/// A part of content - can be text, inline data, function call, or function response -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(untagged)] -pub enum GeminiPart { - Text { - text: String, - }, - InlineData { - #[serde(rename = "inlineData")] - inline_data: GeminiInlineData, - }, - FunctionCall { - #[serde(rename = "functionCall")] - function_call: GeminiFunctionCall, - /// Thought signature for Gemini 3+ models - required for function calling - #[serde(rename = "thoughtSignature", skip_serializing_if = "Option::is_none")] - thought_signature: Option, - }, - FunctionResponse { - #[serde(rename = "functionResponse")] - function_response: GeminiFunctionResponse, - }, -} - -/// A function call from the model -#[derive(Serialize, Deserialize, Clone, Debug)] -pub struct GeminiFunctionCall { - pub name: String, - pub args: serde_json::Value, -} - -/// A function response to send back to the model -#[derive(Serialize, Deserialize, Clone, Debug)] -pub struct GeminiFunctionResponse { - pub name: String, - pub response: serde_json::Value, -} - -// ============================================================================ -// Gemini Text API Request Types -// ============================================================================ - -/// Main request structure for Gemini generateContent -#[derive(Serialize)] -pub struct GeminiTextRequest { - pub contents: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - pub tools: Option>, - #[serde(rename = "toolConfig", skip_serializing_if = "Option::is_none")] - pub tool_config: Option, - #[serde(rename = "systemInstruction", skip_serializing_if = "Option::is_none")] - pub system_instruction: Option, - #[serde(rename = "generationConfig", skip_serializing_if = "Option::is_none")] - pub generation_config: Option, -} - -/// Content message with role and parts -#[derive(Serialize)] -pub struct GeminiContentMessage { - #[serde(skip_serializing_if = "Option::is_none")] - pub role: Option, - pub parts: Vec, -} - -/// Tool definition - either function declarations or Google Search -#[derive(Serialize)] -pub struct GeminiTool { - #[serde( - rename = "functionDeclarations", - skip_serializing_if = "Option::is_none" - )] - pub function_declarations: Option>, - #[serde(rename = "googleSearch", skip_serializing_if = "Option::is_none")] - pub google_search: Option, -} - -/// Function declaration for tool use -#[derive(Serialize)] -pub struct GeminiFunctionDeclaration { - pub name: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub description: Option, - pub parameters: OpenAPISchema, -} - -/// Tool configuration for controlling function calling behavior -#[derive(Serialize)] -pub struct GeminiToolConfig { - #[serde(rename = "functionCallingConfig")] - pub function_calling_config: GeminiFunctionCallingConfig, -} - -/// Function calling configuration -#[derive(Serialize)] -pub struct GeminiFunctionCallingConfig { - pub mode: String, - #[serde( - rename = "allowedFunctionNames", - skip_serializing_if = "Option::is_none" - )] - pub allowed_function_names: Option>, -} - -/// Generation configuration for output format -#[derive(Serialize)] -pub struct GeminiGenerationConfig { - #[serde(skip_serializing_if = "Option::is_none")] - pub temperature: Option, - #[serde(rename = "maxOutputTokens", skip_serializing_if = "Option::is_none")] - pub max_output_tokens: Option, - #[serde(rename = "responseMimeType", skip_serializing_if = "Option::is_none")] - pub response_mime_type: Option, - #[serde(rename = "responseSchema", skip_serializing_if = "Option::is_none")] - pub response_schema: Option, -} - -// ============================================================================ -// Gemini API Response Types -// ============================================================================ - -/// Grounding metadata from Google Search -#[derive(Deserialize)] -#[allow(dead_code)] -pub struct GeminiGroundingMetadata { - #[serde(rename = "webSearchQueries")] - pub web_search_queries: Option>, - #[serde(rename = "groundingChunks")] - pub grounding_chunks: Option>, -} - -// ============================================================================ -// Gemini Image API Types (for Imagen models) -// ============================================================================ - -/// Request for image generation (Imagen models) -#[derive(Serialize)] -pub struct GeminiImageRequest { - #[serde(skip_serializing_if = "Option::is_none")] - pub contents: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub instances: Option>, -} - -/// Content for image generation -#[derive(Serialize)] -pub struct GeminiImageContent { - pub parts: Vec, -} - -/// Content for Imagen predict endpoint -#[derive(Serialize)] -pub struct GeminiPredictContent { - pub prompt: String, -} - -/// Response for image generation -#[derive(Deserialize)] -pub struct GeminiImageResponse { - pub candidates: Option>, - pub predictions: Option>, -} - -/// Image candidate from generateContent -#[derive(Deserialize)] -pub struct GeminiImageCandidate { - pub content: GeminiImageCandidateContent, -} - -/// Content in image candidate -#[derive(Deserialize)] -pub struct GeminiImageCandidateContent { - pub parts: Vec, -} - -/// Part of image candidate -#[derive(Deserialize)] -pub struct GeminiImageCandidatePart { - #[serde(rename = "inlineData", skip_serializing_if = "Option::is_none")] - pub inline_data: Option, -} - -/// Prediction candidate from Imagen -#[derive(Deserialize)] -pub struct GeminiPredictCandidate { - #[serde(rename = "bytesBase64Encoded")] - pub bytes_base64_encoded: String, -} - // ============================================================================ // Query Builder Implementation // ============================================================================ @@ -221,34 +27,24 @@ impl GoogleAIQueryBuilder { Self } - /// Build a text request using the native Gemini API format async fn build_text_request( &self, args: &BuildRequestArgs<'_>, client: &AuthedClient, workspace_id: &str, ) -> Result { - // Convert messages to Gemini format - let contents = self - .convert_messages_to_gemini(args.messages, client, workspace_id) - .await?; + let prepared_messages = + prepare_messages_for_api(args.messages, client, workspace_id).await?; + let (contents, system_instruction) = openai_messages_to_gemini(&prepared_messages); - // Build tools array let tools = self.convert_tools_to_gemini(args.tools, args.has_websearch); - // Build generation config let generation_config = self.build_generation_config(args); - // Build system instruction from system_prompt - let system_instruction = args.system_prompt.map(|s| GeminiContentMessage { - role: None, - parts: vec![GeminiPart::Text { text: s.to_string() }], - }); - let request = GeminiTextRequest { contents, tools, - tool_config: None, // Use AUTO mode by default + tool_config: None, system_instruction, generation_config, }; @@ -257,7 +53,6 @@ impl GoogleAIQueryBuilder { .map_err(|e| Error::internal_err(format!("Failed to serialize request: {}", e))) } - /// Build an image generation request async fn build_image_request( &self, args: &BuildRequestArgs<'_>, @@ -267,7 +62,6 @@ impl GoogleAIQueryBuilder { let is_imagen = args.model.contains("imagen"); let request = if is_imagen { - // For Imagen models, use simple prompt format GeminiImageRequest { instances: Some(vec![GeminiPredictContent { prompt: args.user_message.trim().to_string(), @@ -275,7 +69,6 @@ impl GoogleAIQueryBuilder { contents: None, } } else { - // For Gemini models with image generation, build parts let mut parts = vec![GeminiPart::Text { text: args.user_message.trim().to_string() }]; if let Some(system_prompt) = args.system_prompt { @@ -285,7 +78,6 @@ impl GoogleAIQueryBuilder { ); } - // Add input images if provided if let Some(images) = args.images { for image in images.iter() { if !image.s3.is_empty() { @@ -308,218 +100,39 @@ impl GoogleAIQueryBuilder { .map_err(|e| Error::internal_err(format!("Failed to serialize request: {}", e))) } - /// Convert OpenAI-format messages to Gemini format - async fn convert_messages_to_gemini( - &self, - messages: &[OpenAIMessage], - client: &AuthedClient, - workspace_id: &str, - ) -> Result, Error> { - let mut gemini_messages = Vec::new(); - - for msg in messages { - match msg.role.as_str() { - "system" => { - // Skip - handled via args.system_prompt in build_text_request - } - "tool" => { - // Handle tool responses - if let (Some(tool_call_id), Some(content)) = (&msg.tool_call_id, &msg.content) { - let func_name = self.find_function_name_by_id(messages, tool_call_id); - let response_text = match content { - OpenAIContent::Text(text) => text.clone(), - OpenAIContent::Parts(parts) => parts - .iter() - .filter_map(|p| match p { - ContentPart::Text { text } => Some(text.clone()), - _ => None, - }) - .collect::>() - .join(" "), - }; - - gemini_messages.push(GeminiContentMessage { - role: Some("user".to_string()), - parts: vec![GeminiPart::FunctionResponse { - function_response: GeminiFunctionResponse { - name: func_name, - response: serde_json::json!({ "result": response_text }), - }, - }], - }); - } - } - _ => { - // Handle user/assistant messages - let role = match msg.role.as_str() { - "assistant" => "model", - _ => "user", - }; - - let mut parts = Vec::new(); - - // Handle regular content - if let Some(content) = &msg.content { - let content_parts = self - .convert_content_to_parts(&Some(content.clone()), client, workspace_id) - .await?; - parts.extend(content_parts); - } - - // Handle tool calls from assistant - if let Some(tool_calls) = &msg.tool_calls { - for tc in tool_calls { - let args: serde_json::Value = - serde_json::from_str(&tc.function.arguments).unwrap_or_default(); - // Extract thought_signature from extra_content if present - let thought_signature = tc - .extra_content - .as_ref() - .and_then(|ec| ec.google.as_ref()) - .and_then(|g| g.thought_signature.clone()); - parts.push(GeminiPart::FunctionCall { - function_call: GeminiFunctionCall { - name: tc.function.name.clone(), - args, - }, - thought_signature, - }); - } - } - - if !parts.is_empty() { - gemini_messages - .push(GeminiContentMessage { role: Some(role.to_string()), parts }); - } - } - } - } - - Ok(gemini_messages) - } - - /// Convert OpenAI content to Gemini parts - async fn convert_content_to_parts( - &self, - content: &Option, - client: &AuthedClient, - workspace_id: &str, - ) -> Result, Error> { - let mut parts = Vec::new(); - - if let Some(content) = content { - match content { - OpenAIContent::Text(text) => { - if !text.is_empty() { - parts.push(GeminiPart::Text { text: text.clone() }); - } - } - OpenAIContent::Parts(content_parts) => { - for part in content_parts { - match part { - ContentPart::Text { text } => { - if !text.is_empty() { - parts.push(GeminiPart::Text { text: text.clone() }); - } - } - ContentPart::ImageUrl { image_url } => { - // Parse data URL format: data:mime_type;base64,data - if let Some((mime_type, data)) = parse_data_url(&image_url.url) { - parts.push(GeminiPart::InlineData { - inline_data: GeminiInlineData { mime_type, data }, - }); - } - } - ContentPart::S3Object { s3_object } => { - if !s3_object.s3.is_empty() { - let (mime_type, data) = download_and_encode_s3_image( - s3_object, - client, - workspace_id, - ) - .await?; - parts.push(GeminiPart::InlineData { - inline_data: GeminiInlineData { mime_type, data }, - }); - } - } - } - } - } - } - } - - Ok(parts) - } - - /// Find function name by tool call ID from previous messages - fn find_function_name_by_id(&self, messages: &[OpenAIMessage], tool_call_id: &str) -> String { - for msg in messages { - if let Some(tool_calls) = &msg.tool_calls { - for tc in tool_calls { - if tc.id == tool_call_id { - return tc.function.name.clone(); - } - } - } - } - "unknown_function".to_string() - } - - /// Convert OpenAI tools to Gemini format + /// Convert OpenAI tool definitions to Gemini format. + /// + /// Sanitizes each tool's JSON schema for Google compatibility before delegating + /// to the shared [`openai_tools_to_gemini`] function. fn convert_tools_to_gemini( &self, tools: Option<&[ToolDef]>, has_websearch: bool, ) -> Option> { - let mut gemini_tools = Vec::new(); - - // Add function declarations - if let Some(tool_defs) = tools { - let declarations: Vec = tool_defs - .iter() - .filter_map(|t| { - // Deserialize RawValue into OpenAPISchema, sanitize, then use - let mut schema: OpenAPISchema = - serde_json::from_str(t.function.parameters.get()).ok()?; - schema.sanitize_for_google(); - - Some(GeminiFunctionDeclaration { - name: t.function.name.clone(), - description: t.function.description.clone(), - parameters: schema, - }) - }) - .collect(); - - if !declarations.is_empty() { - gemini_tools.push(GeminiTool { - function_declarations: Some(declarations), - google_search: None, - }); + let Some(tool_defs) = tools else { + if has_websearch { + return Some(vec![GeminiTool { + function_declarations: None, + google_search: Some(serde_json::json!({})), + }]); } - } + return None; + }; - // Add Google Search tool if enabled - if has_websearch { - gemini_tools.push(GeminiTool { - function_declarations: None, - google_search: Some(serde_json::json!({})), - }); - } + let tool_params: Vec = tool_defs + .iter() + .map(|t| { + let mut schema: OpenAPISchema = + serde_json::from_str(t.function.parameters.get()).unwrap_or_default(); + schema.sanitize_for_google(); + serde_json::to_value(&schema).unwrap_or_default() + }) + .collect(); - if gemini_tools.is_empty() { - None - } else { - Some(gemini_tools) - } + openai_tools_to_gemini(tool_defs, &tool_params, has_websearch) } - /// Build generation config for structured output and other settings - fn build_generation_config( - &self, - args: &BuildRequestArgs<'_>, - ) -> Option { + fn build_generation_config(&self, args: &BuildRequestArgs<'_>) -> Option { let has_output_schema = args .output_schema .and_then(|s| s.properties.as_ref()) @@ -529,15 +142,11 @@ impl GoogleAIQueryBuilder { let (response_mime_type, response_schema) = if has_output_schema { let mut schema = args.output_schema.unwrap().clone(); schema.sanitize_for_google(); - ( - Some("application/json".to_string()), - serde_json::to_value(&schema).ok(), - ) + (Some("application/json".to_string()), serde_json::to_value(&schema).ok()) } else { (None, None) }; - // Only create config if there's something to configure if args.temperature.is_some() || args.max_tokens.is_some() || response_mime_type.is_some() { Some(GeminiGenerationConfig { temperature: args.temperature, @@ -554,7 +163,6 @@ impl GoogleAIQueryBuilder { #[async_trait] impl QueryBuilder for GoogleAIQueryBuilder { fn supports_tools_with_output_type(&self, output_type: &OutputType) -> bool { - // Google AI supports tools only for text output matches!(output_type, OutputType::Text) } @@ -578,7 +186,6 @@ impl QueryBuilder for GoogleAIQueryBuilder { Error::internal_err(format!("Failed to parse Gemini image response: {}", e)) })?; - // First, check Gemini models (candidates -> content -> parts -> inline_data) let image_data_from_gemini = gemini_response.candidates.as_ref().and_then(|candidates| { candidates.iter().find_map(|candidate| { candidate @@ -589,13 +196,11 @@ impl QueryBuilder for GoogleAIQueryBuilder { }) }); - // Then, check Imagen models (predictions -> bytes_base64_encoded) let image_data_from_imagen = gemini_response .predictions .as_ref() .and_then(|predictions| predictions.first().map(|p| &p.bytes_base64_encoded)); - // Image data, preferring Gemini first then Imagen models let image_data = image_data_from_gemini.or(image_data_from_imagen); match image_data { @@ -627,7 +232,6 @@ impl QueryBuilder for GoogleAIQueryBuilder { .. } = gemini_sse_parser; - // Send tool call arguments events for accumulated tool calls for tool_call in accumulated_tool_calls.values() { let event = StreamingEvent::ToolCallArguments { call_id: tool_call.id.clone(), @@ -637,7 +241,6 @@ impl QueryBuilder for GoogleAIQueryBuilder { stream_event_processor.send(event, &mut events_str).await?; } - // Convert Gemini usage metadata to TokenUsage let usage = gemini_usage.map(|u| { TokenUsage::new( u.prompt_token_count, @@ -647,11 +250,7 @@ impl QueryBuilder for GoogleAIQueryBuilder { }); Ok(ParsedResponse::Text { - content: if accumulated_content.is_empty() { - None - } else { - Some(accumulated_content) - }, + content: if accumulated_content.is_empty() { None } else { Some(accumulated_content) }, tool_calls: accumulated_tool_calls.into_values().collect(), events_str: Some(events_str), annotations, @@ -663,17 +262,11 @@ 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 - ) + format!("{}/models/{}:streamGenerateContent?alt=sse", base_url, model) } OutputType::Image => { - let url_suffix = if model.contains("imagen") { - "predict" - } else { - "generateContent" - }; + let url_suffix = + if model.contains("imagen") { "predict" } else { "generateContent" }; format!("{}/models/{}:{}", base_url, model, url_suffix) } } @@ -685,7 +278,6 @@ 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())] } } diff --git a/backend/windmill-worker/src/ai/sse.rs b/backend/windmill-worker/src/ai/sse.rs index 77fba3140a..62f13f3494 100644 --- a/backend/windmill-worker/src/ai/sse.rs +++ b/backend/windmill-worker/src/ai/sse.rs @@ -5,15 +5,18 @@ use reqwest::Response; use serde::Deserialize; use serde_json; use tokio_stream::StreamExt; -use windmill_common::{error::Error, utils::rd_string}; +use windmill_common::{ + ai_google::{parse_gemini_sse_event, GeminiUsageMetadata}, + ai_types::{ExtraContent, GoogleExtraContent, OpenAIFunction, OpenAIToolCall}, + error::Error, + utils::rd_string, +}; use crate::ai::{ query_builder::StreamEventProcessor, types::{StreamingEvent, UrlCitation}, }; -use windmill_common::ai_types::{ExtraContent, GoogleExtraContent, OpenAIFunction, OpenAIToolCall}; - #[derive(Deserialize)] pub struct OpenAIChoiceDeltaToolCallFunction { pub name: Option, @@ -457,96 +460,19 @@ impl SSEParser for AnthropicSSEParser { // Gemini SSE Parser // ============================================================================ -/// Gemini streaming response part - can be text or function call -#[derive(Deserialize, Debug)] -pub struct GeminiSSEPart { - #[serde(default)] - pub text: Option, - #[serde(rename = "functionCall")] - pub function_call: Option, - /// Thought signature for Gemini 3+ models - required for function calling - #[serde(rename = "thoughtSignature")] - pub thought_signature: Option, -} - -/// Function call in Gemini streaming response -#[derive(Deserialize, Debug)] -pub struct GeminiSSEFunctionCall { - pub name: String, - pub args: serde_json::Value, -} - -/// Content in Gemini streaming candidate -#[derive(Deserialize, Debug)] -pub struct GeminiSSEContent { - pub parts: Option>, -} - -/// Web reference in Gemini grounding chunk -#[derive(Deserialize, Debug)] -pub struct GeminiGroundingChunkWeb { - pub uri: String, - #[serde(default)] - pub title: Option, -} - -/// Grounding chunk from Gemini web search -#[derive(Deserialize, Debug)] -pub struct GeminiGroundingChunk { - pub web: Option, -} - -/// Grounding metadata from Gemini web search -#[derive(Deserialize, Debug)] -pub struct GeminiGroundingMetadata { - #[serde(rename = "groundingChunks", default)] - pub grounding_chunks: Vec, - #[serde(rename = "webSearchQueries", default)] - pub web_search_queries: Vec, -} - -/// Candidate in Gemini streaming response -#[derive(Deserialize, Debug)] -pub struct GeminiSSECandidate { - pub content: Option, - #[serde(rename = "finishReason")] - #[allow(dead_code)] - pub finish_reason: Option, - #[serde(rename = "groundingMetadata")] - pub grounding_metadata: Option, -} - -/// Gemini usage metadata from SSE response -#[derive(Deserialize, Debug, Clone)] -pub struct GeminiUsageMetadata { - #[serde(rename = "promptTokenCount", default)] - pub prompt_token_count: Option, - #[serde(rename = "candidatesTokenCount", default)] - pub candidates_token_count: Option, - #[serde(rename = "totalTokenCount", default)] - pub total_token_count: Option, -} - -/// Gemini SSE event structure -#[derive(Deserialize, Debug)] -pub struct GeminiSSEEvent { - pub candidates: Option>, - #[serde(rename = "usageMetadata")] - pub usage_metadata: Option, -} - -/// Gemini SSE Parser for streaming responses +/// Accumulates Gemini streaming events and converts them into the worker's +/// internal [`OpenAIToolCall`] / [`StreamingEvent`] representation. +/// +/// The actual SSE parsing is delegated to [`parse_gemini_sse_event`] from +/// `windmill_common::ai_google` so the logic can be shared with the API proxy. pub struct GeminiSSEParser { pub accumulated_content: String, pub accumulated_tool_calls: HashMap, pub events_str: String, pub stream_event_processor: StreamEventProcessor, tool_call_index: i64, - /// Collected URL citation annotations from web search pub annotations: Vec, - /// Whether web search was used in this response pub used_websearch: bool, - /// Token usage from usageMetadata pub usage: Option, } @@ -567,101 +493,57 @@ impl GeminiSSEParser { impl SSEParser for GeminiSSEParser { async fn parse_event_data(&mut self, data: &str) -> Result<(), Error> { - let event: Option = serde_json::from_str(data) - .inspect_err(|e| { - tracing::error!("Failed to parse SSE as a Gemini event {}: {}", data, e); - }) - .ok(); + let Some(parsed) = parse_gemini_sse_event(data)? else { + return Ok(()); + }; - if let Some(event) = event { - if let Some(candidates) = event.candidates { - for candidate in candidates { - if let Some(content) = candidate.content { - if let Some(parts) = content.parts { - for part in parts { - // Handle text content - if let Some(text) = part.text { - if !text.is_empty() { - self.accumulated_content.push_str(&text); - let event = StreamingEvent::TokenDelta { content: text }; - self.stream_event_processor - .send(event, &mut self.events_str) - .await?; - } - } + if let Some(text) = parsed.text { + self.accumulated_content.push_str(&text); + self.stream_event_processor + .send(StreamingEvent::TokenDelta { content: text }, &mut self.events_str) + .await?; + } - // Handle function calls - if let Some(function_call) = part.function_call { - let call_id = format!("call_{}", rd_string(24)); - let idx = self.tool_call_index; - self.tool_call_index += 1; + for tool_call in parsed.tool_calls { + let call_id = format!("call_{}", rd_string(24)); + let idx = self.tool_call_index; + self.tool_call_index += 1; - // Send tool call start event - let event = StreamingEvent::ToolCall { - call_id: call_id.clone(), - function_name: function_call.name.clone(), - }; - self.stream_event_processor - .send(event, &mut self.events_str) - .await?; + self.stream_event_processor + .send( + StreamingEvent::ToolCall { + call_id: call_id.clone(), + function_name: tool_call.name.clone(), + }, + &mut self.events_str, + ) + .await?; - // Build extra_content with thought_signature if present - let extra_content = - part.thought_signature.map(|sig| ExtraContent { - google: Some(GoogleExtraContent { - thought_signature: Some(sig), - }), - }); + let extra_content = tool_call.thought_signature.map(|sig| ExtraContent { + google: Some(GoogleExtraContent { thought_signature: Some(sig) }), + }); - // Store accumulated tool call - self.accumulated_tool_calls.insert( - idx, - OpenAIToolCall { - id: call_id, - function: OpenAIFunction { - name: function_call.name, - arguments: serde_json::to_string( - &function_call.args, - ) - .unwrap_or_else(|_| "{}".to_string()), - }, - r#type: "function".to_string(), - extra_content, - }, - ); - } - } - } - } + self.accumulated_tool_calls.insert( + idx, + OpenAIToolCall { + id: call_id, + function: OpenAIFunction { + name: tool_call.name, + arguments: serde_json::to_string(&tool_call.args) + .unwrap_or_else(|_| "{}".to_string()), + }, + r#type: "function".to_string(), + extra_content, + }, + ); + } - // Handle grounding metadata (web search results) - if let Some(ref grounding_metadata) = candidate.grounding_metadata { - // Set used_websearch if there are search queries or grounding chunks - if !grounding_metadata.web_search_queries.is_empty() - || !grounding_metadata.grounding_chunks.is_empty() - { - self.used_websearch = true; - } - - // Extract citations from grounding chunks - for chunk in &grounding_metadata.grounding_chunks { - if let Some(ref web) = chunk.web { - self.annotations.push(UrlCitation { - start_index: 0, // Gemini doesn't provide character indices - end_index: 0, - url: web.uri.clone(), - title: web.title.clone(), - }); - } - } - } - } - } - - // Extract usage metadata - if let Some(usage_metadata) = event.usage_metadata { - self.usage = Some(usage_metadata); - } + self.annotations.extend(parsed.annotations); + if parsed.used_websearch { + self.used_websearch = true; + } + if let Some(usage) = parsed.usage { + self.usage = Some(usage); } Ok(()) diff --git a/backend/windmill-worker/src/ai/utils.rs b/backend/windmill-worker/src/ai/utils.rs index a374db29c5..d095602b40 100644 --- a/backend/windmill-worker/src/ai/utils.rs +++ b/backend/windmill-worker/src/ai/utils.rs @@ -731,16 +731,3 @@ pub fn extract_text_content(content: &OpenAIContent) -> String { .join(""), } } - -/// Parse a data URL to extract media type and base64 data -/// Format: data:mime_type;base64,data -/// Returns (media_type, data) tuple if successful -pub fn parse_data_url(url: &str) -> Option<(String, String)> { - if !url.starts_with("data:") { - return None; - } - let rest = url.strip_prefix("data:")?; - let (header, data) = rest.split_once(",")?; - let media_type = header.strip_suffix(";base64")?; - Some((media_type.to_string(), data.to_string())) -}