From 3199f9fffd36c84d7cbb4a512935f7eb19fa5049 Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Wed, 17 Sep 2025 18:55:54 +0200 Subject: [PATCH] feat(ai agent): allow multiple images input for ai agent + code cleaning (#6591) * add in frontend * draft openai handling * upload to s3 * simpler output * return s3 directly if any * low quality * implement for gemini * handle imagen model * handle image input * cleaning * remove base64 from output * cleaning * fix timeout * handle openrouter * remove log * allow image input when creating image * cleaning * increase stack size * inline everything * revert stack size * refactor: move AI executor types to separate module - Created ai module structure with types.rs - Moved all type definitions from ai_executor.rs to ai/types.rs - No functional changes, just code organization * refactor: add QueryBuilder trait and provider detection utilities - Created QueryBuilder trait for abstracting provider-specific logic - Added helper functions for provider detection (is_anthropic_provider) - Implemented placeholder QueryBuilder for all providers - Updated OpenAIRequest to use slices instead of Vec references - All providers now have QueryBuilder implementations (using default for now) * feat: implement OpenAI query builder with image support foundation - Created proper OpenAI query builder implementation - Added image_handler module for S3 upload/download utilities - Separated text and image request building logic - Added prepare_messages_for_api to handle S3Object conversion - Foundation laid for supporting tools with image output * refactor(ai): complete AI executor refactoring with query builder pattern - Created modular structure under ai/ module - Moved all types to ai/types.rs - Created QueryBuilder trait for provider abstraction - Implemented OpenAI query builder with image+tools support - Added unified agent runner supporting both text and image outputs with tools - Refactored run_agent to delegate to new unified implementation - Added image handler utilities for S3 operations - Improved code organization and maintainability * cleaning * feat(ai): implement remaining provider query builders - Added Anthropic query builder with proper message conversion - Added Google AI query builder with Gemini API support - Added OpenRouter query builder delegating to OpenAI for compatibility - Added missing Anthropic and Gemini types to types.rs - Fixed type references and compilation errors - All providers now support the unified query builder interface * fixes * fixes * mime type + cleaning * image to images * handle mutlitple images * fix * remove agent_runner file * clean query builder logic * cleaning * cleaning * hide structured_output based on output type * fix * user images and not nested * better descriptions --- backend/Cargo.lock | 1 + backend/windmill-worker/Cargo.toml | 1 + .../windmill-worker/src/ai/image_handler.rs | 68 + backend/windmill-worker/src/ai/mod.rs | 7 + .../src/ai/providers/google_ai.rs | 275 ++ .../windmill-worker/src/ai/providers/mod.rs | 3 + .../src/ai/providers/openai.rs | 349 +++ .../src/ai/providers/openrouter.rs | 202 ++ .../windmill-worker/src/ai/query_builder.rs | 70 + backend/windmill-worker/src/ai/types.rs | 368 +++ backend/windmill-worker/src/ai_executor.rs | 2276 +++++------------ backend/windmill-worker/src/lib.rs | 1 + .../lib/components/InputTransformForm.svelte | 58 +- .../InputTransformSchemaForm.svelte | 3 + .../src/lib/components/flows/flowInfers.ts | 48 +- 15 files changed, 2076 insertions(+), 1654 deletions(-) create mode 100644 backend/windmill-worker/src/ai/image_handler.rs create mode 100644 backend/windmill-worker/src/ai/mod.rs create mode 100644 backend/windmill-worker/src/ai/providers/google_ai.rs create mode 100644 backend/windmill-worker/src/ai/providers/mod.rs create mode 100644 backend/windmill-worker/src/ai/providers/openai.rs create mode 100644 backend/windmill-worker/src/ai/providers/openrouter.rs create mode 100644 backend/windmill-worker/src/ai/query_builder.rs create mode 100644 backend/windmill-worker/src/ai/types.rs diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 005bec2f26..0a427653f9 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -15744,6 +15744,7 @@ dependencies = [ "async-once-cell", "async-recursion", "async-stream", + "async-trait", "backon", "base64 0.22.1", "bit-vec 0.6.3", diff --git a/backend/windmill-worker/Cargo.toml b/backend/windmill-worker/Cargo.toml index de7261a44e..19c8c0d31f 100644 --- a/backend/windmill-worker/Cargo.toml +++ b/backend/windmill-worker/Cargo.toml @@ -67,6 +67,7 @@ serde.workspace = true serde_json.workspace = true futures.workspace = true async-recursion.workspace = true +async-trait.workspace = true anyhow.workspace = true itertools.workspace = true regex.workspace = true diff --git a/backend/windmill-worker/src/ai/image_handler.rs b/backend/windmill-worker/src/ai/image_handler.rs new file mode 100644 index 0000000000..c668ded17f --- /dev/null +++ b/backend/windmill-worker/src/ai/image_handler.rs @@ -0,0 +1,68 @@ +use base64::Engine; +use futures; +use ulid; +use windmill_common::{client::AuthedClient, error::Error, s3_helpers::S3Object}; +use windmill_queue::MiniPulledJob; + +/// Upload image to S3 and return S3Object +pub async fn upload_image_to_s3( + base64_image: &str, + job: &MiniPulledJob, + client: &AuthedClient, +) -> Result { + let image_bytes = base64::engine::general_purpose::STANDARD + .decode(base64_image) + .map_err(|e| Error::internal_err(format!("Failed to decode base64 image: {}", e)))?; + + // Generate unique S3 key + let unique_id = ulid::Ulid::new().to_string(); + let s3_key = format!("ai_images/{}/{}.png", job.id, unique_id); + + // Create byte stream + let byte_stream = futures::stream::once(async move { + Ok::<_, std::convert::Infallible>(bytes::Bytes::from(image_bytes)) + }); + + // Upload to S3 + client + .upload_s3_file( + &job.workspace_id, + s3_key.clone(), + None, // storage - use default + byte_stream, + ) + .await + .map_err(|e| Error::internal_err(format!("Failed to upload image to S3: {}", e)))?; + + Ok(S3Object { + s3: s3_key, + storage: None, + filename: Some("generated_image.png".to_string()), + presigned: None, + }) +} + +/// Download an S3 image and convert it to a base64 data URL +pub async fn download_and_encode_s3_image( + image: &S3Object, + client: &AuthedClient, + workspace_id: &str, +) -> Result<(String, String), Error> { + // Download the image from S3 + let image_bytes = client + .download_s3_file(workspace_id, &image.s3, image.storage.clone()) + .await + .map_err(|e| Error::internal_err(format!("Failed to download S3 image: {}", e)))?; + + // Encode as base64 data URL + let base64_data = base64::engine::general_purpose::STANDARD.encode(&image_bytes); + + // Determine MIME type using mime_guess from file extension, with PNG as fallback + let mime_type = mime_guess::from_path(&image.s3).first(); + let mime_type = mime_type + .as_ref() + .map(|mime| mime.essence_str()) + .unwrap_or("image/png"); + + Ok((mime_type.to_string(), base64_data)) +} diff --git a/backend/windmill-worker/src/ai/mod.rs b/backend/windmill-worker/src/ai/mod.rs new file mode 100644 index 0000000000..b381e51693 --- /dev/null +++ b/backend/windmill-worker/src/ai/mod.rs @@ -0,0 +1,7 @@ +// AI executor module structure +// This module will contain all AI-related execution logic + +pub mod image_handler; +pub mod providers; +pub mod query_builder; +pub mod types; diff --git a/backend/windmill-worker/src/ai/providers/google_ai.rs b/backend/windmill-worker/src/ai/providers/google_ai.rs new file mode 100644 index 0000000000..9b08d40e46 --- /dev/null +++ b/backend/windmill-worker/src/ai/providers/google_ai.rs @@ -0,0 +1,275 @@ +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use serde_json; +use windmill_common::{client::AuthedClient, error::Error}; + +use crate::ai::{ + image_handler::download_and_encode_s3_image, + query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder}, + types::*, +}; + +// Google AI/Gemini-specific types +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct GeminiInlineData { + #[serde(rename = "mimeType")] + pub mime_type: String, + pub data: String, +} + +#[derive(Serialize, Deserialize, Clone)] +#[serde(untagged)] +pub enum GeminiPart { + Text { text: String }, + InlineData { inline_data: GeminiInlineData }, + FunctionCall { function_call: GeminiFunctionCall }, + FunctionResponse { function_response: GeminiFunctionResponse }, +} + +#[derive(Serialize, Deserialize, Clone)] +pub struct GeminiFunctionCall { + pub name: String, + pub args: serde_json::Value, +} + +#[derive(Serialize, Deserialize, Clone)] +pub struct GeminiFunctionResponse { + pub name: String, + pub response: serde_json::Value, +} + +#[derive(Serialize)] +pub struct GeminiContent { + pub parts: Vec, +} + +#[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>, +} + +#[derive(Serialize)] +pub struct GeminiPredictContent { + pub prompt: String, +} + +#[derive(Deserialize)] +pub struct GeminiImageResponse { + #[serde(skip_serializing_if = "Option::is_none")] + pub candidates: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub predictions: Option>, +} + +#[derive(Deserialize)] +pub struct GeminiCandidate { + pub content: GeminiResponseContent, +} + +#[derive(Deserialize)] +pub struct GeminiPredictCandidate { + #[serde(rename = "bytesBase64Encoded")] + pub bytes_base64_encoded: String, // base64 encoded image +} + +#[derive(Deserialize)] +pub struct GeminiResponseContent { + pub parts: Vec, +} + +#[derive(Deserialize)] +pub struct GeminiResponsePart { + #[serde(skip_serializing_if = "Option::is_none")] + #[allow(dead_code)] + pub text: Option, + #[serde(rename = "inlineData", skip_serializing_if = "Option::is_none")] + pub inline_data: Option, + #[serde(rename = "functionCall", skip_serializing_if = "Option::is_none")] + #[allow(dead_code)] + pub function_call: Option, +} + +pub struct GoogleAIQueryBuilder; + +impl GoogleAIQueryBuilder { + pub fn new() -> Self { + Self + } + + async fn build_image_request( + &self, + args: &BuildRequestArgs<'_>, + client: &AuthedClient, + workspace_id: &str, + ) -> Result { + 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(), + }]), + 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 { + parts.insert( + 0, + GeminiPart::Text { text: format!("SYSTEM PROMPT: {}", system_prompt.trim()) }, + ); + } + + // Add input images if provided + if let Some(images) = args.images { + for image in images.iter() { + if !image.s3.is_empty() { + let (mime_type, image_bytes) = + download_and_encode_s3_image(image, client, workspace_id).await?; + parts.push(GeminiPart::InlineData { + inline_data: GeminiInlineData { + mime_type: mime_type, + data: image_bytes, + }, + }); + } + } + } + + GeminiImageRequest { instances: None, contents: Some(vec![GeminiContent { parts }]) } + }; + + serde_json::to_string(&request) + .map_err(|e| Error::internal_err(format!("Failed to serialize request: {}", e))) + } +} + +#[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) + } + + async fn build_request( + &self, + args: &BuildRequestArgs<'_>, + client: &AuthedClient, + workspace_id: &str, + ) -> Result { + match args.output_type { + OutputType::Text => { + // For text output, use OpenAI-compatible format + let openai_builder = super::openai::OpenAIQueryBuilder::new(); + openai_builder + .build_request(args, client, workspace_id) + .await + } + OutputType::Image => self.build_image_request(args, client, workspace_id).await, + } + } + + async fn parse_response(&self, response: reqwest::Response) -> Result { + let url = response.url().path(); + + // For chat completions (text), use OpenAI parser + if url.contains("/chat/completions") { + let openai_builder = super::openai::OpenAIQueryBuilder::new(); + return openai_builder.parse_response(response).await; + } + + // Check if this is an image generation response + if url.contains(":predict") || url.contains(":generateContent") { + let response_text = response + .text() + .await + .map_err(|e| Error::internal_err(format!("Failed to read response text: {}", e)))?; + + let gemini_response: GeminiImageResponse = serde_json::from_str(&response_text) + .map_err(|e| { + Error::internal_err(format!( + "Failed to parse Gemini response: {}. Raw response: {}", + e, response_text + )) + })?; + + // Find image data in response + let image_data = gemini_response + .candidates + .as_ref() + .and_then(|candidates| { + candidates.iter().find_map(|candidate| { + candidate + .content + .parts + .iter() + .find_map(|part| part.inline_data.as_ref().map(|data| &data.data)) + }) + }) + .or_else(|| { + gemini_response + .predictions + .as_ref() + .and_then(|predictions| { + predictions + .iter() + .find_map(|prediction| Some(&prediction.bytes_base64_encoded)) + }) + }); + + if let Some(base64_image) = image_data { + Ok(ParsedResponse::Image { base64_data: base64_image.clone() }) + } else { + Err(Error::internal_err( + "No image data received from Gemini".to_string(), + )) + } + } else { + // This should not happen as we use OpenAI format for text + Err(Error::internal_err( + "Unexpected text response in Google AI parser".to_string(), + )) + } + } + + fn get_endpoint(&self, base_url: &str, model: &str, output_type: &OutputType) -> String { + match output_type { + OutputType::Text => format!("{}/chat/completions", base_url), // Use OpenAI-compatible endpoint + OutputType::Image => { + // For image generation, build the full URL with model name + let url_suffix = if model.contains("imagen") { + "predict" + } else { + "generateContent" + }; + format!( + "https://generativelanguage.googleapis.com/v1beta/models/{}:{}", + model, url_suffix + ) + } + } + } + + fn get_auth_headers( + &self, + api_key: &str, + output_type: &OutputType, + ) -> Vec<(&'static str, String)> { + match output_type { + OutputType::Text => { + // For text output, use Bearer token (OpenAI-compatible) + vec![("Authorization", format!("Bearer {}", api_key))] + } + OutputType::Image => { + // For image generation, use Google API key header + vec![("x-goog-api-key", api_key.to_string())] + } + } + } +} diff --git a/backend/windmill-worker/src/ai/providers/mod.rs b/backend/windmill-worker/src/ai/providers/mod.rs new file mode 100644 index 0000000000..13cf766e28 --- /dev/null +++ b/backend/windmill-worker/src/ai/providers/mod.rs @@ -0,0 +1,3 @@ +pub mod google_ai; +pub mod openai; +pub mod openrouter; diff --git a/backend/windmill-worker/src/ai/providers/openai.rs b/backend/windmill-worker/src/ai/providers/openai.rs new file mode 100644 index 0000000000..26bc804aed --- /dev/null +++ b/backend/windmill-worker/src/ai/providers/openai.rs @@ -0,0 +1,349 @@ +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use serde_json; +use windmill_common::{client::AuthedClient, error::Error}; + +use crate::ai::{ + image_handler::download_and_encode_s3_image, + query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder}, + types::*, +}; + +// OpenAI-specific types +#[derive(Deserialize, Serialize, Clone, Debug)] +pub struct OpenAIFunction { + pub name: String, + pub arguments: String, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub struct OpenAIToolCall { + pub id: String, + pub function: OpenAIFunction, + pub r#type: String, +} + +#[derive(Deserialize)] +pub struct OpenAIChoice { + pub message: OpenAIMessage, +} + +#[derive(Deserialize)] +pub struct OpenAIResponse { + pub choices: Vec, +} + +#[derive(Serialize)] +pub struct ImageGenerationTool { + pub r#type: String, + pub quality: Option, + pub background: Option, +} + +// Input content for image generation - supports both text and images +#[derive(Serialize, Clone, Debug)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ImageGenerationContent { + #[serde(rename = "input_text")] + InputText { text: String }, + #[serde(rename = "input_image")] + InputImage { image_url: String }, +} + +#[derive(Serialize)] +pub struct ImageGenerationMessage { + pub role: String, + pub content: Vec, +} + +#[derive(Serialize)] +pub struct ImageGenerationRequest<'a> { + pub model: &'a str, + pub input: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub instructions: Option<&'a str>, + pub tools: Vec, +} + +#[derive(Deserialize)] +pub struct OpenAIImageResponse { + pub output: Vec, +} + +#[derive(Deserialize)] +pub struct OpenAIImageOutput { + pub r#type: String, // Expected to be "image_generation_call" + pub status: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub result: Option, // Base64 encoded image, None if not completed +} + +#[derive(Serialize)] +pub struct OpenAIRequest<'a> { + pub model: &'a str, + pub messages: &'a [OpenAIMessage], + #[serde(skip_serializing_if = "Option::is_none")] + pub tools: Option<&'a [ToolDef]>, + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub max_completion_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub response_format: Option, +} + +pub struct OpenAIQueryBuilder; + +impl OpenAIQueryBuilder { + pub fn new() -> Self { + Self + } + + pub async fn prepare_messages_for_api( + &self, + messages: &[OpenAIMessage], + client: &AuthedClient, + workspace_id: &str, + ) -> Result, Error> { + let mut prepared_messages = Vec::new(); + + for message in messages { + let mut prepared_message = message.clone(); + + if let Some(content) = &message.content { + match content { + OpenAIContent::Text(text) => { + prepared_message.content = Some(OpenAIContent::Text(text.clone())); + } + OpenAIContent::Parts(parts) => { + let mut prepared_content = Vec::new(); + + for part in parts { + match part { + ContentPart::S3Object { s3_object } => { + // Convert S3Object to base64 image URL + let (mime_type, image_bytes) = download_and_encode_s3_image( + s3_object, + client, + workspace_id, + ) + .await?; + prepared_content.push(ContentPart::ImageUrl { + image_url: ImageUrlData { + url: format!( + "data:{};base64,{}", + mime_type, image_bytes + ), + }, + }); + } + other => { + // Keep Text and ImageUrl as-is + prepared_content.push(other.clone()); + } + } + } + + prepared_message.content = Some(OpenAIContent::Parts(prepared_content)); + } + } + } + + prepared_messages.push(prepared_message); + } + + Ok(prepared_messages) + } + + async fn build_text_request( + &self, + args: &BuildRequestArgs<'_>, + client: &AuthedClient, + workspace_id: &str, + ) -> Result { + let prepared_messages = self + .prepare_messages_for_api(args.messages, client, workspace_id) + .await?; + + // Check if we need to add response_format for structured output + let has_output_properties = args + .output_schema + .and_then(|schema| schema.properties.as_ref()) + .map(|props| !props.is_empty()) + .unwrap_or(false); + + let response_format = if has_output_properties && args.output_schema.is_some() { + let schema = args.output_schema.unwrap(); + let strict_schema = schema.clone().make_strict(); + Some(ResponseFormat { + r#type: "json_schema".to_string(), + json_schema: JsonSchemaFormat { + name: "structured_output".to_string(), + schema: strict_schema, + strict: Some(true), + }, + }) + } else { + None + }; + + let request = OpenAIRequest { + model: args.model, + messages: &prepared_messages, + tools: args.tools, + temperature: args.temperature, + max_completion_tokens: args.max_tokens, + response_format, + }; + + serde_json::to_string(&request) + .map_err(|e| Error::internal_err(format!("Failed to serialize request: {}", e))) + } + + async fn build_image_request( + &self, + args: &BuildRequestArgs<'_>, + client: &AuthedClient, + workspace_id: &str, + ) -> Result { + // Build content array with text and optional image + let mut content = + vec![ImageGenerationContent::InputText { text: args.user_message.to_string() }]; + + // Add images if provided + if let Some(images) = args.images { + for image in images.iter() { + if !image.s3.is_empty() { + let (mime_type, image_bytes) = + download_and_encode_s3_image(image, client, workspace_id).await?; + content.push(ImageGenerationContent::InputImage { + image_url: format!("data:{};base64,{}", mime_type, image_bytes), + }); + } + } + } + + // Build the request with tools if provided + let tools = vec![ImageGenerationTool { + r#type: "image_generation".to_string(), + quality: Some("low".to_string()), + background: None, + }]; + + // TODO: OpenAI's image generation API doesn't support custom tools in the same way as chat completions + // This would require a different approach, potentially using chat completions with image output + // For now, we'll use the standard image generation without custom tools + + let image_request = ImageGenerationRequest { + model: args.model, + input: vec![ImageGenerationMessage { role: "user".to_string(), content }], + instructions: args.system_prompt, + tools, + }; + + serde_json::to_string(&image_request) + .map_err(|e| Error::internal_err(format!("Failed to serialize request: {}", e))) + } +} + +#[async_trait] +impl QueryBuilder for OpenAIQueryBuilder { + fn supports_tools_with_output_type(&self, _output_type: &OutputType) -> bool { + // OpenAI supports tools for both text and image output + true + } + + async fn build_request( + &self, + args: &BuildRequestArgs<'_>, + client: &AuthedClient, + workspace_id: &str, + ) -> Result { + match args.output_type { + OutputType::Text => self.build_text_request(args, client, workspace_id).await, + OutputType::Image => self.build_image_request(args, client, workspace_id).await, + } + } + + async fn parse_response(&self, response: reqwest::Response) -> Result { + // Check if this is an image response + let url = response.url().path(); + if url.contains("/responses") { + // Parse image generation response + let response_text = response + .text() + .await + .map_err(|e| Error::internal_err(format!("Failed to read response text: {}", e)))?; + + let image_response: OpenAIImageResponse = serde_json::from_str(&response_text) + .map_err(|e| { + Error::internal_err(format!( + "Failed to parse OpenAI image response: {}. Raw response: {}", + e, response_text + )) + })?; + + // Find the first completed image generation output + let image_generation_call = image_response + .output + .iter() + .find(|output| { + output.r#type == "image_generation_call" && output.status == "completed" + }) + .and_then(|output| output.result.as_ref()); + + if let Some(base64_image) = image_generation_call { + Ok(ParsedResponse::Image { base64_data: base64_image.clone() }) + } else { + Err(Error::internal_err( + "No completed image output received from OpenAI".to_string(), + )) + } + } else { + // Parse text/chat completion response + let openai_response: OpenAIResponse = response + .json() + .await + .map_err(|e| Error::internal_err(format!("Failed to parse response: {}", e)))?; + + let first_choice = openai_response + .choices + .into_iter() + .next() + .ok_or_else(|| Error::internal_err("No response from API"))?; + + Ok(ParsedResponse::Text { + content: first_choice.message.content.map(|c| match c { + OpenAIContent::Text(text) => text, + OpenAIContent::Parts(parts) => { + // Extract text from parts + parts + .into_iter() + .filter_map(|part| match part { + ContentPart::Text { text } => Some(text), + _ => None, + }) + .collect::>() + .join(" ") + } + }), + tool_calls: first_choice.message.tool_calls.unwrap_or_default(), + }) + } + } + + fn get_endpoint(&self, base_url: &str, _model: &str, output_type: &OutputType) -> String { + match output_type { + OutputType::Text => format!("{}/chat/completions", base_url), + OutputType::Image => format!("{}/responses", base_url), + } + } + + fn get_auth_headers( + &self, + api_key: &str, + _output_type: &OutputType, + ) -> Vec<(&'static str, String)> { + vec![("Authorization", format!("Bearer {}", api_key))] + } +} diff --git a/backend/windmill-worker/src/ai/providers/openrouter.rs b/backend/windmill-worker/src/ai/providers/openrouter.rs new file mode 100644 index 0000000000..9fc6c6f051 --- /dev/null +++ b/backend/windmill-worker/src/ai/providers/openrouter.rs @@ -0,0 +1,202 @@ +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use serde_json; +use windmill_common::{client::AuthedClient, error::Error}; + +use crate::ai::{ + providers::openai::{OpenAIQueryBuilder, OpenAIResponse}, + query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder}, + types::*, +}; + +// OpenRouter-specific types +#[derive(Serialize)] +pub struct OpenRouterChatRequest<'a> { + pub model: &'a str, + pub messages: &'a [OpenAIMessage], + #[serde(skip_serializing_if = "Option::is_none")] + pub tools: Option<&'a [ToolDef]>, + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub max_completion_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub response_format: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub modalities: Option>, +} + +#[derive(Deserialize)] +pub struct OpenRouterImageResponse { + pub choices: Vec, +} + +#[derive(Deserialize)] +pub struct OpenRouterImageChoice { + pub message: OpenRouterImageResponseMessage, +} + +#[derive(Deserialize)] +pub struct OpenRouterImageResponseMessage { + #[serde(skip_serializing_if = "Option::is_none")] + pub images: Option>, +} + +#[derive(Deserialize)] +pub struct OpenRouterImageData { + pub image_url: OpenRouterImageUrl, +} + +#[derive(Deserialize)] +pub struct OpenRouterImageUrl { + pub url: String, // data:image/png;base64,... format +} + +pub struct OpenRouterQueryBuilder { + // OpenRouter uses OpenAI-compatible API, so we delegate most work to OpenAI builder + openai_builder: OpenAIQueryBuilder, +} + +impl OpenRouterQueryBuilder { + pub fn new() -> Self { + Self { openai_builder: OpenAIQueryBuilder::new() } + } +} + +#[async_trait] +impl QueryBuilder for OpenRouterQueryBuilder { + fn supports_tools_with_output_type(&self, _output_type: &OutputType) -> bool { + // OpenRouter supports tools for both text and image output (via OpenAI-compatible API) + true + } + + async fn build_request( + &self, + args: &BuildRequestArgs<'_>, + client: &AuthedClient, + workspace_id: &str, + ) -> Result { + match args.output_type { + OutputType::Text => { + // For text, use standard OpenAI format without modalities + self.openai_builder + .build_request(args, client, workspace_id) + .await + } + OutputType::Image => { + // For image generation, we need to add modalities field + // First, prepare the messages using the OpenAI builder's logic + let openai_builder = &self.openai_builder; + let prepared_messages = openai_builder + .prepare_messages_for_api(args.messages, client, workspace_id) + .await?; + + // Check if we need to add response_format for structured output + let has_output_properties = args + .output_schema + .and_then(|schema| schema.properties.as_ref()) + .map(|props| !props.is_empty()) + .unwrap_or(false); + + let response_format = if has_output_properties && args.output_schema.is_some() { + let schema = args.output_schema.unwrap(); + let strict_schema = schema.clone().make_strict(); + Some(ResponseFormat { + r#type: "json_schema".to_string(), + json_schema: JsonSchemaFormat { + name: "structured_output".to_string(), + schema: strict_schema, + strict: Some(true), + }, + }) + } else { + None + }; + + // Build OpenRouter-specific request with modalities + let request = OpenRouterChatRequest { + model: args.model, + messages: &prepared_messages, + tools: args.tools, + temperature: args.temperature, + max_completion_tokens: args.max_tokens, + response_format, + modalities: Some(vec!["image", "text"]), + }; + + serde_json::to_string(&request) + .map_err(|e| Error::internal_err(format!("Failed to serialize request: {}", e))) + } + } + } + + async fn parse_response(&self, response: reqwest::Response) -> Result { + let response_text = response + .text() + .await + .map_err(|e| Error::internal_err(format!("Failed to read response text: {}", e)))?; + + // First try to parse as OpenRouter image response + if let Ok(image_response) = serde_json::from_str::(&response_text) + { + // Extract base64 image from the first choice + let image_url = image_response + .choices + .get(0) + .and_then(|choice| choice.message.images.as_ref()) + .and_then(|images| images.get(0)) + .map(|image| &image.image_url.url); + + if let Some(data_url) = image_url { + // Extract base64 data from data URL format: data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAA... + if let Some(base64_start) = data_url.find("base64,") { + let base64_data = &data_url[base64_start + 7..]; // Skip "base64," prefix + return Ok(ParsedResponse::Image { base64_data: base64_data.to_string() }); + } + } + } + + // If not an image response or parsing failed, try as regular OpenAI response + let openai_response: OpenAIResponse = + serde_json::from_str(&response_text).map_err(|e| { + Error::internal_err(format!( + "Failed to parse response: {}. Raw response: {}", + e, response_text + )) + })?; + + let first_choice = openai_response + .choices + .into_iter() + .next() + .ok_or_else(|| Error::internal_err("No response from API"))?; + + Ok(ParsedResponse::Text { + content: first_choice.message.content.map(|c| match c { + OpenAIContent::Text(text) => text, + OpenAIContent::Parts(parts) => parts + .into_iter() + .filter_map(|part| match part { + ContentPart::Text { text } => Some(text), + _ => None, + }) + .collect::>() + .join(" "), + }), + tool_calls: first_choice.message.tool_calls.unwrap_or_default(), + }) + } + + fn get_endpoint(&self, base_url: &str, _model: &str, _output_type: &OutputType) -> String { + // OpenRouter uses the same endpoint for both text and image generation + format!("{}/chat/completions", base_url) + } + + fn get_auth_headers( + &self, + api_key: &str, + _output_type: &OutputType, + ) -> Vec<(&'static str, String)> { + vec![("Authorization", format!("Bearer {}", api_key))] + } +} diff --git a/backend/windmill-worker/src/ai/query_builder.rs b/backend/windmill-worker/src/ai/query_builder.rs new file mode 100644 index 0000000000..07b9c96e4e --- /dev/null +++ b/backend/windmill-worker/src/ai/query_builder.rs @@ -0,0 +1,70 @@ +use async_trait::async_trait; +use windmill_common::{client::AuthedClient, error::Error, s3_helpers::S3Object}; + +use crate::ai::{ + providers::{ + google_ai::GoogleAIQueryBuilder, + openai::{OpenAIQueryBuilder, OpenAIToolCall}, + openrouter::OpenRouterQueryBuilder, + }, + types::*, +}; + +/// Arguments for building an AI request +pub struct BuildRequestArgs<'a> { + pub messages: &'a [OpenAIMessage], + pub tools: Option<&'a [ToolDef]>, + pub model: &'a str, + pub temperature: Option, + pub max_tokens: Option, + pub output_schema: Option<&'a OpenAPISchema>, + pub output_type: &'a OutputType, + pub system_prompt: Option<&'a str>, + pub user_message: &'a str, + pub images: Option<&'a [S3Object]>, +} + +/// Response from AI provider +pub enum ParsedResponse { + Text { content: Option, tool_calls: Vec }, + Image { base64_data: String }, +} + +/// Trait for building provider-specific AI requests +#[async_trait] +pub trait QueryBuilder: Send + Sync { + /// Check if this provider supports tools with the given output type + fn supports_tools_with_output_type(&self, output_type: &OutputType) -> bool; + + /// Build the request body for the provider + async fn build_request( + &self, + args: &BuildRequestArgs<'_>, + client: &AuthedClient, + workspace_id: &str, + ) -> Result; + + /// Parse the response from the provider + async fn parse_response(&self, response: reqwest::Response) -> Result; + + /// Get the API endpoint for this provider + fn get_endpoint(&self, base_url: &str, model: &str, output_type: &OutputType) -> String; + + /// Get the authentication headers for this provider + fn get_auth_headers( + &self, + api_key: &str, + output_type: &OutputType, + ) -> Vec<(&'static str, String)>; +} + +/// Factory function to create the appropriate query builder for a provider +pub fn create_query_builder(provider: &ProviderWithResource) -> Box { + use windmill_common::ai_providers::AIProvider; + + match provider.kind { + AIProvider::GoogleAI => Box::new(GoogleAIQueryBuilder::new()), + AIProvider::OpenRouter => Box::new(OpenRouterQueryBuilder::new()), + _ => Box::new(OpenAIQueryBuilder::new()), // Use OpenAI as default for all other providers + } +} diff --git a/backend/windmill-worker/src/ai/types.rs b/backend/windmill-worker/src/ai/types.rs new file mode 100644 index 0000000000..cc1417a92d --- /dev/null +++ b/backend/windmill-worker/src/ai/types.rs @@ -0,0 +1,368 @@ +use serde::{Deserialize, Serialize}; +use serde_json::value::RawValue; +use std::collections::HashMap; +use windmill_common::{ + ai_providers::AIProvider, db::DB, error::Error, flow_status::AgentAction, flows::FlowModule, + s3_helpers::S3Object, +}; +use windmill_parser::Typ; + +use crate::ai::providers::openai::OpenAIToolCall; + +// Shared types used across multiple providers + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ContentPart { + Text { + text: String, + }, + #[serde(rename = "image_url")] + ImageUrl { + image_url: ImageUrlData, + }, + #[serde(rename = "s3_object")] + S3Object { + s3_object: S3Object, + }, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct ImageUrlData { + pub url: String, // data:image/png;base64,... or https://... +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(untagged)] +pub enum OpenAIContent { + Text(String), + Parts(Vec), +} + +#[derive(Deserialize, Serialize, Clone, Default, Debug)] +pub struct OpenAIMessage { + pub role: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_call_id: Option, + #[serde(skip_serializing)] + pub agent_action: Option, +} + +/// same as OpenAIMessage but with agent_action field included in the serialization +#[derive(Serialize)] +pub struct Message<'a> { + #[serde(flatten)] + pub message: &'a OpenAIMessage, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_action: Option<&'a AgentAction>, +} + +#[derive(Serialize, Clone, Debug)] +pub struct ResponseFormat { + pub r#type: String, + pub json_schema: JsonSchemaFormat, +} + +#[derive(Serialize, Clone, Debug)] +pub struct JsonSchemaFormat { + pub name: String, + pub schema: OpenAPISchema, + #[serde(skip_serializing_if = "Option::is_none")] + pub strict: Option, +} + +#[derive(Serialize, Clone, Debug)] +pub struct ToolDefFunction { + pub name: String, + pub description: Option, + pub parameters: Box, +} + +#[derive(Serialize, Clone, Debug)] +pub struct ToolDef { + pub r#type: String, + pub function: ToolDefFunction, +} + +pub struct Tool { + pub module: FlowModule, + pub def: ToolDef, +} + +#[derive(Deserialize, Serialize, Debug, Clone, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum OutputType { + Text, + Image, +} + +impl Default for OutputType { + fn default() -> Self { + OutputType::Text + } +} + +#[derive(Deserialize, Debug)] +pub struct AIAgentArgs { + pub provider: ProviderWithResource, + pub system_prompt: Option, + pub user_message: String, + pub temperature: Option, + pub max_completion_tokens: Option, + pub output_schema: Option, + pub output_type: Option, + pub user_images: Option>, +} + +#[derive(Deserialize, Debug)] +pub struct ProviderResource { + #[serde(alias = "apiKey")] + pub api_key: String, + #[serde(alias = "baseUrl")] + pub base_url: Option, +} + +#[derive(Deserialize, Debug)] +pub struct ProviderWithResource { + pub kind: AIProvider, + pub resource: ProviderResource, + pub model: String, +} + +impl ProviderWithResource { + pub fn get_api_key(&self) -> &str { + &self.resource.api_key + } + + pub fn get_model(&self) -> &str { + &self.model + } + + pub async fn get_base_url(&self, db: &DB) -> Result { + self.kind + .get_base_url(self.resource.base_url.clone(), db) + .await + } +} + +#[derive(Serialize)] +pub struct AIAgentResult<'a> { + pub output: Box, + pub messages: Vec>, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(untagged)] +pub enum SchemaType { + Single(String), + Multiple(Vec), +} + +impl Default for SchemaType { + fn default() -> Self { + SchemaType::Single("object".to_string()) + } +} + +#[derive(Serialize, Deserialize, Default, Clone, Debug)] +pub struct OpenAPISchema { + #[serde(skip_serializing_if = "Option::is_none")] + pub r#type: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub items: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub properties: Option>>, + #[serde(skip_serializing_if = "Option::is_none")] + pub required: Option>, + #[serde(skip_serializing_if = "Option::is_none", rename = "oneOf")] + pub one_of: Option>>, + #[serde(skip_serializing_if = "Option::is_none")] + pub format: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub r#enum: Option>, + #[serde( + skip_serializing_if = "Option::is_none", + rename = "additionalProperties" + )] + pub additional_properties: Option, +} + +impl OpenAPISchema { + pub fn from_str(typ: &str) -> Self { + OpenAPISchema { r#type: Some(SchemaType::Single(typ.to_string())), ..Default::default() } + } + + pub fn from_str_with_enum(typ: &str, enu: &Option>) -> Self { + OpenAPISchema { + r#type: Some(SchemaType::Single(typ.to_string())), + r#enum: enu.clone(), + ..Default::default() + } + } + + pub fn datetime() -> Self { + Self { + r#type: Some(SchemaType::Single("string".to_string())), + format: Some("date-time".to_string()), + ..Default::default() + } + } + + pub fn from_typ(typ: &Typ) -> Self { + match typ { + Typ::Str(enu) => Self::from_str_with_enum("string", enu), + Typ::Int => Self::from_str("integer"), + Typ::Float => Self::from_str("number"), + Typ::Bool => Self::from_str("boolean"), + Typ::Bytes => Self::from_str("string"), + Typ::Datetime => Self::datetime(), + Typ::Resource(_) => Self::from_str("string"), + Typ::Email => Self::from_str("string"), + Typ::Sql => Self::from_str("string"), + Typ::DynSelect(_) => Self::from_str("string"), + Typ::DynMultiselect(_) => Self::from_str("string"), + Typ::List(typ) => OpenAPISchema { + r#type: Some(SchemaType::Single("array".to_string())), + items: Some(Box::new(Self::from_typ(typ))), + ..Default::default() + }, + Typ::Object(typ) => OpenAPISchema { + r#type: Some(SchemaType::Single("object".to_string())), + items: None, + properties: typ.props.as_ref().map(|props| { + props + .iter() + .map(|prop| (prop.key.clone(), Box::new(Self::from_typ(&prop.typ)))) + .collect() + }), + required: typ + .props + .as_ref() + .map(|props| props.iter().map(|prop| prop.key.clone()).collect()), + ..Default::default() + }, + Typ::OneOf(variants) => OpenAPISchema { + r#type: Some(SchemaType::Single("object".to_string())), + one_of: Some( + variants + .iter() + .map(|variant| { + let schema = OpenAPISchema { + r#type: Some(SchemaType::Single("object".to_string())), + properties: Some( + variant + .properties + .iter() + .map(|prop| { + ( + prop.key.clone(), + Box::new( + if prop.key == "label" || prop.key == "kind" { + Self::from_str_with_enum( + "string", + &Some(vec![variant.label.clone()]), + ) + } else { + Self::from_typ(&prop.typ) + }, + ), + ) + }) + .collect(), + ), + required: Some( + variant + .properties + .iter() + .map(|prop| prop.key.clone()) + .collect(), + ), + ..Default::default() + }; + Box::new(schema) + }) + .collect(), + ), + ..Default::default() + }, + Typ::Unknown => Self::from_str("object"), + } + } + + /// Makes this schema compatible with OpenAI's strict mode by: + /// - Adding additionalProperties: false to all object types + /// - Making non-required properties nullable + /// - Ensuring all properties are in the required array + pub fn make_strict(mut self) -> Self { + // Handle this schema if it's an object type + if let Some(SchemaType::Single(ref type_str)) = self.r#type { + if type_str == "object" { + // Set additionalProperties to false + self.additional_properties = Some(false); + + if let Some(properties) = self.properties.as_mut() { + // Get original required fields + let original_required = self.required.as_ref(); + + if let Some(required) = original_required { + // Update properties to make non-required fields nullable + for (key, prop) in properties.iter_mut() { + let mut new_prop = (**prop).clone(); + // Make non-required fields nullable + if !required.contains(key) { + new_prop = new_prop.make_nullable(); + } + // Recursively make nested schemas strict + new_prop = new_prop.make_strict(); + *prop = Box::new(new_prop); + } + } + + // All properties must be in required array for strict mode + self.required = Some(properties.keys().cloned().collect()); + } + } + } + + // Recursively process nested schemas + if let Some(ref mut items) = self.items { + **items = items.as_ref().clone().make_strict(); + } + + if let Some(ref mut one_of) = self.one_of { + *one_of = one_of + .iter() + .map(|schema| Box::new(schema.as_ref().clone().make_strict())) + .collect(); + } + + self + } + + /// Makes this property nullable by converting its type to a union with null + pub fn make_nullable(mut self) -> Self { + match self.r#type.take() { + Some(SchemaType::Single(type_str)) => { + if type_str != "null" { + self.r#type = Some(SchemaType::Multiple(vec![type_str, "null".into()])); + } else { + self.r#type = Some(SchemaType::Single("null".into())); + } + } + Some(SchemaType::Multiple(mut types)) => { + if !types.iter().any(|t| t == "null") { + types.push("null".into()); + } + self.r#type = Some(SchemaType::Multiple(types)); + } + None => { + self.r#type = Some(SchemaType::Single("null".into())); + } + } + self + } +} diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 7f5df68bc8..ddc2cc480f 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -1,35 +1,34 @@ use async_recursion::async_recursion; -use base64::Engine; -use mime_guess; use regex::Regex; -use serde::{Deserialize, Serialize}; use serde_json::value::RawValue; use std::{collections::HashMap, sync::Arc}; -#[cfg(feature = "benchmark")] -use windmill_common::bench::BenchmarkIter; +use ulid; +use uuid::Uuid; use windmill_common::{ ai_providers::AIProvider, - auth::get_job_perms, cache, client::AuthedClient, db::DB, error::{self, to_anyhow, Error}, flow_status::AgentAction, - flows::{FlowModule, FlowModuleValue, Step}, + flows::{FlowModuleValue, Step}, get_latest_hash_for_path, jobs::JobKind, - s3_helpers::S3Object, scripts::{get_full_hub_script_by_path, ScriptHash, ScriptLang}, utils::{StripPath, HTTP_CLIENT}, worker::{to_raw_value, Connection}, }; -use windmill_parser::Typ; use windmill_queue::{ flow_status::get_step_of_flow_status, get_mini_pulled_job, push, CanceledBy, JobCompleted, MiniPulledJob, PushArgs, PushIsolationLevel, }; use crate::{ + ai::{ + image_handler::upload_image_to_s3, + query_builder::{create_query_builder, BuildRequestArgs, ParsedResponse}, + types::*, + }, common::{build_args_map, error_to_value, OccupancyMetrics}, create_job_dir, handle_child::run_future_with_polling_update_job_poller, @@ -39,1093 +38,12 @@ use crate::{ JobCompletedSender, SendResult, SendResultPayload, }; -const MAX_AGENT_ITERATIONS: usize = 10; -const REQUEST_TIMEOUT: u64 = 120; - lazy_static::lazy_static! { static ref TOOL_NAME_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_]+$").unwrap(); } -#[derive(Deserialize, Serialize, Clone, Debug)] -struct OpenAIFunction { - name: String, - arguments: String, -} - -#[derive(Deserialize, Serialize, Clone, Debug)] -struct OpenAIToolCall { - id: String, - function: OpenAIFunction, - r#type: String, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(tag = "type", rename_all = "snake_case")] -enum ContentPart { - Text { - text: String, - }, - #[serde(rename = "image_url")] - ImageUrl { - image_url: ImageUrlData, - }, - #[serde(rename = "s3_object")] - S3Object { - s3_object: S3Object, - }, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -struct ImageUrlData { - url: String, // data:image/png;base64,... or https://... -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(untagged)] -enum OpenAIContent { - Text(String), - Parts(Vec), -} - -#[derive(Deserialize, Serialize, Clone, Default, Debug)] -struct OpenAIMessage { - role: String, - #[serde(skip_serializing_if = "Option::is_none")] - content: Option, - #[serde(skip_serializing_if = "Option::is_none")] - tool_calls: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - tool_call_id: Option, - #[serde(skip_serializing)] - agent_action: Option, -} - -/// same as OpenAIMessage but with agent_action field included in the serialization -#[derive(Serialize)] -struct Message<'a> { - #[serde(flatten)] - message: &'a OpenAIMessage, - #[serde(skip_serializing_if = "Option::is_none")] - agent_action: Option<&'a AgentAction>, -} - -#[derive(Deserialize)] -struct OpenAIChoice { - message: OpenAIMessage, -} - -#[derive(Deserialize)] -struct OpenAIResponse { - choices: Vec, -} - -#[derive(Serialize)] -struct ImageGenerationTool { - r#type: String, - quality: Option, - background: Option, -} - -// Input content for image generation - supports both text and images -#[derive(Serialize, Clone, Debug)] -#[serde(tag = "type", rename_all = "snake_case")] -enum ImageGenerationContent { - #[serde(rename = "input_text")] - InputText { text: String }, - #[serde(rename = "input_image")] - InputImage { image_url: String }, -} - -#[derive(Serialize)] -struct ImageGenerationMessage { - role: String, - content: Vec, -} - -#[derive(Serialize)] -struct ImageGenerationRequest<'a> { - model: &'a str, - input: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - instructions: Option<&'a str>, - tools: Vec, -} - -#[derive(Deserialize)] -struct OpenAIImageResponse { - output: Vec, -} - -#[derive(Deserialize)] -struct OpenAIImageOutput { - r#type: String, // Expected to be "image_generation_call" - #[serde(default)] - result: Option, // Base64 encoded image -} - -// Gemini API structures -#[derive(Serialize, Deserialize, Clone, Debug)] -struct GeminiInlineData { - #[serde(rename = "mimeType")] - mime_type: String, - data: String, -} - -#[derive(Serialize)] -#[serde(untagged)] -enum GeminiPart { - Text { text: String }, - InlineData { inline_data: GeminiInlineData }, -} - -#[derive(Serialize)] -struct GeminiContent { - parts: Vec, -} - -#[derive(Serialize)] -struct GeminiImageRequest { - #[serde(skip_serializing_if = "Option::is_none")] - contents: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - instances: Option>, -} - -#[derive(Serialize)] -struct GeminiPredictContent { - prompt: String, -} - -#[derive(Deserialize)] -struct GeminiImageResponse { - #[serde(skip_serializing_if = "Option::is_none")] - candidates: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - predictions: Option>, -} - -#[derive(Deserialize)] -struct GeminiCandidate { - content: GeminiResponseContent, -} - -#[derive(Deserialize)] -struct GeminiPredictCandidate { - #[serde(rename = "bytesBase64Encoded")] - bytes_base64_encoded: String, // base64 encoded image -} - -#[derive(Deserialize)] -struct GeminiResponseContent { - parts: Vec, -} - -#[derive(Deserialize)] -struct GeminiResponsePart { - #[serde(rename = "inlineData")] - inline_data: Option, -} - -// OpenRouter image generation structures -#[derive(Serialize)] -struct OpenRouterImageRequest<'a> { - model: &'a str, - messages: Vec, - modalities: Vec<&'a str>, -} - -#[derive(Serialize)] -struct OpenRouterImageMessage { - role: String, - content: String, -} - -#[derive(Deserialize)] -struct OpenRouterImageResponse { - choices: Vec, -} - -#[derive(Deserialize)] -struct OpenRouterImageChoice { - message: OpenRouterImageResponseMessage, -} - -#[derive(Deserialize)] -struct OpenRouterImageResponseMessage { - #[serde(skip_serializing_if = "Option::is_none")] - images: Option>, -} - -#[derive(Deserialize)] -struct OpenRouterImageData { - image_url: OpenRouterImageUrl, -} - -#[derive(Deserialize)] -struct OpenRouterImageUrl { - url: String, // data:image/png;base64,... format -} - -#[derive(Serialize)] -struct OpenAIRequest<'a> { - model: &'a str, - messages: &'a Vec, - #[serde(skip_serializing_if = "Option::is_none")] - tools: Option<&'a Vec>, - #[serde(skip_serializing_if = "Option::is_none")] - temperature: Option, - #[serde(skip_serializing_if = "Option::is_none")] - max_completion_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - response_format: Option, -} - -#[derive(Serialize, Clone, Debug)] -struct ResponseFormat { - r#type: String, - json_schema: JsonSchemaFormat, -} - -#[derive(Serialize, Clone, Debug)] -struct JsonSchemaFormat { - name: String, - schema: OpenAPISchema, - #[serde(skip_serializing_if = "Option::is_none")] - strict: Option, -} - -#[derive(Serialize, Clone, Debug)] -struct ToolDefFunction { - name: String, - description: Option, - parameters: Box, -} - -#[derive(Serialize, Clone, Debug)] -struct ToolDef { - r#type: String, - function: ToolDefFunction, -} - -struct Tool { - module: FlowModule, - def: ToolDef, -} - -#[derive(Deserialize, Serialize, Debug, Clone, PartialEq)] -#[serde(rename_all = "lowercase")] -enum OutputType { - Text, - Image, -} - -impl Default for OutputType { - fn default() -> Self { - OutputType::Text - } -} - -#[derive(Deserialize, Debug)] -struct AIAgentArgs { - provider: ProviderWithResource, - system_prompt: Option, - user_message: String, - temperature: Option, - max_completion_tokens: Option, - output_schema: Option, - output_type: Option, - image: Option, -} - -#[derive(Deserialize, Debug)] -struct ProviderResource { - #[serde(alias = "apiKey")] - api_key: String, - #[serde(alias = "baseUrl")] - base_url: Option, -} - -#[derive(Deserialize, Debug)] -struct ProviderWithResource { - kind: AIProvider, - resource: ProviderResource, - model: String, -} - -impl ProviderWithResource { - fn get_api_key(&self) -> &str { - &self.resource.api_key - } - - fn get_model(&self) -> &str { - &self.model - } - - async fn get_base_url(&self, db: &DB) -> Result { - self.kind - .get_base_url(self.resource.base_url.clone(), db) - .await - } -} - -#[derive(Serialize)] -struct AIAgentResult<'a> { - output: Box, - messages: Vec>, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(untagged)] -enum SchemaType { - Single(String), - Multiple(Vec), -} - -impl Default for SchemaType { - fn default() -> Self { - SchemaType::Single("object".to_string()) - } -} - -#[derive(Serialize, Deserialize, Default, Clone, Debug)] -struct OpenAPISchema { - #[serde(skip_serializing_if = "Option::is_none")] - r#type: Option, - #[serde(skip_serializing_if = "Option::is_none")] - items: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - properties: Option>>, - #[serde(skip_serializing_if = "Option::is_none")] - required: Option>, - #[serde(skip_serializing_if = "Option::is_none", rename = "oneOf")] - one_of: Option>>, - #[serde(skip_serializing_if = "Option::is_none")] - format: Option, - #[serde(skip_serializing_if = "Option::is_none")] - r#enum: Option>, - #[serde( - skip_serializing_if = "Option::is_none", - rename = "additionalProperties" - )] - additional_properties: Option, -} - -impl OpenAPISchema { - fn from_str(typ: &str) -> Self { - OpenAPISchema { r#type: Some(SchemaType::Single(typ.to_string())), ..Default::default() } - } - - fn from_str_with_enum(typ: &str, enu: &Option>) -> Self { - OpenAPISchema { - r#type: Some(SchemaType::Single(typ.to_string())), - r#enum: enu.clone(), - ..Default::default() - } - } - - fn datetime() -> Self { - Self { - r#type: Some(SchemaType::Single("string".to_string())), - format: Some("date-time".to_string()), - ..Default::default() - } - } - - fn from_typ(typ: &Typ) -> Self { - match typ { - Typ::Str(enu) => Self::from_str_with_enum("string", enu), - Typ::Int => Self::from_str("integer"), - Typ::Float => Self::from_str("number"), - Typ::Bool => Self::from_str("boolean"), - Typ::Bytes => Self::from_str("string"), - Typ::Datetime => Self::datetime(), - Typ::Resource(_) => Self::from_str("string"), - Typ::Email => Self::from_str("string"), - Typ::Sql => Self::from_str("string"), - Typ::DynSelect(_) => Self::from_str("string"), - Typ::DynMultiselect(_) => Self::from_str("string"), - Typ::List(typ) => OpenAPISchema { - r#type: Some(SchemaType::Single("array".to_string())), - items: Some(Box::new(Self::from_typ(typ))), - ..Default::default() - }, - Typ::Object(typ) => OpenAPISchema { - r#type: Some(SchemaType::Single("object".to_string())), - items: None, - properties: typ.props.as_ref().map(|props| { - props - .iter() - .map(|prop| (prop.key.clone(), Box::new(Self::from_typ(&prop.typ)))) - .collect() - }), - required: typ - .props - .as_ref() - .map(|props| props.iter().map(|prop| prop.key.clone()).collect()), - ..Default::default() - }, - Typ::OneOf(variants) => OpenAPISchema { - r#type: Some(SchemaType::Single("object".to_string())), - one_of: Some( - variants - .iter() - .map(|variant| { - let schema = OpenAPISchema { - r#type: Some(SchemaType::Single("object".to_string())), - properties: Some( - variant - .properties - .iter() - .map(|prop| { - ( - prop.key.clone(), - Box::new( - if prop.key == "label" || prop.key == "kind" { - Self::from_str_with_enum( - "string", - &Some(vec![variant.label.clone()]), - ) - } else { - Self::from_typ(&prop.typ) - }, - ), - ) - }) - .collect(), - ), - required: Some( - variant - .properties - .iter() - .map(|prop| prop.key.clone()) - .collect(), - ), - ..Default::default() - }; - Box::new(schema) - }) - .collect(), - ), - ..Default::default() - }, - Typ::Unknown => Self::from_str("object"), - } - } - - /// Makes this schema compatible with OpenAI's strict mode by: - /// - Adding additionalProperties: false to all object types - /// - Making non-required properties nullable - /// - Ensuring all properties are in the required array - fn make_strict(mut self) -> Self { - // Handle this schema if it's an object type - if let Some(SchemaType::Single(ref type_str)) = self.r#type { - if type_str == "object" { - // Set additionalProperties to false - self.additional_properties = Some(false); - - if let Some(properties) = self.properties.as_mut() { - // Get original required fields - let original_required = self.required.as_ref(); - - if let Some(required) = original_required { - // Update properties to make non-required fields nullable - for (key, prop) in properties.iter_mut() { - let mut new_prop = (**prop).clone(); - // Make non-required fields nullable - if !required.contains(key) { - new_prop = new_prop.make_nullable(); - } - // Recursively make nested schemas strict - new_prop = new_prop.make_strict(); - *prop = Box::new(new_prop); - } - } - - // All properties must be in required array for strict mode - self.required = Some(properties.keys().cloned().collect()); - } - } - } - - // Recursively process nested schemas - if let Some(ref mut items) = self.items { - **items = items.as_ref().clone().make_strict(); - } - - if let Some(ref mut one_of) = self.one_of { - *one_of = one_of - .iter() - .map(|schema| Box::new(schema.as_ref().clone().make_strict())) - .collect(); - } - - self - } - - /// Makes this property nullable by converting its type to a union with null - fn make_nullable(mut self) -> Self { - match self.r#type.take() { - Some(SchemaType::Single(type_str)) => { - if type_str != "null" { - self.r#type = Some(SchemaType::Multiple(vec![type_str, "null".into()])); - } else { - self.r#type = Some(SchemaType::Single("null".into())); - } - } - Some(SchemaType::Multiple(mut types)) => { - if !types.iter().any(|t| t == "null") { - types.push("null".into()); - } - self.r#type = Some(SchemaType::Multiple(types)); - } - None => { - self.r#type = Some(SchemaType::Single("null".into())); - } - } - self - } -} - -/// Find a unique tool name to avoid collisions with user-provided tools -fn find_unique_tool_name(base_name: &str, existing_tools: Option<&[ToolDef]>) -> String { - let Some(tools) = existing_tools else { - return base_name.to_string(); - }; - - if !tools.iter().any(|t| t.function.name == base_name) { - return base_name.to_string(); - } - - for i in 1..100 { - let candidate = format!("{}_{}", base_name, i); - if !tools.iter().any(|t| t.function.name == candidate) { - return candidate; - } - } - - // Fallback with process id if somehow we can't find a unique name - format!("{}_{}_fallback", base_name, std::process::id()) -} - -/// Helper function to download an S3 image and convert it to a base64 data URL -async fn download_and_encode_s3_image( - image: &S3Object, - client: &AuthedClient, - workspace_id: &str, -) -> error::Result<(String, String)> { - // Download the image from S3 - let image_bytes = client - .download_s3_file(workspace_id, &image.s3, image.storage.clone()) - .await - .map_err(|e| Error::internal_err(format!("Failed to download S3 image: {}", e)))?; - - // Encode as base64 data URL - let base64_data = base64::engine::general_purpose::STANDARD.encode(&image_bytes); - - // Determine MIME type using mime_guess from file extension, with PNG as fallback - let mime_type = mime_guess::from_path(&image.s3).first(); - let mime_type = mime_type - .as_ref() - .map(|mime| mime.essence_str()) - .unwrap_or("image/png"); - - Ok((mime_type.to_string(), base64_data)) -} - -/// Convert messages with S3Objects to messages with base64 image URLs for API calls -async fn prepare_messages_for_api( - messages: &[OpenAIMessage], - client: &AuthedClient, - workspace_id: &str, -) -> error::Result> { - let mut prepared_messages = Vec::new(); - - for message in messages { - let mut prepared_message = message.clone(); - - if let Some(content) = &message.content { - match content { - OpenAIContent::Text(text) => { - prepared_message.content = Some(OpenAIContent::Text(text.clone())); - } - OpenAIContent::Parts(parts) => { - let mut prepared_content = Vec::new(); - - for part in parts { - match part { - ContentPart::S3Object { s3_object } => { - // Convert S3Object to base64 image URL - let (mime_type, image_data_url) = - download_and_encode_s3_image(s3_object, client, workspace_id) - .await?; - prepared_content.push(ContentPart::ImageUrl { - image_url: ImageUrlData { - url: format!( - "data:{};base64,{}", - mime_type, image_data_url - ), - }, - }); - } - other => { - // Keep Text and ImageUrl as-is - prepared_content.push(other.clone()); - } - } - } - - prepared_message.content = Some(OpenAIContent::Parts(prepared_content)); - } - } - } - - prepared_messages.push(prepared_message); - } - - Ok(prepared_messages) -} - -/// Generate image from provider and extract base64 data -async fn generate_image_from_provider( - provider: &ProviderWithResource, - user_message: &str, - system_prompt: Option<&str>, - base_url: &str, - api_key: &str, - image: Option<&S3Object>, - client: &AuthedClient, - workspace_id: &str, -) -> error::Result { - match provider.kind { - AIProvider::OpenAI => { - // Build content array with text and optional image - let mut content = - vec![ImageGenerationContent::InputText { text: user_message.to_string() }]; - - // Add image if provided - if let Some(image) = image { - if !image.s3.is_empty() { - // Download and encode S3 image to base64 - let (mime_type, bytes64) = - download_and_encode_s3_image(image, client, workspace_id).await?; - content.push(ImageGenerationContent::InputImage { - image_url: format!("data:{};base64,{}", mime_type, bytes64), - }); - } - } - - let image_request = ImageGenerationRequest { - model: provider.get_model(), - input: vec![ImageGenerationMessage { role: "user".to_string(), content }], - instructions: system_prompt, - tools: vec![ImageGenerationTool { - r#type: "image_generation".to_string(), - quality: Some("low".to_string()), - background: None, - }], - }; - - let resp = HTTP_CLIENT - .post(format!("{}/responses", base_url)) - .timeout(std::time::Duration::from_secs(REQUEST_TIMEOUT)) - .bearer_auth(api_key) - .json(&image_request) - .send() - .await - .map_err(|e| Error::internal_err(format!("Failed to call OpenAI API: {}", e)))?; - - match resp.error_for_status_ref() { - Ok(_) => { - let image_response = resp.json::().await.map_err(|e| { - Error::internal_err(format!("Failed to parse OpenAI response: {}", e)) - })?; - - // Find the first image generation output - let image_generation_call = image_response - .output - .iter() - .find(|output| output.r#type == "image_generation_call") - .and_then(|output| output.result.as_ref()); - - if let Some(base64_image) = image_generation_call { - Ok(base64_image.to_string()) - } else { - Err(Error::internal_err( - "No image output received from OpenAI".to_string(), - )) - } - } - Err(e) => { - let _status = resp.status(); - let text = resp - .text() - .await - .unwrap_or_else(|_| "".to_string()); - Err(Error::internal_err(format!( - "OpenAI API error: {} - {}", - e, text - ))) - } - } - } - AIProvider::GoogleAI => { - let is_imagen = provider.get_model().contains("imagen"); - - let gemini_request = if is_imagen { - // For Imagen models, we keep the simple prompt format (no image support) - GeminiImageRequest { - instances: Some(vec![GeminiPredictContent { - prompt: user_message.trim().to_string(), - }]), - contents: None, - } - } else { - // For Gemini models, build parts array with text and optional image - let mut parts = vec![GeminiPart::Text { text: user_message.trim().to_string() }]; - - if let Some(system_prompt) = system_prompt { - parts.insert( - 0, - GeminiPart::Text { - text: format!("SYSTEM PROMPT: {}", system_prompt.trim().to_string()), - }, - ); - } - - // Add image if provided - if let Some(image) = image { - if !image.s3.is_empty() { - // Download and encode S3 image to base64 - let (mime_type, bytes64) = - download_and_encode_s3_image(image, client, workspace_id).await?; - - parts.push(GeminiPart::InlineData { - inline_data: GeminiInlineData { mime_type, data: bytes64 }, - }); - } - } - - GeminiImageRequest { - instances: None, - contents: Some(vec![GeminiContent { parts }]), - } - }; - - let url_suffix = if is_imagen { - "predict" - } else { - "generateContent" - }; - let gemini_url = format!( - "https://generativelanguage.googleapis.com/v1beta/models/{}:{}", - provider.get_model(), - url_suffix - ); - - let resp = HTTP_CLIENT - .post(&gemini_url) - .timeout(std::time::Duration::from_secs(REQUEST_TIMEOUT)) - .header("x-goog-api-key", api_key) - .header("Content-Type", "application/json") - .json(&gemini_request) - .send() - .await - .map_err(|e| Error::internal_err(format!("Failed to call Gemini API: {}", e)))?; - - match resp.error_for_status_ref() { - Ok(_) => { - let response_text = resp.text().await.map_err(|e| { - Error::internal_err(format!("Failed to read response text: {}", e)) - })?; - - let gemini_response: GeminiImageResponse = serde_json::from_str(&response_text) - .map_err(|e| { - Error::internal_err(format!( - "Failed to parse Gemini response: {}. Raw response: {}", - e, response_text - )) - })?; - - // Find the first candidate with inline image data - let mut image_data = - gemini_response.candidates.as_ref().and_then(|candidates| { - candidates.iter().find_map(|candidate| { - candidate.content.parts.iter().find_map(|part| { - part.inline_data.as_ref().map(|data| &data.data) - }) - }) - }); - - if image_data.is_none() { - image_data = gemini_response - .predictions - .as_ref() - .and_then(|predictions| { - predictions - .iter() - .find_map(|prediction| Some(&prediction.bytes_base64_encoded)) - }); - } - - if let Some(base64_image) = image_data { - Ok(base64_image.clone()) - } else { - Err(Error::internal_err( - "No image data received from Gemini".to_string(), - )) - } - } - Err(e) => { - let _status = resp.status(); - let text = resp - .text() - .await - .unwrap_or_else(|_| "".to_string()); - Err(Error::internal_err(format!( - "Gemini API error: {} - {}", - e, text - ))) - } - } - } - AIProvider::OpenRouter => { - let mut messages = Vec::new(); - - // Add system message if provided - if let Some(system_prompt) = system_prompt { - messages.push(OpenRouterImageMessage { - role: "system".to_string(), - content: system_prompt.to_string(), - }); - } - - // Add user message - messages.push(OpenRouterImageMessage { - role: "user".to_string(), - content: user_message.to_string(), - }); - - let openrouter_request = OpenRouterImageRequest { - model: provider.get_model(), - messages, - modalities: vec!["image", "text"], - }; - - let resp = HTTP_CLIENT - .post(format!("{}/chat/completions", base_url)) - .timeout(std::time::Duration::from_secs(REQUEST_TIMEOUT)) - .bearer_auth(api_key) - .json(&openrouter_request) - .send() - .await - .map_err(|e| { - Error::internal_err(format!("Failed to call OpenRouter API: {}", e)) - })?; - - match resp.error_for_status_ref() { - Ok(_) => { - let openrouter_response = - resp.json::().await.map_err(|e| { - Error::internal_err(format!( - "Failed to parse OpenRouter response: {}", - e - )) - })?; - - // Extract base64 image from the first choice - let image_url = openrouter_response - .choices - .get(0) - .and_then(|choice| choice.message.images.as_ref()) - .and_then(|images| images.get(0)) - .map(|image| &image.image_url.url); - - if let Some(data_url) = image_url { - // Extract base64 data from data URL format: data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAA... - if let Some(base64_start) = data_url.find("base64,") { - let base64_data = &data_url[base64_start + 7..]; // Skip "base64," prefix - Ok(base64_data.to_string()) - } else { - Err(Error::internal_err( - "Invalid data URL format received from OpenRouter".to_string(), - )) - } - } else { - Err(Error::internal_err( - "No image data received from OpenRouter".to_string(), - )) - } - } - Err(e) => { - let _status = resp.status(); - let text = resp - .text() - .await - .unwrap_or_else(|_| "".to_string()); - Err(Error::internal_err(format!( - "OpenRouter API error: {} - {}", - e, text - ))) - } - } - } - _ => Err(Error::BadRequest(format!( - "Image generation is not supported for provider: {:?}", - provider.kind - ))), - } -} - -/// Upload image to S3 and return S3Object -async fn upload_image_to_s3( - base64_image: &str, - job: &MiniPulledJob, - client: &AuthedClient, -) -> error::Result { - let image_bytes = base64::engine::general_purpose::STANDARD - .decode(base64_image) - .map_err(|e| Error::internal_err(format!("Failed to decode base64 image: {}", e)))?; - - // Generate unique S3 key - let unique_id = ulid::Ulid::new().to_string(); - let s3_key = format!("ai_images/{}/{}.png", job.id, unique_id); - - // Create byte stream - let byte_stream = futures::stream::once(async move { - Ok::<_, std::convert::Infallible>(bytes::Bytes::from(image_bytes)) - }); - - // Upload to S3 - client - .upload_s3_file( - &job.workspace_id, - s3_key.clone(), - None, // storage - use default - byte_stream, - ) - .await - .map_err(|e| Error::internal_err(format!("Failed to upload image to S3: {}", e)))?; - - Ok(S3Object { - s3: s3_key, - storage: None, - filename: Some("generated_image.png".to_string()), - presigned: None, - }) -} - -/// Handle image output generation and return S3 object and messages -async fn handle_image_output( - args: &AIAgentArgs, - job: &MiniPulledJob, - client: &AuthedClient, - db: &DB, -) -> error::Result<(Option, Vec)> { - let base_url = args.provider.get_base_url(db).await?; - let api_key = args.provider.get_api_key(); - - let mut messages = - if let Some(system_prompt) = args.system_prompt.clone().filter(|s| !s.is_empty()) { - vec![OpenAIMessage { - role: "system".to_string(), - content: Some(OpenAIContent::Text(system_prompt)), - ..Default::default() - }] - } else { - vec![] - }; - - // Generate image from provider - let base64_image = generate_image_from_provider( - &args.provider, - &args.user_message, - args.system_prompt.as_deref(), - &base_url, - api_key, - args.image.as_ref(), - client, - &job.workspace_id, - ) - .await?; - - // Add assistant success message - messages.push(OpenAIMessage { - role: "assistant".to_string(), - content: Some(OpenAIContent::Text( - "Image created successfully".to_string(), - )), - ..Default::default() - }); - - // Upload to S3 - let s3_object = upload_image_to_s3(&base64_image, job, client).await?; - - Ok((Some(s3_object), messages)) -} - -async fn update_flow_status_module_with_actions( - db: &DB, - parent_job: &uuid::Uuid, - actions: &[AgentAction], -) -> Result<(), Error> { - let step = get_step_of_flow_status(db, parent_job.to_owned()).await?; - match step { - Step::Step(step) => { - sqlx::query!( - r#" - UPDATE v2_job_status SET - flow_status = jsonb_set( - flow_status, - array['modules', $3::TEXT, 'agent_actions'], - $2 - ) - WHERE id = $1 - "#, - parent_job, - sqlx::types::Json(actions) as _, - step as i32 - ) - .execute(db) - .await?; - } - _ => {} - } - Ok(()) -} - -async fn update_flow_status_module_with_actions_success( - db: &DB, - parent_job: &uuid::Uuid, - action_success: bool, -) -> Result<(), Error> { - let step = get_step_of_flow_status(db, parent_job.to_owned()).await?; - match step { - Step::Step(step) => { - // Append the new bool to the existing array, or create a new array if it doesn't exist - sqlx::query!( - r#" - UPDATE v2_job_status SET - flow_status = jsonb_set( - flow_status, - array['modules', $2::TEXT, 'agent_actions_success'], - COALESCE( - flow_status->'modules'->$2->'agent_actions_success', - to_jsonb(ARRAY[]::bool[]) - ) || to_jsonb(ARRAY[$3::bool]) - ) - WHERE id = $1 - "#, - parent_job, - step as i32, - action_success - ) - .execute(db) - .await?; - } - _ => {} - } - Ok(()) -} +const MAX_AGENT_ITERATIONS: usize = 10; +const REQUEST_TIMEOUT_SECONDS: u64 = 120; fn parse_raw_script_schema(content: &str, language: &ScriptLang) -> Result, Error> { let main_arg_signature = parse_sig_of_lang(content, Some(&language), None)?.unwrap(); // safe to unwrap as langauge is some @@ -1156,546 +74,6 @@ fn parse_raw_script_schema(content: &str, language: &ScriptLang) -> Result, -) -> error::Result> { - let output_type = args.output_type.as_ref().unwrap_or(&OutputType::Text); - - match *output_type { - OutputType::Image => { - let (s3_result, messages) = handle_image_output(&args, job, client, db).await?; - - let final_messages: Vec = messages - .iter() - .map(|m| Message { message: m, agent_action: m.agent_action.as_ref() }) - .collect(); - - if let Some(s3_output) = s3_result { - Ok(to_raw_value(&s3_output)) - } else { - Ok(to_raw_value(&AIAgentResult { - output: to_raw_value(&None::), - messages: final_messages, - })) - } - } - OutputType::Text => { - let base_url = args.provider.get_base_url(db).await?; - let api_key = args.provider.get_api_key(); - - let mut messages = - if let Some(system_prompt) = args.system_prompt.clone().filter(|s| !s.is_empty()) { - vec![OpenAIMessage { - role: "system".to_string(), - content: Some(OpenAIContent::Text(system_prompt)), - ..Default::default() - }] - } else { - vec![] - }; - - // Create user message with optional image - let user_content = if let Some(image) = &args.image { - if !image.s3.is_empty() { - OpenAIContent::Parts(vec![ - ContentPart::Text { text: args.user_message.clone() }, - ContentPart::S3Object { s3_object: image.clone() }, - ]) - } else { - OpenAIContent::Text(args.user_message.clone()) - } - } else { - OpenAIContent::Text(args.user_message.clone()) - }; - - messages.push(OpenAIMessage { - role: "user".to_string(), - content: Some(user_content), - ..Default::default() - }); - - let mut actions = vec![]; - let mut content = None; - - let mut tool_defs: Option> = if tools.is_empty() { - None - } else { - Some(tools.iter().map(|t| t.def.clone()).collect()) - }; - - let has_output_properties = args - .output_schema - .as_ref() - .and_then(|schema| schema.properties.as_ref()) - .map(|props| !props.is_empty()) - .unwrap_or(false); - let provider_is_anthropic = args.provider.kind.is_anthropic(); - let is_openrouter_anthropic = args.provider.kind == AIProvider::OpenRouter - && args.provider.model.starts_with("anthropic/"); - let is_anthropic = provider_is_anthropic || is_openrouter_anthropic; - let mut response_format: Option = None; - let mut used_structured_output_tool = false; - let mut structured_output_tool_name: Option = None; - - if has_output_properties { - let schema = args.output_schema.as_ref().unwrap(); // we know it's some because of the check above - if is_anthropic { - // if output schema is provided, and provider is anthropic, add a structured_output tool in the list of tools - let unique_tool_name = - find_unique_tool_name("structured_output", tool_defs.as_deref()); - structured_output_tool_name = Some(unique_tool_name.clone()); - - let output_tool = ToolDef { - r#type: "function".to_string(), - function: ToolDefFunction { - name: unique_tool_name, - description: Some( - "This tool MUST be used last to return a structured JSON object as the final output." - .to_string(), - ), - parameters: to_raw_value(&schema), - }, - }; - if let Some(ref mut existing_tools) = tool_defs { - existing_tools.push(output_tool); - } else { - tool_defs = Some(vec![output_tool]); - } - } else { - // if output schema is provided, and provider is openai, add a response_format with json_schema - let strict_schema = schema.clone().make_strict(); - response_format = Some(ResponseFormat { - r#type: "json_schema".to_string(), - json_schema: JsonSchemaFormat { - name: "structured_output".to_string(), - schema: strict_schema, - strict: Some(true), - }, - }); - } - } - - for i in 0..MAX_AGENT_ITERATIONS { - if used_structured_output_tool { - break; - } - - let response = { - // Convert messages with S3Objects to base64 image URLs for API request - let prepared_messages = - prepare_messages_for_api(&messages, client, &job.workspace_id).await?; - - let resp = HTTP_CLIENT - .post(format!("{}/chat/completions", base_url)) - .timeout(std::time::Duration::from_secs(REQUEST_TIMEOUT)) - .bearer_auth(api_key) - .json(&OpenAIRequest { - model: args.provider.get_model(), - messages: &prepared_messages, - tools: tool_defs.as_ref(), - temperature: args.temperature, - max_completion_tokens: args.max_completion_tokens, - response_format: if has_output_properties && !is_anthropic { - response_format.clone() - } else { - None - }, - }) - .send() - .await - .map_err(|e| Error::internal_err(format!("Failed to call API: {}", e)))?; - - match resp.error_for_status_ref() { - Ok(_) => resp, - Err(e) => { - let status = resp.status(); - let text = resp - .text() - .await - .unwrap_or_else(|_| "".to_string()); - tracing::error!( - "Non 200 response from API: status: {}, body: {}", - status, - text - ); - return Err(Error::internal_err(format!( - "Non 200 response from API: {} - {}", - e, text - ))); - } - } - }; - - let mut response = response.json::().await.map_err(|e| { - Error::internal_err(format!("Failed to parse API response: {}", e)) - })?; - - let first_choice = response - .choices - .pop() - .ok_or_else(|| Error::internal_err("No response from API"))?; - - content = first_choice.message.content; - let tool_calls = first_choice.message.tool_calls.unwrap_or_default(); - - if let Some(ref response_content) = content { - actions.push(AgentAction::Message {}); - messages.push(OpenAIMessage { - role: "assistant".to_string(), - content: Some(response_content.clone()), - agent_action: Some(AgentAction::Message {}), - ..Default::default() - }); - - update_flow_status_module_with_actions(db, parent_job, &actions).await?; - update_flow_status_module_with_actions_success(db, parent_job, true).await?; - } - - if tool_calls.is_empty() { - break; - } else if i == MAX_AGENT_ITERATIONS - 1 { - return Err(Error::internal_err( - "AI agent reached max iterations, but there are still tool calls" - .to_string(), - )); - } - - messages.push(OpenAIMessage { - role: "assistant".to_string(), - tool_calls: Some(tool_calls.clone()), - ..Default::default() - }); - - for tool_call in tool_calls.iter() { - // Structured output tool is used, we stop here as this will be the final output - if structured_output_tool_name - .as_ref() - .map_or(false, |name| tool_call.function.name == *name) - { - used_structured_output_tool = true; - messages.push(OpenAIMessage { - role: "tool".to_string(), - content: Some(OpenAIContent::Text( - "Successfully ran structured_output tool".to_string(), - )), - tool_call_id: Some(tool_call.id.clone()), - ..Default::default() - }); - messages.push(OpenAIMessage { - role: "assistant".to_string(), - content: Some(OpenAIContent::Text( - tool_call.function.arguments.clone(), - )), - agent_action: Some(AgentAction::Message {}), - ..Default::default() - }); - content = Some(OpenAIContent::Text(tool_call.function.arguments.clone())); - break; - } - - let tool = tools - .iter() - .find(|t| t.def.function.name == tool_call.function.name); - if let Some(tool) = tool { - let job_id = ulid::Ulid::new().into(); - actions.push(AgentAction::ToolCall { - job_id, - function_name: tool_call.function.name.clone(), - module_id: tool.module.id.clone(), - }); - - update_flow_status_module_with_actions(db, parent_job, &actions).await?; - - let tool_call_args = serde_json::from_str::>>( - &tool_call.function.arguments, - )?; - - let job_payload = match tool.module.get_value()? { - FlowModuleValue::Script { - path: script_path, - hash: script_hash, - tag_override, - .. - } => { - let payload = script_to_payload( - script_hash, - script_path, - db, - job, - &tool.module, - tag_override, - tool.module.apply_preprocessor, - ) - .await?; - payload - } - FlowModuleValue::RawScript { - path, - content, - language, - lock, - tag, - custom_concurrency_key, - concurrent_limit, - concurrency_time_window_s, - .. - } => { - let path = path.unwrap_or_else(|| { - format!("{}/tools/{}", job.runnable_path(), tool.module.id) - }); - - let payload = raw_script_to_payload( - path, - content, - language, - lock, - custom_concurrency_key, - concurrent_limit, - concurrency_time_window_s, - &tool.module, - tag, - tool.module.delete_after_use.unwrap_or(false), - ); - payload - } - _ => { - return Err(Error::internal_err(format!( - "Unsupported tool: {}", - tool_call.function.name - ))); - } - }; - - let mut tx = db.begin().await?; - - let job_perms = get_job_perms(&mut *tx, &job.id, &job.workspace_id) - .await? - .map(|x| x.into()); - - let (email, permissioned_as) = - if let Some(on_behalf_of) = job_payload.on_behalf_of.as_ref() { - (&on_behalf_of.email, on_behalf_of.permissioned_as.clone()) - } else { - (&job.permissioned_as_email, job.permissioned_as.to_owned()) - }; - - let job_priority = tool.module.priority.or(job.priority); - - let tx = PushIsolationLevel::Transaction(tx); - let (uuid, tx) = push( - db, - tx, - &job.workspace_id, - job_payload.payload, - PushArgs { args: &tool_call_args, extra: None }, - &job.created_by, - email, - permissioned_as, - Some(&format!("job-span-{}", job.id)), - None, - job.schedule_path(), - Some(job.id), - None, - None, - Some(job_id), - false, - false, - None, - job.visible_to_owner, - Some(job.tag.clone()), // we reuse the same tag as the agent job because it's run on the same worker - job_payload.timeout, - None, - job_priority, - job_perms.as_ref(), - true, - ) - .await?; - - tx.commit().await?; - - let tool_job = get_mini_pulled_job(db, &uuid).await?; - - let Some(tool_job) = tool_job else { - return Err(Error::internal_err("Tool job not found".to_string())); - }; - - let tool_job = Arc::new(tool_job); - - let job_dir = create_job_dir(&worker_dir, job.id).await; - - let (inner_job_completed_tx, inner_job_completed_rx) = - JobCompletedSender::new(&conn, 1); - - let inner_job_completed_rx = inner_job_completed_rx.expect( - "inner_job_completed_tx should be set as agent jobs are not supported on agent workers", - ); - - #[cfg(feature = "benchmark")] - let mut bench = BenchmarkIter::new(); - - match handle_queued_job( - tool_job.clone(), - None, - None, - None, - None, - conn, - client, - hostname, - worker_name, - worker_dir, - &job_dir, - None, - base_internal_url, - inner_job_completed_tx, - occupancy_metrics, - killpill_rx, - None, - #[cfg(feature = "benchmark")] - &mut bench, - ) - .await - { - Err(err) => { - let err_string = format!("{}: {}", err.name(), err.to_string()); - let err_json = error_to_value(&err); - let _ = handle_non_flow_job_error( - db, - &tool_job, - 0, - None, - err_string.clone(), - err_json, - worker_name, - ) - .await; - messages.push(OpenAIMessage { - role: "tool".to_string(), - content: Some(OpenAIContent::Text(format!( - "Error running tool: {}", - err_string - ))), - tool_call_id: Some(tool_call.id.clone()), - agent_action: Some(AgentAction::ToolCall { - job_id, - function_name: tool_call.function.name.clone(), - module_id: tool.module.id.clone(), - }), - ..Default::default() - }); - update_flow_status_module_with_actions_success( - db, parent_job, false, - ) - .await?; - } - Ok(success) => { - let send_result = inner_job_completed_rx.bounded_rx.try_recv().ok(); - - let result = if let Some(SendResult { - result: - SendResultPayload::JobCompleted(JobCompleted { result, .. }), - .. - }) = send_result.as_ref() - { - job_completed_tx - .send(send_result.as_ref().unwrap().result.clone(), true) - .await - .map_err(to_anyhow)?; - result - } else { - if let Some(send_result) = send_result { - job_completed_tx - .send(send_result.result, true) - .await - .map_err(to_anyhow)?; - } - return Err(Error::internal_err( - "Tool job completed but no result".to_string(), - )); - }; - - messages.push(OpenAIMessage { - role: "tool".to_string(), - content: Some(OpenAIContent::Text(result.get().to_string())), - tool_call_id: Some(tool_call.id.clone()), - agent_action: Some(AgentAction::ToolCall { - job_id, - function_name: tool_call.function.name.clone(), - module_id: tool.module.id.clone(), - }), - ..Default::default() - }); - update_flow_status_module_with_actions_success( - db, parent_job, success, - ) - .await?; - } - } - } else { - return Err(Error::internal_err(format!( - "Tool not found: {}", - tool_call.function.name - ))); - } - } - } - - let final_messages: Vec = messages - .iter() - .map(|m| Message { message: m, agent_action: m.agent_action.as_ref() }) - .collect(); - - // Parse content as JSON, fallback to string if it fails - let output_value = match content { - Some(content_str) => match has_output_properties { - true => match content_str { - OpenAIContent::Text(text) => serde_json::from_str::>(&text) - .map_err(|_e| { - Error::internal_err(format!( - "Failed to parse structured output: {}", - text - )) - }), - // No need to handle this, it will always be a text string - OpenAIContent::Parts(_parts) => Err(Error::internal_err( - "Failed to parse structured output".to_string(), - )), - }, - false => Ok(match content_str { - OpenAIContent::Text(text) => to_raw_value(&text), - OpenAIContent::Parts(parts) => to_raw_value(&parts), - }), - }?, - None => to_raw_value(&""), - }; - - Ok(to_raw_value(&AIAgentResult { - output: output_value, - messages: final_messages, - })) - } - } -} - pub struct FlowJobRunnableIdAndRawFlow { pub runnable_id: Option, pub raw_flow: Option>>, @@ -1918,3 +296,637 @@ pub async fn handle_ai_agent_job( Ok(result) } + +/// Find a unique tool name to avoid collisions with user-provided tools +fn find_unique_tool_name(base_name: &str, existing_tools: Option<&[ToolDef]>) -> String { + let Some(tools) = existing_tools else { + return base_name.to_string(); + }; + + if !tools.iter().any(|t| t.function.name == base_name) { + return base_name.to_string(); + } + + for i in 1..100 { + let candidate = format!("{}_{}", base_name, i); + if !tools.iter().any(|t| t.function.name == candidate) { + return candidate; + } + } + + // Fallback with process id if somehow we can't find a unique name + format!("{}_{}_fallback", base_name, std::process::id()) +} + +async fn update_flow_status_module_with_actions( + db: &DB, + parent_job: &Uuid, + actions: &[AgentAction], +) -> Result<(), Error> { + let step = get_step_of_flow_status(db, parent_job.to_owned()).await?; + match step { + Step::Step(step) => { + sqlx::query!( + r#" + UPDATE v2_job_status SET + flow_status = jsonb_set( + flow_status, + array['modules', $3::TEXT, 'agent_actions'], + $2 + ) + WHERE id = $1 + "#, + parent_job, + sqlx::types::Json(actions) as _, + step as i32 + ) + .execute(db) + .await?; + } + _ => {} + } + Ok(()) +} + +async fn update_flow_status_module_with_actions_success( + db: &DB, + parent_job: &Uuid, + action_success: bool, +) -> Result<(), Error> { + let step = get_step_of_flow_status(db, parent_job.to_owned()).await?; + match step { + Step::Step(step) => { + // Append the new bool to the existing array, or create a new array if it doesn't exist + sqlx::query!( + r#" + UPDATE v2_job_status SET + flow_status = jsonb_set( + flow_status, + array['modules', $2::TEXT, 'agent_actions_success'], + COALESCE( + flow_status->'modules'->$2->'agent_actions_success', + to_jsonb(ARRAY[]::bool[]) + ) || to_jsonb(ARRAY[$3::bool]) + ) + WHERE id = $1 + "#, + parent_job, + step as i32, + action_success + ) + .execute(db) + .await?; + } + _ => {} + } + Ok(()) +} + +/// Check if the provider is Anthropic (either direct or through OpenRouter) +fn is_anthropic_provider(provider: &ProviderWithResource) -> bool { + let provider_is_anthropic = provider.kind.is_anthropic(); + let is_openrouter_anthropic = + provider.kind == AIProvider::OpenRouter && provider.model.starts_with("anthropic/"); + provider_is_anthropic || is_openrouter_anthropic +} + +#[async_recursion] +pub async fn run_agent( + // connection + db: &DB, + conn: &Connection, + + // agent job and flow data + job: &MiniPulledJob, + parent_job: &Uuid, + args: &AIAgentArgs, + tools: &[Tool], + + // job execution context + client: &AuthedClient, + occupancy_metrics: &mut OccupancyMetrics, + job_completed_tx: &JobCompletedSender, + worker_dir: &str, + base_internal_url: &str, + worker_name: &str, + hostname: &str, + killpill_rx: &mut tokio::sync::broadcast::Receiver<()>, +) -> error::Result> { + let output_type = args.output_type.as_ref().unwrap_or(&OutputType::Text); + let base_url = args.provider.get_base_url(db).await?; + let api_key = args.provider.get_api_key(); + + // Create the query builder for the provider + let query_builder = create_query_builder(&args.provider); + + // Initialize messages + let mut messages = + if let Some(system_prompt) = args.system_prompt.clone().filter(|s| !s.is_empty()) { + vec![OpenAIMessage { + role: "system".to_string(), + content: Some(OpenAIContent::Text(system_prompt)), + ..Default::default() + }] + } else { + vec![] + }; + + // Create user message with optional images + let mut parts = vec![ContentPart::Text { text: args.user_message.clone() }]; + if let Some(images) = &args.user_images { + for image in images.iter() { + if !image.s3.is_empty() { + parts.push(ContentPart::S3Object { s3_object: image.clone() }); + } + } + } + let user_content = OpenAIContent::Parts(parts); + + messages.push(OpenAIMessage { + role: "user".to_string(), + content: Some(user_content), + ..Default::default() + }); + + let mut actions = vec![]; + let mut content = None; + + // Check if this provider supports tools with the current output type + let supports_tools = query_builder.supports_tools_with_output_type(output_type); + + let mut tool_defs: Option> = if tools.is_empty() || !supports_tools { + None + } else { + Some(tools.iter().map(|t| t.def.clone()).collect()) + }; + + // Handle structured output schema + let has_output_properties = args + .output_schema + .as_ref() + .and_then(|schema| schema.properties.as_ref()) + .map(|props| !props.is_empty()) + .unwrap_or(false); + + let is_anthropic = is_anthropic_provider(&args.provider); + let mut used_structured_output_tool = false; + let mut structured_output_tool_name: Option = None; + + // For text output with schema, handle structured output + if has_output_properties && output_type == &OutputType::Text { + let schema = args.output_schema.as_ref().unwrap(); + if is_anthropic { + // Anthropic uses a tool for structured output + let unique_tool_name = find_unique_tool_name("structured_output", tool_defs.as_deref()); + structured_output_tool_name = Some(unique_tool_name.clone()); + + let output_tool = ToolDef { + r#type: "function".to_string(), + function: ToolDefFunction { + name: unique_tool_name, + description: Some( + "This tool MUST be used last to return a structured JSON object as the final output." + .to_string(), + ), + parameters: to_raw_value(&schema), + }, + }; + if let Some(ref mut existing_tools) = tool_defs { + existing_tools.push(output_tool); + } else { + tool_defs = Some(vec![output_tool]); + } + } + // For non-Anthropic providers, response_format is handled by the query builder + } + + // Main agent loop + for i in 0..MAX_AGENT_ITERATIONS { + if used_structured_output_tool { + break; + } + + // For text output or image output with tools + let build_args = BuildRequestArgs { + messages: &messages, + tools: tool_defs.as_deref(), + model: args.provider.get_model(), + temperature: args.temperature, + max_tokens: args.max_completion_tokens, + output_schema: args.output_schema.as_ref(), + output_type, + system_prompt: args.system_prompt.as_deref(), + user_message: &args.user_message, + images: args.user_images.as_deref(), + }; + + let request_body = query_builder + .build_request(&build_args, client, &job.workspace_id) + .await?; + + let endpoint = + query_builder.get_endpoint(&base_url, args.provider.get_model(), output_type); + let auth_headers = query_builder.get_auth_headers(api_key, output_type); + + let mut request = HTTP_CLIENT + .post(&endpoint) + .timeout(std::time::Duration::from_secs(REQUEST_TIMEOUT_SECONDS)) + .header("Content-Type", "application/json"); + + // Apply authentication headers + for (header_name, header_value) in auth_headers { + request = request.header(header_name, header_value); + } + + let resp = request + .body(request_body) + .send() + .await + .map_err(|e| Error::internal_err(format!("Failed to call API: {}", e)))?; + + match resp.error_for_status_ref() { + Ok(_) => { + let parsed = query_builder.parse_response(resp).await?; + + match parsed { + ParsedResponse::Text { content: response_content, tool_calls } => { + if let Some(ref response_content) = response_content { + actions.push(AgentAction::Message {}); + messages.push(OpenAIMessage { + role: "assistant".to_string(), + content: Some(OpenAIContent::Text(response_content.clone())), + agent_action: Some(AgentAction::Message {}), + ..Default::default() + }); + + update_flow_status_module_with_actions(db, parent_job, &actions) + .await?; + update_flow_status_module_with_actions_success(db, parent_job, true) + .await?; + + content = Some(OpenAIContent::Text(response_content.clone())); + } + + if tool_calls.is_empty() { + break; + } else if i == MAX_AGENT_ITERATIONS - 1 { + return Err(Error::internal_err( + "AI agent reached max iterations, but there are still tool calls" + .to_string(), + )); + } + + messages.push(OpenAIMessage { + role: "assistant".to_string(), + tool_calls: Some(tool_calls.clone()), + ..Default::default() + }); + + // Handle tool calls (keeping existing tool execution logic) + for tool_call in tool_calls.iter() { + // Check if this is the structured output tool + if structured_output_tool_name + .as_ref() + .map_or(false, |name| tool_call.function.name == *name) + { + used_structured_output_tool = true; + messages.push(OpenAIMessage { + role: "tool".to_string(), + content: Some(OpenAIContent::Text( + "Successfully ran structured_output tool".to_string(), + )), + tool_call_id: Some(tool_call.id.clone()), + ..Default::default() + }); + messages.push(OpenAIMessage { + role: "assistant".to_string(), + content: Some(OpenAIContent::Text( + tool_call.function.arguments.clone(), + )), + agent_action: Some(AgentAction::Message {}), + ..Default::default() + }); + content = + Some(OpenAIContent::Text(tool_call.function.arguments.clone())); + break; + } + + // Execute regular tool + let tool = tools + .iter() + .find(|t| t.def.function.name == tool_call.function.name); + if let Some(tool) = tool { + let job_id = ulid::Ulid::new().into(); + actions.push(AgentAction::ToolCall { + job_id, + function_name: tool_call.function.name.clone(), + module_id: tool.module.id.clone(), + }); + + update_flow_status_module_with_actions(db, parent_job, &actions) + .await?; + + let tool_call_args = + serde_json::from_str::>>( + &tool_call.function.arguments, + )?; + + let job_payload = match tool.module.get_value()? { + FlowModuleValue::Script { + path: script_path, + hash: script_hash, + tag_override, + .. + } => { + let payload = script_to_payload( + script_hash, + script_path, + db, + job, + &tool.module, + tag_override, + tool.module.apply_preprocessor, + ) + .await?; + payload + } + FlowModuleValue::RawScript { + path, + content, + language, + lock, + tag, + custom_concurrency_key, + concurrent_limit, + concurrency_time_window_s, + .. + } => { + let path = path.unwrap_or_else(|| { + format!( + "{}/tools/{}", + job.runnable_path(), + tool.module.id + ) + }); + + let payload = raw_script_to_payload( + path, + content, + language, + lock, + custom_concurrency_key, + concurrent_limit, + concurrency_time_window_s, + &tool.module, + tag, + tool.module.delete_after_use.unwrap_or(false), + ); + payload + } + _ => { + return Err(Error::internal_err(format!( + "Unsupported tool: {}", + tool_call.function.name + ))); + } + }; + + let mut tx = db.begin().await?; + + let job_perms = windmill_common::auth::get_job_perms( + &mut *tx, + &job.id, + &job.workspace_id, + ) + .await? + .map(|x| x.into()); + + let (email, permissioned_as) = + if let Some(on_behalf_of) = job_payload.on_behalf_of.as_ref() { + (&on_behalf_of.email, on_behalf_of.permissioned_as.clone()) + } else { + (&job.permissioned_as_email, job.permissioned_as.to_owned()) + }; + + let job_priority = tool.module.priority.or(job.priority); + + let tx = PushIsolationLevel::Transaction(tx); + let (uuid, tx) = push( + db, + tx, + &job.workspace_id, + job_payload.payload, + PushArgs { args: &tool_call_args, extra: None }, + &job.created_by, + email, + permissioned_as, + Some(&format!("job-span-{}", job.id)), + None, + job.schedule_path(), + Some(job.id), + None, + None, + Some(job_id), + false, + false, + None, + job.visible_to_owner, + Some(job.tag.clone()), + job_payload.timeout, + None, + job_priority, + job_perms.as_ref(), + true, + ) + .await?; + + tx.commit().await?; + + let tool_job = get_mini_pulled_job(db, &uuid).await?; + + let Some(tool_job) = tool_job else { + return Err(Error::internal_err( + "Tool job not found".to_string(), + )); + }; + + let tool_job = Arc::new(tool_job); + + let job_dir = create_job_dir(&worker_dir, job.id).await; + + let (inner_job_completed_tx, inner_job_completed_rx) = + JobCompletedSender::new(&conn, 1); + + let inner_job_completed_rx = inner_job_completed_rx.expect( + "inner_job_completed_tx should be set as agent jobs are not supported on agent workers", + ); + + #[cfg(feature = "benchmark")] + let mut bench = windmill_common::bench::BenchmarkIter::new(); + + match handle_queued_job( + tool_job.clone(), + None, + None, + None, + None, + conn, + client, + hostname, + worker_name, + worker_dir, + &job_dir, + None, + base_internal_url, + inner_job_completed_tx, + occupancy_metrics, + killpill_rx, + None, + #[cfg(feature = "benchmark")] + &mut bench, + ) + .await + { + Err(err) => { + let err_string = + format!("{}: {}", err.name(), err.to_string()); + let err_json = error_to_value(&err); + let _ = handle_non_flow_job_error( + db, + &tool_job, + 0, + None, + err_string.clone(), + err_json, + worker_name, + ) + .await; + messages.push(OpenAIMessage { + role: "tool".to_string(), + content: Some(OpenAIContent::Text(format!( + "Error running tool: {}", + err_string + ))), + tool_call_id: Some(tool_call.id.clone()), + agent_action: Some(AgentAction::ToolCall { + job_id, + function_name: tool_call.function.name.clone(), + module_id: tool.module.id.clone(), + }), + ..Default::default() + }); + update_flow_status_module_with_actions_success( + db, parent_job, false, + ) + .await?; + } + Ok(success) => { + let send_result = + inner_job_completed_rx.bounded_rx.try_recv().ok(); + + let result = if let Some(SendResult { + result: + SendResultPayload::JobCompleted(JobCompleted { + result, + .. + }), + .. + }) = send_result.as_ref() + { + job_completed_tx + .send( + send_result.as_ref().unwrap().result.clone(), + true, + ) + .await + .map_err(to_anyhow)?; + result + } else { + if let Some(send_result) = send_result { + job_completed_tx + .send(send_result.result, true) + .await + .map_err(to_anyhow)?; + } + return Err(Error::internal_err( + "Tool job completed but no result".to_string(), + )); + }; + + messages.push(OpenAIMessage { + role: "tool".to_string(), + content: Some(OpenAIContent::Text( + result.get().to_string(), + )), + tool_call_id: Some(tool_call.id.clone()), + agent_action: Some(AgentAction::ToolCall { + job_id, + function_name: tool_call.function.name.clone(), + module_id: tool.module.id.clone(), + }), + ..Default::default() + }); + update_flow_status_module_with_actions_success( + db, parent_job, success, + ) + .await?; + } + } + } else { + return Err(Error::internal_err(format!( + "Tool not found: {}", + tool_call.function.name + ))); + } + } + } + ParsedResponse::Image { base64_data } => { + // For image output with tools, we got an image response + let s3_object = upload_image_to_s3(&base64_data, job, client).await?; + return Ok(to_raw_value(&s3_object)); + } + } + } + Err(e) => { + let _status = resp.status(); + let text = resp + .text() + .await + .unwrap_or_else(|_| "".to_string()); + return Err(Error::internal_err(format!("API error: {} - {}", e, text))); + } + } + } + + // Return the final result + let final_messages: Vec = messages + .iter() + .map(|m| Message { message: m, agent_action: m.agent_action.as_ref() }) + .collect(); + + // Parse content as JSON for structured output, fallback to string if it fails + let output_value = match content { + Some(content_str) => match has_output_properties { + true => match content_str { + OpenAIContent::Text(text) => { + serde_json::from_str::>(&text).map_err(|_e| { + Error::internal_err(format!("Failed to parse structured output: {}", text)) + }) + } + OpenAIContent::Parts(_parts) => Err(Error::internal_err( + "Failed to parse structured output".to_string(), + )), + }, + false => Ok(match content_str { + OpenAIContent::Text(text) => to_raw_value(&text), + OpenAIContent::Parts(parts) => to_raw_value(&parts), + }), + }?, + None => to_raw_value(&""), + }; + + Ok(to_raw_value(&AIAgentResult { + output: output_value, + messages: final_messages, + })) +} diff --git a/backend/windmill-worker/src/lib.rs b/backend/windmill-worker/src/lib.rs index ace218c207..87f8282e9e 100644 --- a/backend/windmill-worker/src/lib.rs +++ b/backend/windmill-worker/src/lib.rs @@ -16,6 +16,7 @@ mod java_executor; #[cfg(feature = "ruby")] mod ruby_executor; +mod ai; mod ai_executor; mod bun_executor; pub mod common; diff --git a/frontend/src/lib/components/InputTransformForm.svelte b/frontend/src/lib/components/InputTransformForm.svelte index 617f822f46..fdf354da96 100644 --- a/frontend/src/lib/components/InputTransformForm.svelte +++ b/frontend/src/lib/components/InputTransformForm.svelte @@ -11,6 +11,7 @@ import type { Schema } from '$lib/common' import type { InputCat } from '$lib/utils' import { createEventDispatcher, getContext, untrack } from 'svelte' + import { computeShow } from '$lib/utils' import ArgInput from './ArgInput.svelte' import FieldHeader from './FieldHeader.svelte' @@ -58,6 +59,7 @@ hideHelpButton?: boolean class?: string editor?: SimpleEditor | undefined + otherArgs?: Record } let { @@ -79,7 +81,8 @@ enableAi = false, hideHelpButton = false, class: className = '', - editor = $bindable(undefined) + editor = $bindable(undefined), + otherArgs = {} }: Props = $props() let monaco: SimpleEditor | undefined = $state(undefined) @@ -87,6 +90,8 @@ let argInput: ArgInput | undefined = $state(undefined) let focusedPrev = false + let hidden = $state(false) + const variableMatch = (value: string): RegExpMatchArray | null => value.match(/^variable\('([^']+)'\)$/) const resourceMatch = (value: string): RegExpMatchArray | null => @@ -259,6 +264,47 @@ } } + function handleFieldVisibility( + schema: Schema | any, + arg: InputTransform | any, + otherArgs: Record + ) { + const schemaProperty = schema?.properties?.[argName] + if (schemaProperty?.showExpr) { + // Build args object with current field value and other context + const currentValue = propertyType === 'static' ? arg?.value : arg?.expr + + // Convert otherArgs from InputTransform objects to their actual values + const contextArgs = { + [argName]: currentValue + } + + // Extract values from InputTransform objects in otherArgs + Object.keys(otherArgs ?? {}).forEach((key) => { + const otherArg = otherArgs[key] + const otherArgValue = otherArg.type === 'static' ? otherArg.value : otherArg.expr + contextArgs[key] = otherArgValue + }) + + const shouldShow = computeShow(argName, schemaProperty.showExpr, contextArgs) + if (shouldShow) { + hidden = false + } else if (!hidden) { + hidden = true + // Clear the arg value when hidden (following SchemaForm pattern) + if (arg) { + arg.value = undefined + arg.expr = undefined + } + // Make sure validation passes when hidden + inputCheck = true + } + } else { + // No showExpr, always show + hidden = false + } + } + function onFocus() { focused = true if (isStaticTemplate(inputCat)) { @@ -359,12 +405,20 @@ $effect(() => { schema?.properties?.[argName]?.default && untrack(() => setDefaultCode()) }) + $effect.pre(() => { + // Monitor changes that affect field visibility + JSON.stringify(schema) + JSON.stringify(arg) + JSON.stringify(otherArgs) + + untrack(() => handleFieldVisibility(schema, arg, otherArgs)) + }) let connecting = $derived( $propPickerConfig?.propName == argName && $propPickerConfig?.insertionMode == 'connect' ) -{#if arg != undefined} +{#if arg != undefined && !hidden}
key !== argName) + )} />
{/if} diff --git a/frontend/src/lib/components/flows/flowInfers.ts b/frontend/src/lib/components/flows/flowInfers.ts index 8f419fce5c..7c24278cc2 100644 --- a/frontend/src/lib/components/flows/flowInfers.ts +++ b/frontend/src/lib/components/flows/flowInfers.ts @@ -62,50 +62,58 @@ export async function loadSchemaFromModule(module: FlowModule): Promise<{ type: 'object', format: 'ai-provider' }, + output_type: { + type: 'string', + description: + 'The type of output the AI agent will generate (text or image). Image output requires a configured workspace S3 storage, will ignore tools, and only works with OpenAI, Google AI and OpenRouter gemini-image-preview model.', + enum: ['text', 'image'], + default: 'text' + }, user_message: { - type: 'string' + type: 'string', + description: 'The message to give as input to the AI agent.' }, system_prompt: { - type: 'string' + type: 'string', + description: 'The system prompt to give as input to the AI agent.' }, - image: { - type: 'object', - description: 'Image to send to the AI agent (optional)', - format: 'resource-s3_object' + user_images: { + type: 'array', + description: + 'Array of images to give as input to the AI agent. Requires a configured workspace S3 storage.', + items: { + type: 'object' as const, + resourceType: 's3object' + } }, max_completion_tokens: { - type: 'number' + type: 'number', + description: 'The maximum number of output tokens.' }, temperature: { type: 'number', description: - 'Controls randomness in text generation. Range: 0.0 (deterministic) to 2.0 (random).' - }, - output_type: { - type: 'string', - description: - 'The type of output the AI agent will generate (text or image). Image output will ignore tools, and only works with OpenAI, Google AI and OpenRouter gemini-image-preview model.', - enum: ['text', 'image'], - default: 'text' + 'Controls randomness in text generation. Range: 0.0 (deterministic) to 2.0 (random).', + showExpr: "fields.output_type === 'text'" }, output_schema: { type: 'object', description: 'JSON schema that the AI agent will follow for its response format (only used if output_type is text)', - format: 'json-schema' + format: 'json-schema', + showExpr: "fields.output_type === 'text'" } }, - required: ['provider', 'model', 'user_message'], + required: ['provider', 'user_message', 'output_type'], type: 'object', order: [ 'provider', - 'model', + 'output_type', 'user_message', 'system_prompt', - 'image', + 'user_images', 'max_completion_tokens', 'temperature', - 'output_type', 'output_schema' ] }