feat(ai agent): handle images in ai agent (#6572)

* 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

* cleaning

* fix for openai

* better mime type

* add descriptions

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",

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;

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 {

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<String>,
+    ) -> anyhow::Result<bytes::Bytes> {
+        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())),
+        }
+    }
 }

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 }

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'
 			]
 		}