From 6f7d31e56ba7a1357f0ffeda221b20d4e0cf2b78 Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Mon, 11 May 2026 16:59:40 +0200 Subject: [PATCH] refactor: move ai image handling to windmill-ai (#9098) --- backend/Cargo.lock | 4 +++ backend/windmill-ai/Cargo.toml | 4 +++ .../ai => windmill-ai/src}/image_handler.rs | 27 ++++++++++++------- backend/windmill-ai/src/lib.rs | 1 + backend/windmill-worker/src/ai/mod.rs | 1 - .../src/ai/providers/anthropic.rs | 3 +-- .../src/ai/providers/bedrock.rs | 2 +- .../src/ai/providers/google_ai.rs | 3 +-- .../src/ai/providers/openai.rs | 3 +-- .../src/ai/providers/openrouter.rs | 3 ++- .../windmill-worker/src/ai/providers/other.rs | 3 +-- backend/windmill-worker/src/ai_executor.rs | 9 +++---- 12 files changed, 38 insertions(+), 25 deletions(-) rename backend/{windmill-worker/src/ai => windmill-ai/src}/image_handler.rs (86%) diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 842799d5f9..719890fd97 100644 --- a/backend/Cargo.lock +++ b/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", diff --git a/backend/windmill-ai/Cargo.toml b/backend/windmill-ai/Cargo.toml index 3cfea846b9..542c390cc0 100644 --- a/backend/windmill-ai/Cargo.toml +++ b/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 } diff --git a/backend/windmill-worker/src/ai/image_handler.rs b/backend/windmill-ai/src/image_handler.rs similarity index 86% rename from backend/windmill-worker/src/ai/image_handler.rs rename to backend/windmill-ai/src/image_handler.rs index 089d73dc42..5bd9794da5 100644 --- a/backend/windmill-worker/src/ai/image_handler.rs +++ b/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 { 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, diff --git a/backend/windmill-ai/src/lib.rs b/backend/windmill-ai/src/lib.rs index 8b786cb741..945091613b 100644 --- a/backend/windmill-ai/src/lib.rs +++ b/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; diff --git a/backend/windmill-worker/src/ai/mod.rs b/backend/windmill-worker/src/ai/mod.rs index 2423fca372..b004b0ec20 100644 --- a/backend/windmill-worker/src/ai/mod.rs +++ b/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; diff --git a/backend/windmill-worker/src/ai/providers/anthropic.rs b/backend/windmill-worker/src/ai/providers/anthropic.rs index f2c310923f..0a3c5df6e4 100644 --- a/backend/windmill-worker/src/ai/providers/anthropic.rs +++ b/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 diff --git a/backend/windmill-worker/src/ai/providers/bedrock.rs b/backend/windmill-worker/src/ai/providers/bedrock.rs index ec499f5fcd..8c37433cab 100644 --- a/backend/windmill-worker/src/ai/providers/bedrock.rs +++ b/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}, }; diff --git a/backend/windmill-worker/src/ai/providers/google_ai.rs b/backend/windmill-worker/src/ai/providers/google_ai.rs index 2a00011de7..19c50d06de 100644 --- a/backend/windmill-worker/src/ai/providers/google_ai.rs +++ b/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 // ============================================================================ diff --git a/backend/windmill-worker/src/ai/providers/openai.rs b/backend/windmill-worker/src/ai/providers/openai.rs index 691f8fd5e3..e0767d5dda 100644 --- a/backend/windmill-worker/src/ai/providers/openai.rs +++ b/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)] diff --git a/backend/windmill-worker/src/ai/providers/openrouter.rs b/backend/windmill-worker/src/ai/providers/openrouter.rs index 9bfc47ed8f..ede541cd13 100644 --- a/backend/windmill-worker/src/ai/providers/openrouter.rs +++ b/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)] diff --git a/backend/windmill-worker/src/ai/providers/other.rs b/backend/windmill-worker/src/ai/providers/other.rs index 7982dc739d..a840517133 100644 --- a/backend/windmill-worker/src/ai/providers/other.rs +++ b/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 { diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 4dfeee3571..cdf2995f9c 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/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);