windmill/backend/windmill-ai/src/image_handler.rs

use crate::types::*;
use base64::Engine;
use futures;
use ulid;
use uuid::Uuid;
use windmill_common::{client::AuthedClient, error::Error};
use windmill_types::s3::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,
    workspace_id: &str,
    job_id: &Uuid,
    client: &AuthedClient,
) -> Result<S3Object, Error> {
    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(
            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.
///
/// The caller must provide an AuthedClient authorized for `workspace_id`.
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))
}

/// 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,
    workspace_id: &str,
) -> Result<ContentPart, Error> {
    let (mime_type, file_bytes) =
        download_and_encode_s3_image(s3_object, client, workspace_id).await?;
    let data_url = format!("data:{};base64,{}", mime_type, file_bytes);

    if crate::ai_types::is_document_mime(&mime_type) {
        let filename = s3_object
            .s3
            .rsplit('/')
            .next()
            .unwrap_or("document.pdf")
            .to_string();
        Ok(ContentPart::File { file: FileData { filename, file_data: data_url } })
    } else {
        Ok(ContentPart::ImageUrl { image_url: ImageUrlData { url: data_url } })
    }
}

/// 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,
    workspace_id: &str,
) -> Result<Vec<OpenAIMessage>, 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 } => {
                                prepared_content.push(
                                    s3_object_to_content_part(s3_object, client, workspace_id)
                                        .await?,
                                );
                            }
                            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)
}