diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 426b75ebbd..2ed7ee5507 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -15773,6 +15773,7 @@ dependencies = [ "lazy_static", "libloading 0.8.8", "mappable-rc", + "mime_guess", "mysql_async", "native-tls", "nix 0.27.1", diff --git a/backend/src/main.rs b/backend/src/main.rs index 78c2a1e637..849c4a7774 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -435,17 +435,16 @@ async fn windmill_main() -> anyhow::Result<()> { environment } else { load_base_url(&conn) - .await - .unwrap_or_else(|_| "local".to_string()) - .trim_start_matches("https://") - .trim_start_matches("http://") - .split(".") - .next() - .unwrap_or_else(|| "local") - .to_string() + .await + .unwrap_or_else(|_| "local".to_string()) + .trim_start_matches("https://") + .trim_start_matches("http://") + .split(".") + .next() + .unwrap_or_else(|| "local") + .to_string() }; - let _guard = windmill_common::tracing_init::initialize_tracing(&hostname, &mode, &environment); let is_agent = mode == Mode::Agent; diff --git a/backend/windmill-api/src/job_helpers_oss.rs b/backend/windmill-api/src/job_helpers_oss.rs index 61c9f0606b..bee7c2ec81 100644 --- a/backend/windmill-api/src/job_helpers_oss.rs +++ b/backend/windmill-api/src/job_helpers_oss.rs @@ -5,8 +5,6 @@ pub use crate::job_helpers_ee::*; #[cfg(not(feature = "private"))] use axum::Router; #[cfg(not(feature = "private"))] -use serde::Serialize; -#[cfg(not(feature = "private"))] use uuid::Uuid; #[cfg(not(feature = "private"))] use windmill_common::s3_helpers::StorageResourceType; @@ -32,12 +30,6 @@ use axum::response::Response; #[cfg(all(feature = "parquet", not(feature = "private")))] use serde::Deserialize; -#[derive(Serialize)] -#[cfg(not(feature = "private"))] -pub struct UploadFileResponse { - pub file_key: String, -} - #[derive(Deserialize)] #[cfg(not(feature = "private"))] pub struct LoadImagePreviewQuery { diff --git a/backend/windmill-common/src/client.rs b/backend/windmill-common/src/client.rs index 9507100eb5..a95ccbe292 100644 --- a/backend/windmill-common/src/client.rs +++ b/backend/windmill-common/src/client.rs @@ -206,4 +206,42 @@ impl AuthedClient { _ => Err(anyhow::anyhow!(response.text().await.unwrap_or_default()))?, } } + + pub async fn download_s3_file( + &self, + workspace_id: &str, + file_key: &str, + storage: Option, + ) -> anyhow::Result { + let mut query = vec![("file_key", file_key.to_string())]; + if let Some(storage) = storage { + query.push(("storage", storage)); + } + let response = self + .force_client + .as_ref() + .unwrap_or(&HTTP_CLIENT) + .get(&format!( + "{}/api/w/{}/job_helpers/download_s3_file", + self.base_internal_url, workspace_id + )) + .query(&query) + .header( + reqwest::header::AUTHORIZATION, + reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.token)) + .map_err(|e| anyhow::anyhow!(e.to_string()))?, + ) + .send() + .await + .context("Failed to send download_s3_file request") + .map_err(|e| anyhow::anyhow!(e.to_string()))?; + + match response.status().as_u16() { + 200u16 => Ok(response + .bytes() + .await + .context("Failed to read response bytes")?), + _ => Err(anyhow::anyhow!(response.text().await.unwrap_or_default())), + } + } } diff --git a/backend/windmill-worker/Cargo.toml b/backend/windmill-worker/Cargo.toml index 772a6195f9..de7261a44e 100644 --- a/backend/windmill-worker/Cargo.toml +++ b/backend/windmill-worker/Cargo.toml @@ -112,6 +112,7 @@ nix.workspace = true bytes.workspace = true reqwest.workspace = true reqwest-middleware.workspace = true +mime_guess.workspace = true hex.workspace = true tiberius = { workspace = true, optional = true } tokio-util = { workspace = true, optional = true } diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 9778884553..7f5df68bc8 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -1,4 +1,6 @@ use async_recursion::async_recursion; +use base64::Engine; +use mime_guess; use regex::Regex; use serde::{Deserialize, Serialize}; use serde_json::value::RawValue; @@ -16,6 +18,7 @@ use windmill_common::{ flows::{FlowModule, 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}, @@ -37,32 +40,12 @@ use crate::{ }; 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(); } -/// 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()) -} - #[derive(Deserialize, Serialize, Clone, Debug)] struct OpenAIFunction { name: String, @@ -76,11 +59,39 @@ struct OpenAIToolCall { r#type: String, } -#[derive(Deserialize, Serialize, Clone, Default)] +#[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, + content: Option, #[serde(skip_serializing_if = "Option::is_none")] tool_calls: Option>, #[serde(skip_serializing_if = "Option::is_none")] @@ -108,6 +119,153 @@ 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, @@ -154,6 +312,19 @@ struct Tool { 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, @@ -162,6 +333,8 @@ struct AIAgentArgs { temperature: Option, max_completion_tokens: Option, output_schema: Option, + output_type: Option, + image: Option, } #[derive(Deserialize, Debug)] @@ -413,6 +586,483 @@ impl OpenAPISchema { } } +/// 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, @@ -506,220 +1156,7 @@ fn parse_raw_script_schema(content: &str, language: &ScriptLang) -> Result, -) -> error::Result<(bool, Arc>)> { - 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, - agent_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/{}", agent_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, &agent_job.id, &agent_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 { - ( - &agent_job.permissioned_as_email, - agent_job.permissioned_as.to_owned(), - ) - }; - - let job_priority = tool_module.priority.or(agent_job.priority); - - let tx = PushIsolationLevel::Transaction(tx); - let (uuid, tx) = push( - db, - tx, - &agent_job.workspace_id, - job_payload.payload, - PushArgs { args: &tool_call_args, extra: None }, - &agent_job.created_by, - email, - permissioned_as, - Some(&format!("job-span-{}", agent_job.id)), - None, - agent_job.schedule_path(), - Some(agent_job.id), - None, - None, - Some(job_id), - false, - false, - None, - agent_job.visible_to_owner, - Some(agent_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, agent_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, - err_json, - worker_name, - ) - .await; - Err(err) - } - 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(), - )); - }; - - Ok((success, result.clone())) - } - } -} - +#[async_recursion] async fn run_agent( // connection db: &DB, @@ -728,8 +1165,8 @@ async fn run_agent( // agent job and flow data job: &MiniPulledJob, parent_job: &uuid::Uuid, - args: AIAgentArgs, - tools: Vec, + args: &AIAgentArgs, + tools: &[Tool], // job execution context client: &AuthedClient, @@ -741,290 +1178,522 @@ async fn run_agent( hostname: &str, killpill_rx: &mut tokio::sync::broadcast::Receiver<()>, ) -> error::Result> { - let mut messages = if let Some(system_prompt) = args.system_prompt.filter(|s| !s.is_empty()) { - vec![OpenAIMessage { - role: "system".to_string(), - content: Some(system_prompt), - ..Default::default() - }] - } else { - vec![] - }; + let output_type = args.output_type.as_ref().unwrap_or(&OutputType::Text); - messages.push(OpenAIMessage { - role: "user".to_string(), - content: Some(args.user_message), - ..Default::default() - }); + match *output_type { + OutputType::Image => { + let (s3_result, messages) = handle_image_output(&args, job, client, db).await?; - let mut actions = vec![]; + let final_messages: Vec = messages + .iter() + .map(|m| Message { message: m, agent_action: m.agent_action.as_ref() }) + .collect(); - let mut content = None; - - let base_url = args.provider.get_base_url(db).await?; - let api_key = args.provider.get_api_key(); - - 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); + if let Some(s3_output) = s3_result { + Ok(to_raw_value(&s3_output)) } else { - tool_defs = Some(vec![output_tool]); + Ok(to_raw_value(&AIAgentResult { + output: to_raw_value(&None::), + messages: final_messages, + })) } - } 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), - }, - }); } - } + OutputType::Text => { + let base_url = args.provider.get_base_url(db).await?; + let api_key = args.provider.get_api_key(); - for i in 0..MAX_AGENT_ITERATIONS { - if used_structured_output_tool { - break; - } + 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![] + }; - let response = { - let resp = HTTP_CLIENT - .post(format!("{}/chat/completions", base_url)) - .bearer_auth(api_key) - .json(&OpenAIRequest { - model: args.provider.get_model(), - messages: &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 - ))); + // 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()) + }; - 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 content) = content { - actions.push(AgentAction::Message {}); messages.push(OpenAIMessage { - role: "assistant".to_string(), - content: Some(content.clone()), - agent_action: Some(AgentAction::Message {}), + role: "user".to_string(), + content: Some(user_content), ..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?; - } + let mut actions = vec![]; + let mut content = None; - 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(), - )); - } + let mut tool_defs: Option> = if tools.is_empty() { + None + } else { + Some(tools.iter().map(|t| t.def.clone()).collect()) + }; - 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 + let has_output_properties = args + .output_schema .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("Successfully ran structured_output tool".to_string()), - tool_call_id: Some(tool_call.id.clone()), - ..Default::default() - }); + .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(), - content: Some(tool_call.function.arguments.clone()), - agent_action: Some(AgentAction::Message {}), + tool_calls: Some(tool_calls.clone()), ..Default::default() }); - content = Some(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?; - - match call_tool( - db, - conn, - job, - &tool.module, - &tool_call, - job_id, - client, - occupancy_metrics, - base_internal_url, - worker_dir, - worker_name, - hostname, - job_completed_tx, - killpill_rx, - ) - .await - { - Ok((success, result)) => { + 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(result.get().to_string()), + content: Some(OpenAIContent::Text( + "Successfully ran structured_output tool".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?; + 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; } - Err(err) => { - let err_string = format!("{}: {}", err.name(), err.to_string()); - messages.push(OpenAIMessage { - role: "tool".to_string(), - content: Some(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() + + 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_success(db, parent_job, false) - .await?; + + 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 + ))); } } - } 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, + })) } } - - 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 => serde_json::from_str::>(&content_str).map_err(|_e| { - Error::internal_err(format!( - "Failed to parse structured output: {}", - content_str - )) - })?, - false => to_raw_value(&content_str), - }, - None => to_raw_value(&""), - }; - - Ok(to_raw_value(&AIAgentResult { - output: output_value, - messages: final_messages, - })) } pub struct FlowJobRunnableIdAndRawFlow { @@ -1221,8 +1890,8 @@ pub async fn handle_ai_agent_job( conn, job, parent_job, - args, - tools, + &args, + &tools, client, &mut inner_occupancy_metrics, job_completed_tx, diff --git a/frontend/src/lib/components/flows/flowInfers.ts b/frontend/src/lib/components/flows/flowInfers.ts index 3b4f36aebe..8f419fce5c 100644 --- a/frontend/src/lib/components/flows/flowInfers.ts +++ b/frontend/src/lib/components/flows/flowInfers.ts @@ -68,6 +68,11 @@ export async function loadSchemaFromModule(module: FlowModule): Promise<{ system_prompt: { type: 'string' }, + image: { + type: 'object', + description: 'Image to send to the AI agent (optional)', + format: 'resource-s3_object' + }, max_completion_tokens: { type: 'number' }, @@ -76,9 +81,17 @@ export async function loadSchemaFromModule(module: FlowModule): Promise<{ 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' + }, output_schema: { type: 'object', - description: 'JSON schema that the AI agent will follow for its response format', + description: + 'JSON schema that the AI agent will follow for its response format (only used if output_type is text)', format: 'json-schema' } }, @@ -89,8 +102,10 @@ export async function loadSchemaFromModule(module: FlowModule): Promise<{ 'model', 'user_message', 'system_prompt', + 'image', 'max_completion_tokens', 'temperature', + 'output_type', 'output_schema' ] }