refactor: move ai image handling to windmill-ai (#9098)

backend/Cargo.lock

@@ -16114,8 +16114,11 @@ dependencies = [
  "aws-sdk-bedrockruntime",
  "aws-smithy-types",
  "base64 0.22.1",
+ "bytes",
  "eventsource-stream",
+ "futures",
  "lazy_static",
+ "mime_guess",
  "reqwest 0.13.1",
  "serde",
  "serde_json",
@@ -16123,6 +16126,7 @@ dependencies = [
  "tokio",
  "tokio-stream",
  "tracing",
+ "ulid",
  "uuid",
  "windmill-common",
  "windmill-mcp",

backend/windmill-ai/Cargo.toml

@@ -21,7 +21,10 @@ windmill-mcp = { workspace = true, optional = true }
 
 async-trait.workspace = true
 base64.workspace = true
+bytes.workspace = true
 eventsource-stream.workspace = true
+futures.workspace = true
+mime_guess.workspace = true
 reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true
@@ -31,6 +34,7 @@ lazy_static.workspace = true
 tracing.workspace = true
 tokio.workspace = true
 tokio-stream.workspace = true
+ulid.workspace = true
 
 # Bedrock (optional)
 aws-config = { workspace = true, optional = true }

backend/windmill-ai/src/image_handler.rs

@@ -1,15 +1,18 @@
+use crate::types::*;
 use base64::Engine;
 use futures;
 use ulid;
-use windmill_ai::types::*;
+use uuid::Uuid;
 use windmill_common::{client::AuthedClient, error::Error};
-use windmill_queue::MiniPulledJob;
 use windmill_types::s3::S3Object;
 
-/// Upload image to S3 and return S3Object
+/// Upload image to S3 and return S3Object.
+///
+/// The caller must provide an AuthedClient authorized for `workspace_id`.
 pub async fn upload_image_to_s3(
     base64_image: &str,
-    job: &MiniPulledJob,
+    workspace_id: &str,
+    job_id: &Uuid,
     client: &AuthedClient,
 ) -> Result<S3Object, Error> {
     let image_bytes = base64::engine::general_purpose::STANDARD
@@ -18,7 +21,7 @@ pub async fn upload_image_to_s3(
 
     // Generate unique S3 key
     let unique_id = ulid::Ulid::new().to_string();
-    let s3_key = format!("ai_images/{}/{}.png", job.id, unique_id);
+    let s3_key = format!("ai_images/{}/{}.png", job_id, unique_id);
 
     // Create byte stream
     let byte_stream = futures::stream::once(async move {
@@ -28,7 +31,7 @@ pub async fn upload_image_to_s3(
     // Upload to S3
     client
         .upload_s3_file(
-            &job.workspace_id,
+            workspace_id,
             s3_key.clone(),
             None, // storage - use default
             byte_stream,
@@ -44,7 +47,9 @@ pub async fn upload_image_to_s3(
     })
 }
 
-/// Download an S3 image and convert it to a base64 data URL
+/// Download an S3 image and convert it to a base64 data URL.
+///
+/// The caller must provide an AuthedClient authorized for `workspace_id`.
 pub async fn download_and_encode_s3_image(
     image: &S3Object,
     client: &AuthedClient,
@@ -70,6 +75,8 @@ pub async fn download_and_encode_s3_image(
 }
 
 /// Convert an S3Object to the appropriate ContentPart based on MIME type.
+///
+/// The caller must provide an AuthedClient authorized for `workspace_id`.
 pub async fn s3_object_to_content_part(
     s3_object: &S3Object,
     client: &AuthedClient,
@@ -79,7 +86,7 @@ pub async fn s3_object_to_content_part(
         download_and_encode_s3_image(s3_object, client, workspace_id).await?;
     let data_url = format!("data:{};base64,{}", mime_type, file_bytes);
 
-    if windmill_ai::ai_types::is_document_mime(&mime_type) {
+    if crate::ai_types::is_document_mime(&mime_type) {
         let filename = s3_object
             .s3
             .rsplit('/')
@@ -92,7 +99,9 @@ pub async fn s3_object_to_content_part(
     }
 }
 
-/// Prepare messages for API by converting S3Objects to base64 ImageUrls
+/// Prepare messages for API by converting S3Objects to base64 ImageUrls.
+///
+/// The caller must provide an AuthedClient authorized for `workspace_id`.
 pub async fn prepare_messages_for_api(
     messages: &[OpenAIMessage],
     client: &AuthedClient,

backend/windmill-ai/src/lib.rs

@@ -4,6 +4,7 @@ pub mod ai_cache;
 pub mod ai_google;
 pub mod ai_providers;
 pub mod ai_types;
+pub mod image_handler;
 pub mod query_builder;
 pub mod sse;
 pub mod types;

backend/windmill-worker/src/ai/mod.rs

@@ -1,7 +1,6 @@
 // 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 tools;

backend/windmill-worker/src/ai/providers/anthropic.rs

@@ -4,6 +4,7 @@ use serde_json::value::RawValue;
 use windmill_ai::{
     ai_google::parse_data_url,
     ai_providers::AIProvider,
+    image_handler::prepare_messages_for_api,
     query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink},
     sse::{AnthropicSSEParser, SSEParser},
     types::*,
@@ -11,8 +12,6 @@ use windmill_ai::{
 };
 use windmill_common::{client::AuthedClient, error::Error};
 
-use crate::ai::image_handler::prepare_messages_for_api;
-
 /// Anthropic API version for standard API
 const ANTHROPIC_VERSION_STANDARD: &str = "2023-06-01";
 /// Anthropic API version for Google Vertex AI

backend/windmill-worker/src/ai/providers/bedrock.rs

@@ -6,9 +6,9 @@
 //! - Stream event parsing
 //! - Helper utilities
 
-use crate::ai::image_handler::prepare_messages_for_api;
 use std::collections::HashMap;
 use windmill_ai::{
+    image_handler::prepare_messages_for_api,
     query_builder::{ParsedResponse, StreamEventSink},
     types::{OpenAIMessage, StreamingEvent, TokenUsage, ToolDef},
 };

backend/windmill-worker/src/ai/providers/google_ai.rs

@@ -5,14 +5,13 @@ use windmill_ai::{
         GeminiImageContent, GeminiImageRequest, GeminiImageResponse, GeminiInlineData, GeminiPart,
         GeminiPredictContent, GeminiTextRequest, GeminiTool,
     },
+    image_handler::{download_and_encode_s3_image, prepare_messages_for_api},
     query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink},
     sse::{GeminiSSEParser, SSEParser},
     types::*,
 };
 use windmill_common::{client::AuthedClient, error::Error};
 
-use crate::ai::image_handler::{download_and_encode_s3_image, prepare_messages_for_api};
-
 // ============================================================================
 // Query Builder Implementation
 // ============================================================================

backend/windmill-worker/src/ai/providers/openai.rs

@@ -4,6 +4,7 @@ use serde_json::value::RawValue;
 use windmill_ai::{
     ai_providers::AIProvider,
     ai_types::OpenAIToolCall,
+    image_handler::{prepare_messages_for_api, s3_object_to_content_part},
     query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink},
     sse::{OpenAIResponsesSSEParser, SSEParser},
     types::*,
@@ -11,8 +12,6 @@ use windmill_ai::{
 };
 use windmill_common::{client::AuthedClient, error::Error};
 
-use crate::ai::image_handler::{prepare_messages_for_api, s3_object_to_content_part};
-
 // Responses API structures
 #[derive(Deserialize)]
 #[allow(dead_code)]

backend/windmill-worker/src/ai/providers/openrouter.rs

@@ -3,12 +3,13 @@ use serde::{Deserialize, Serialize};
 use serde_json;
 use windmill_ai::{
     ai_providers::AIProvider,
+    image_handler::prepare_messages_for_api,
     query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink},
     types::*,
 };
 use windmill_common::{client::AuthedClient, error::Error};
 
-use crate::ai::{image_handler::prepare_messages_for_api, providers::other::OtherQueryBuilder};
+use crate::ai::providers::other::OtherQueryBuilder;
 
 // OpenRouter-specific types
 #[derive(Serialize)]

backend/windmill-worker/src/ai/providers/other.rs

@@ -3,6 +3,7 @@ use serde::Serialize;
 use serde_json;
 use windmill_ai::{
     ai_providers::AIProvider,
+    image_handler::prepare_messages_for_api,
     query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink},
     sse::{OpenAISSEParser, SSEParser},
     types::*,
@@ -10,8 +11,6 @@ use windmill_ai::{
 };
 use windmill_common::{client::AuthedClient, error::Error};
 
-use crate::ai::image_handler::prepare_messages_for_api;
-
 #[derive(Serialize, Debug, Clone)]
 #[serde(rename_all = "lowercase")]
 pub enum ToolChoice {

backend/windmill-worker/src/ai_executor.rs

@@ -22,6 +22,7 @@ use windmill_mcp::McpClient;
 use crate::ai::tools::McpClientStub as McpClient;
 use windmill_ai::{
     ai_providers::AIProvider,
+    image_handler::upload_image_to_s3,
     query_builder::{BuildRequestArgs, ParsedResponse},
     types::*,
     utils::{should_use_structured_output_tool, AI_HTTP_HEADERS},
@@ -43,10 +44,7 @@ use windmill_common::{
 use windmill_queue::{cancel_single_job, CanceledBy, MiniPulledJob};
 
 use crate::{
-    ai::{
-        image_handler::upload_image_to_s3,
-        query_builder::{create_query_builder, StreamEventProcessor},
-    },
+    ai::query_builder::{create_query_builder, StreamEventProcessor},
     common::{build_args_map, resolve_job_timeout, OccupancyMetrics, StreamNotifier},
     handle_child::{run_future_with_polling_update_job_poller_graceful, GracefulPollOutcome},
 };
@@ -1206,7 +1204,8 @@ pub async fn run_agent(
             }
             ParsedResponse::Image { base64_data } => {
                 // For image output, upload to S3 and track in conversation
-                let s3_object = upload_image_to_s3(&base64_data, job, client).await?;
+                let s3_object =
+                    upload_image_to_s3(&base64_data, &job.workspace_id, &job.id, client).await?;
 
                 let content = to_raw_value(&s3_object);