diff --git a/crates/api_clients/artcraft/artcraft_api_defs/src/generate/image/multi_function/gpt_image_1p5_multi_function_image_gen.rs b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/image/multi_function/gpt_image_1p5_multi_function_image_gen.rs index d53b1e68a4..e95901f746 100644 --- a/crates/api_clients/artcraft/artcraft_api_defs/src/generate/image/multi_function/gpt_image_1p5_multi_function_image_gen.rs +++ b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/image/multi_function/gpt_image_1p5_multi_function_image_gen.rs @@ -20,7 +20,7 @@ pub struct GptImage1p5MultiFunctionImageGenRequest { /// Only for image editing, which turns this into an inpainting task. /// Text to image should not set this. - pub mask_image_token: Option, + pub mask_image_token: Option, /// Number of images to generate. Default is one. pub num_images: Option, diff --git a/crates/schema/public/enums/src/common/model_type.rs b/crates/schema/public/enums/src/common/model_type.rs index e1f0aaf115..9e2fed9b6b 100644 --- a/crates/schema/public/enums/src/common/model_type.rs +++ b/crates/schema/public/enums/src/common/model_type.rs @@ -29,6 +29,8 @@ pub enum ModelType { FluxProKontextMax, #[serde(rename = "gpt_image_1")] GptImage1, + #[serde(rename = "gpt_image_1p5")] + GptImage1p5, // Generic grok image model without a version #[serde(rename = "grok_image")] GrokImage, @@ -117,6 +119,7 @@ impl ModelType { Self::FluxPro11Ultra => "flux_pro_1p1_ultra", Self::FluxProKontextMax => "flux_pro_kontext_max", Self::GptImage1 => "gpt_image_1", + Self::GptImage1p5 => "gpt_image_1p5", Self::GrokImage => "grok_image", Self::Recraft3 => "recraft_3", Self::SeedEdit3 => "seededit_3", @@ -162,6 +165,7 @@ impl ModelType { "flux_pro_1p1_ultra" => Ok(Self::FluxPro11Ultra), "flux_pro_kontext_max" => Ok(Self::FluxProKontextMax), "gpt_image_1" => Ok(Self::GptImage1), + "gpt_image_1p5" => Ok(Self::GptImage1p5), "grok_image" => Ok(Self::GrokImage), "recraft_3" => Ok(Self::Recraft3), "seededit_3" => Ok(Self::SeedEdit3), @@ -211,6 +215,7 @@ impl ModelType { Self::FluxPro11Ultra, Self::FluxProKontextMax, Self::GptImage1, + Self::GptImage1p5, Self::GrokImage, Self::Recraft3, Self::SeedEdit3, @@ -265,6 +270,7 @@ mod tests { assert_serialization(ModelType::FluxPro11Ultra, "flux_pro_1p1_ultra"); assert_serialization(ModelType::FluxProKontextMax, "flux_pro_kontext_max"); assert_serialization(ModelType::GptImage1, "gpt_image_1"); + assert_serialization(ModelType::GptImage1p5, "gpt_image_1p5"); assert_serialization(ModelType::GrokImage, "grok_image"); assert_serialization(ModelType::Recraft3, "recraft_3"); assert_serialization(ModelType::SeedEdit3, "seededit_3"); @@ -307,6 +313,7 @@ mod tests { assert_eq!(ModelType::FluxPro11Ultra.to_str(), "flux_pro_1p1_ultra"); assert_eq!(ModelType::FluxProKontextMax.to_str(), "flux_pro_kontext_max"); assert_eq!(ModelType::GptImage1.to_str(), "gpt_image_1"); + assert_eq!(ModelType::GptImage1p5.to_str(), "gpt_image_1p5"); assert_eq!(ModelType::GrokImage.to_str(), "grok_image"); assert_eq!(ModelType::Recraft3.to_str(), "recraft_3"); assert_eq!(ModelType::SeedEdit3.to_str(), "seededit_3"); @@ -351,6 +358,7 @@ mod tests { assert_eq!(ModelType::from_str("flux_pro_1p1_ultra").unwrap(), ModelType::FluxPro11Ultra); assert_eq!(ModelType::from_str("flux_pro_kontext_max").unwrap(), ModelType::FluxProKontextMax); assert_eq!(ModelType::from_str("gpt_image_1").unwrap(), ModelType::GptImage1); + assert_eq!(ModelType::from_str("gpt_image_1p5").unwrap(), ModelType::GptImage1p5); assert_eq!(ModelType::from_str("grok_image").unwrap(), ModelType::GrokImage); assert_eq!(ModelType::from_str("recraft_3").unwrap(), ModelType::Recraft3); assert_eq!(ModelType::from_str("seededit_3").unwrap(), ModelType::SeedEdit3); @@ -385,7 +393,7 @@ mod tests { #[test] fn all_variants() { let mut variants = ModelType::all_variants(); - assert_eq!(variants.len(), 35); + assert_eq!(variants.len(), 36); // Image models assert_eq!(variants.pop_first(), Some(ModelType::Flux1Dev)); assert_eq!(variants.pop_first(), Some(ModelType::Flux1Schnell)); @@ -395,6 +403,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(ModelType::FluxPro11Ultra)); assert_eq!(variants.pop_first(), Some(ModelType::FluxProKontextMax)); assert_eq!(variants.pop_first(), Some(ModelType::GptImage1)); + assert_eq!(variants.pop_first(), Some(ModelType::GptImage1p5)); assert_eq!(variants.pop_first(), Some(ModelType::GrokImage)); assert_eq!(variants.pop_first(), Some(ModelType::Recraft3)); assert_eq!(variants.pop_first(), Some(ModelType::SeedEdit3)); diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/gpt_image_1p5_multi_function_image_gen_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/gpt_image_1p5_multi_function_image_gen_handler.rs new file mode 100644 index 0000000000..520c373b09 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/gpt_image_1p5_multi_function_image_gen_handler.rs @@ -0,0 +1,355 @@ +use std::collections::{HashMap, HashSet}; +use std::iter::FromIterator; +use std::sync::Arc; + +use crate::http_server::common_responses::common_web_error::CommonWebError; +use crate::http_server::common_responses::media::media_links_builder::MediaLinksBuilder; +use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test; +use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain; +use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format; +use crate::state::server_state::ServerState; +use crate::util::lookup::lookup_image_urls_as_map::lookup_image_urls_as_map; +use actix_web::web::Json; +use actix_web::{web, HttpRequest}; +use artcraft_api_defs::generate::image::multi_function::gpt_image_1p5_multi_function_image_gen::{GptImage1p5MultiFunctionImageGenBackground, GptImage1p5MultiFunctionImageGenInputFidelity, GptImage1p5MultiFunctionImageGenNumImages, GptImage1p5MultiFunctionImageGenQuality, GptImage1p5MultiFunctionImageGenRequest, GptImage1p5MultiFunctionImageGenResponse, GptImage1p5MultiFunctionImageGenSize}; +use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath; +use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType; +use enums::by_table::prompts::prompt_type::PromptType; +use enums::common::generation_provider::GenerationProvider; +use enums::common::model_type::ModelType; +use enums::common::visibility::Visibility; +use fal_client::creds::open_ai_api_key::OpenAiApiKey; +use fal_client::requests::webhook::image::edit::enqueue_gpt_image_1p5_edit_image_webhook::{enqueue_gpt_image_1p5_image_edit_webhook, EnqueueGptImage1p5EditImageArgs, EnqueueGptImage1p5EditImageBackground, EnqueueGptImage1p5EditImageInputFidelity, EnqueueGptImage1p5EditImageNumImages, EnqueueGptImage1p5EditImageQuality, EnqueueGptImage1p5EditImageSize}; +use fal_client::requests::webhook::image::text::enqueue_gpt_image_1p5_text_to_image_webhook::{enqueue_gpt_image_1p5_text_to_image_webhook, EnqueueGptImage1p5TextToImageArgs, EnqueueGptImage1p5TextToImageBackground, EnqueueGptImage1p5TextToImageNumImages, EnqueueGptImage1p5TextToImageQuality, EnqueueGptImage1p5TextToImageSize}; +use http_server_common::request::get_request_ip::get_request_ip; +use log::{error, info, warn}; +use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::insert_generic_inference_job_for_fal_queue; +use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::FalCategory; +use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::InsertGenericInferenceForFalArgs; +use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; +use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::{batch_get_media_files_by_tokens, batch_get_media_files_by_tokens_with_connection}; +use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem}; +use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs}; +use server_environment::ServerEnvironment; +use sqlx::pool::PoolConnection; +use sqlx::{Acquire, MySql}; +use tokens::tokens::media_files::MediaFileToken; +use utoipa::ToSchema; + +/// Gpt Image 1.5 Multi-Function (generate + edit + inpaint) +#[utoipa::path( + post, + tag = "Generate Images (Multi-Function)", + path = "/v1/generate/image/multi_function/gpt_image_1p5", + responses( + (status = 200, description = "Success", body = GptImage1p5MultiFunctionImageGenResponse), + ), + params( + ("request" = GptImage1p5MultiFunctionImageGenRequest, description = "Payload for Request"), + ) +)] +pub async fn nano_banana_pro_multi_function_image_gen_handler( + http_request: HttpRequest, + request: Json, + server_state: web::Data> +) -> Result, CommonWebError> { + + payments_error_test(&request.prompt.as_deref().unwrap_or(""))?; + + let mut mysql_connection = server_state.mysql_pool + .acquire() + .await?; + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session_from_connection(&http_request, &mut mysql_connection) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + CommonWebError::ServerError + })?; + + let maybe_avt_token = server_state + .avt_cookie_manager + .get_avt_token_from_request(&http_request); + + // TODO: Limit usage for new accounts. Billing, free credits metering, etc. + + //let user_session = match maybe_user_session { + // Some(session) => session, + // None => { + // warn!("not logged in"); + // return Err(CommonWebError::NotAuthorized); + // } + //}; + + if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) { + return Err(CommonWebError::BadInputWithSimpleMessage(reason)); + } + + let mut query_media_tokens = None; + + if let Some(media_tokens) = request.image_media_tokens.as_ref() { + if !media_tokens.is_empty() { + let mut tokens = Vec::new(); + tokens.extend(media_tokens.iter().map(|token| token.to_owned())); + + if let Some(mask_media_token) = request.mask_image_token.as_ref() { + // NB: Masks are only relevant when this is an image editing exercise. + tokens.push(mask_media_token.to_owned()); + } + + query_media_tokens = Some(tokens); + } + } + + let image_urls_by_token = match query_media_tokens { + None => HashMap::new(), + Some(media_tokens) => { + info!("Looking up image media tokens: {:?}", media_tokens); + lookup_image_urls_as_map( + &http_request, + &mut mysql_connection, + server_state.server_environment, + &media_tokens, + ).await? + } + }; + + let maybe_image_urls = match request.image_media_tokens.as_ref() { + None => None, + Some(tokens) => { + let urls = tokens.iter() + .filter_map(|token| image_urls_by_token.get(token)) + .map(|url| url.to_string()) + .collect::>(); + Some(urls) + } + }; + + let maybe_mask_url = match request.mask_image_token.as_ref() { + None => None, + Some(token) => { + image_urls_by_token.get(&token).map(|url| url.to_string()) + } + }; + + insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection) + .await + .map_err(|err| { + error!("Error inserting idempotency token: {:?}", err); + CommonWebError::BadInputWithSimpleMessage("repeated idempotency token".to_string()) + })?; + + info!("Fal webhook URL: {}", server_state.fal.webhook_url); + + let fal_result; + + if let Some(input_image_urls) = maybe_image_urls.as_deref() { + info!("gpt image 1.5 edit image"); + + let num_images = match request.num_images { + Some(GptImage1p5MultiFunctionImageGenNumImages::One) => EnqueueGptImage1p5EditImageNumImages::One, + Some(GptImage1p5MultiFunctionImageGenNumImages::Two) => EnqueueGptImage1p5EditImageNumImages::Two, + Some(GptImage1p5MultiFunctionImageGenNumImages::Three) => EnqueueGptImage1p5EditImageNumImages::Three, + Some(GptImage1p5MultiFunctionImageGenNumImages::Four) => EnqueueGptImage1p5EditImageNumImages::Four, + None => EnqueueGptImage1p5EditImageNumImages::One, // Default to One + }; + + let image_size = match request.image_size { + Some(GptImage1p5MultiFunctionImageGenSize::Square) => EnqueueGptImage1p5EditImageSize::Square, + Some(GptImage1p5MultiFunctionImageGenSize::Wide) => EnqueueGptImage1p5EditImageSize::Wide, + Some(GptImage1p5MultiFunctionImageGenSize::Tall) => EnqueueGptImage1p5EditImageSize::Tall, + None => EnqueueGptImage1p5EditImageSize::Square, + }; + + let background = match request.background { + Some(GptImage1p5MultiFunctionImageGenBackground::Auto) => EnqueueGptImage1p5EditImageBackground::Auto, + Some(GptImage1p5MultiFunctionImageGenBackground::Transparent) => EnqueueGptImage1p5EditImageBackground::Transparent, + Some(GptImage1p5MultiFunctionImageGenBackground::Opaque) => EnqueueGptImage1p5EditImageBackground::Opaque, + None => EnqueueGptImage1p5EditImageBackground::Auto, + }; + + let quality = match request.quality { + Some(GptImage1p5MultiFunctionImageGenQuality::Low) => EnqueueGptImage1p5EditImageQuality::Low, + Some(GptImage1p5MultiFunctionImageGenQuality::Medium) => EnqueueGptImage1p5EditImageQuality::Medium, + Some(GptImage1p5MultiFunctionImageGenQuality::High) => EnqueueGptImage1p5EditImageQuality::High, + None => EnqueueGptImage1p5EditImageQuality::High, + }; + + let input_fidelity = match request.input_fidelity { + Some(GptImage1p5MultiFunctionImageGenInputFidelity::Low) => EnqueueGptImage1p5EditImageInputFidelity::Low, + Some(GptImage1p5MultiFunctionImageGenInputFidelity::High) => EnqueueGptImage1p5EditImageInputFidelity::High, + None => EnqueueGptImage1p5EditImageInputFidelity::High, + }; + + let args = EnqueueGptImage1p5EditImageArgs { + prompt: request.prompt.as_deref().unwrap_or(""), + image_urls: input_image_urls.to_owned(), + num_images, + mask_image_url: maybe_mask_url, + image_size: Some(image_size), + background: Some(background), + quality: Some(quality), + input_fidelity: Some(input_fidelity), + output_format: None, + webhook_url: &server_state.fal.webhook_url, + api_key: &server_state.fal.api_key, + }; + + fal_result = enqueue_gpt_image_1p5_image_edit_webhook(args) + .await + .map_err(|err| { + warn!("Error calling enqueue_gpt_image_1p5_image_edit_webhook: {:?}", err); + CommonWebError::ServerError + })?; + + } else { + info!("gpt image 1.5 text-to-image"); + + let num_images = match request.num_images { + Some(GptImage1p5MultiFunctionImageGenNumImages::One) => EnqueueGptImage1p5TextToImageNumImages::One, + Some(GptImage1p5MultiFunctionImageGenNumImages::Two) => EnqueueGptImage1p5TextToImageNumImages::Two, + Some(GptImage1p5MultiFunctionImageGenNumImages::Three) => EnqueueGptImage1p5TextToImageNumImages::Three, + Some(GptImage1p5MultiFunctionImageGenNumImages::Four) => EnqueueGptImage1p5TextToImageNumImages::Four, + None => EnqueueGptImage1p5TextToImageNumImages::One, // Default to One + }; + + let image_size = match request.image_size { + Some(GptImage1p5MultiFunctionImageGenSize::Square) => EnqueueGptImage1p5TextToImageSize::Square, + Some(GptImage1p5MultiFunctionImageGenSize::Wide) => EnqueueGptImage1p5TextToImageSize::Wide, + Some(GptImage1p5MultiFunctionImageGenSize::Tall) => EnqueueGptImage1p5TextToImageSize::Tall, + None => EnqueueGptImage1p5TextToImageSize::Square, + }; + + let background = match request.background { + Some(GptImage1p5MultiFunctionImageGenBackground::Auto) => EnqueueGptImage1p5TextToImageBackground::Auto, + Some(GptImage1p5MultiFunctionImageGenBackground::Transparent) => EnqueueGptImage1p5TextToImageBackground::Transparent, + Some(GptImage1p5MultiFunctionImageGenBackground::Opaque) => EnqueueGptImage1p5TextToImageBackground::Opaque, + None => EnqueueGptImage1p5TextToImageBackground::Auto, + }; + + let quality = match request.quality { + Some(GptImage1p5MultiFunctionImageGenQuality::Low) => EnqueueGptImage1p5TextToImageQuality::Low, + Some(GptImage1p5MultiFunctionImageGenQuality::Medium) => EnqueueGptImage1p5TextToImageQuality::Medium, + Some(GptImage1p5MultiFunctionImageGenQuality::High) => EnqueueGptImage1p5TextToImageQuality::High, + None => EnqueueGptImage1p5TextToImageQuality::High, + }; + + let args = EnqueueGptImage1p5TextToImageArgs { + prompt: request.prompt.as_deref().unwrap_or(""), + num_images, + output_format: None, + image_size: Some(image_size), + background: Some(background), + quality: Some(quality), + webhook_url: &server_state.fal.webhook_url, + api_key: &server_state.fal.api_key, + }; + + fal_result = enqueue_gpt_image_1p5_text_to_image_webhook(args) + .await + .map_err(|err| { + warn!("Error calling enqueue_gpt_image_1p5_text_to_image_webhook: {:?}", err); + CommonWebError::ServerError + })?; + } + + let external_job_id = fal_result.request_id + .ok_or_else(|| { + warn!("Fal request_id is None"); + CommonWebError::ServerError + })?; + + info!("Fal request_id: {}", external_job_id); + + let ip_address = get_request_ip(&http_request); + + let mut transaction = mysql_connection + .begin() + .await + .map_err(|err| { + error!("Error starting MySQL transaction: {:?}", err); + CommonWebError::ServerError + })?; + + // NB: Don't fail the job if the query fails. + let prompt_result = insert_prompt(InsertPromptArgs { + maybe_apriori_prompt_token: None, + prompt_type: PromptType::ArtcraftApp, + maybe_creator_user_token: maybe_user_session + .as_ref() + .map(|s| &s.user_token), + maybe_model_type: Some(ModelType::GptImage1p5), + maybe_generation_provider: Some(GenerationProvider::Artcraft), + maybe_positive_prompt: request.prompt.as_deref(), + maybe_negative_prompt: None, + maybe_other_args: None, + creator_ip_address: &ip_address, + mysql_executor: &mut *transaction, + phantom: Default::default(), + }).await; + + let prompt_token = match prompt_result { + Ok(token) => Some(token), + Err(err) => { + warn!("Error inserting prompt: {:?}", err); + None // Don't fail the job if the prompt insertion fails. + } + }; + + if let Some(media_tokens) = &request.image_media_tokens { + if let Some(token) = prompt_token.as_ref() { + let result = insert_batch_prompt_context_items(InsertBatchArgs { + prompt_token: token.clone(), + items: media_tokens.iter().map(|token| { + PromptContextItem { + media_token: token.clone(), + context_semantic_type: PromptContextSemanticType::Imgref, + } + }).collect(), + transaction: &mut transaction, + }).await; + + if let Err(err) = result { + // NB: Fail open. + warn!("Error inserting batch prompt context items: {:?}", err); + } + } + } + + let db_result = insert_generic_inference_job_for_fal_queue(InsertGenericInferenceForFalArgs { + uuid_idempotency_token: &request.uuid_idempotency_token, + maybe_external_third_party_id: &external_job_id, + fal_category: FalCategory::ImageGeneration, + maybe_inference_args: None, + maybe_prompt_token: prompt_token.as_ref(), + maybe_creator_user_token: maybe_user_session.as_ref().map(|s| &s.user_token), + maybe_avt_token: maybe_avt_token.as_ref(), + creator_ip_address: &ip_address, + creator_set_visibility: Visibility::Public, + mysql_executor: &mut *transaction, + phantom: Default::default(), + }).await; + + let job_token = match db_result { + Ok(token) => token, + Err(err) => { + warn!("Error inserting generic inference job for FAL queue: {:?}", err); + return Err(CommonWebError::ServerError); + } + }; + + let _r = transaction + .commit() + .await + .map_err(|err| { + error!("Error committing MySQL transaction: {:?}", err); + CommonWebError::ServerError + })?; + + Ok(Json(GptImage1p5MultiFunctionImageGenResponse { + success: true, + inference_job_token: job_token, + })) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/mod.rs index cebfe76101..0d025caa70 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/mod.rs @@ -1 +1,2 @@ +pub mod gpt_image_1p5_multi_function_image_gen_handler; pub mod nano_banana_pro_multi_function_image_gen_handler; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/nano_banana_pro_multi_function_image_gen_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/nano_banana_pro_multi_function_image_gen_handler.rs index 09b113b14c..99b952438f 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/nano_banana_pro_multi_function_image_gen_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/nano_banana_pro_multi_function_image_gen_handler.rs @@ -8,6 +8,7 @@ use crate::http_server::endpoints::generate::common::payments_error_test::paymen use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain; use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format; use crate::state::server_state::ServerState; +use crate::util::lookup::lookup_image_urls_as_optional_list::lookup_image_urls_as_optional_list; use actix_web::web::Json; use actix_web::{web, HttpRequest}; use artcraft_api_defs::generate::image::multi_function::nano_banana_pro_multi_function_image_gen::{NanoBananaProMultiFunctionImageGenAspectRatio, NanoBananaProMultiFunctionImageGenImageResolution, NanoBananaProMultiFunctionImageGenNumImages, NanoBananaProMultiFunctionImageGenRequest, NanoBananaProMultiFunctionImageGenResponse}; @@ -89,7 +90,7 @@ pub async fn nano_banana_pro_multi_function_image_gen_handler( let image_urls = match request.image_media_tokens.as_ref() { Some(media_tokens) => { info!("Looking up image media tokens: {:?}", media_tokens); - lookup_image_urls( + lookup_image_urls_as_optional_list( &http_request, &mut mysql_connection, server_state.server_environment, @@ -251,7 +252,6 @@ pub async fn nano_banana_pro_multi_function_image_gen_handler( } }; - if let Some(media_tokens) = &request.image_media_tokens { if let Some(token) = prompt_token.as_ref() { let result = insert_batch_prompt_context_items(InsertBatchArgs { @@ -307,66 +307,3 @@ pub async fn nano_banana_pro_multi_function_image_gen_handler( inference_job_token: job_token, })) } - - -async fn lookup_image_urls( - http_request: &HttpRequest, - mysql_connection: &mut PoolConnection, - server_environment: ServerEnvironment, - tokens: &[MediaFileToken], -) -> Result>, CommonWebError> { - const CAN_SEE_DELETED: bool = false; - - let result = batch_get_media_files_by_tokens_with_connection( - mysql_connection, - tokens, - CAN_SEE_DELETED, - ).await; - - let media_files = match result { - Ok(files) => files, - Err(err) => { - error!("Error getting media files by tokens: {:?}", err); - return Err(CommonWebError::ServerError); - } - }; - - if media_files.len() != tokens.len() { - warn!("Wrong number of media files returned for tokens: {} found for {} tokens", media_files.len(), tokens.len()); - - let requested : HashSet<&MediaFileToken> = HashSet::from_iter(tokens.iter()); - let returned : HashSet<&MediaFileToken> = HashSet::from_iter(media_files.iter().map(|m| &m.token)); - - let diff = requested.difference(&returned) - .cloned() - .collect::>(); - - return Err(CommonWebError::BadInputWithSimpleMessage( - format!("Not all media files could be found. Media files found: {}, tokens provided: {}, in original: {:?}, req {:?}, ret {:?}", - media_files.len(), tokens.len(), diff, requested, returned))); - } - - let media_domain = get_media_domain(&http_request); - - let image_urls = media_files.iter() - .map(|file| { - let public_bucket_path = MediaFileBucketPath::from_object_hash( - &file.public_bucket_directory_hash, - file.maybe_public_bucket_prefix.as_deref(), - file.maybe_public_bucket_extension.as_deref()); - - let media_links = MediaLinksBuilder::from_media_path_and_env( - media_domain, - server_environment, - &public_bucket_path); - - media_links.cdn_url.to_string() - }) - .collect::>(); - - if image_urls.is_empty() { - return Ok(None) - } - - Ok(Some(image_urls)) -} diff --git a/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_map.rs b/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_map.rs new file mode 100644 index 0000000000..1a52a3b64e --- /dev/null +++ b/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_map.rs @@ -0,0 +1,71 @@ +use crate::http_server::common_responses::common_web_error::CommonWebError; +use crate::http_server::common_responses::media::media_links_builder::MediaLinksBuilder; +use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain; +use actix_web::HttpRequest; +use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath; +use log::{error, warn}; +use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::batch_get_media_files_by_tokens_with_connection; +use server_environment::ServerEnvironment; +use sqlx::pool::PoolConnection; +use sqlx::MySql; +use std::collections::{HashMap, HashSet}; +use std::iter::FromIterator; +use tokens::tokens::media_files::MediaFileToken; + +pub async fn lookup_image_urls_as_map( + http_request: &HttpRequest, + mysql_connection: &mut PoolConnection, + server_environment: ServerEnvironment, + tokens: &[MediaFileToken], +) -> Result, CommonWebError> { + const CAN_SEE_DELETED: bool = false; + + let result = batch_get_media_files_by_tokens_with_connection( + mysql_connection, + tokens, + CAN_SEE_DELETED, + ).await; + + let media_files = match result { + Ok(files) => files, + Err(err) => { + error!("Error getting media files by tokens: {:?}", err); + return Err(CommonWebError::ServerError); + } + }; + + if media_files.len() != tokens.len() { + warn!("Wrong number of media files returned for tokens: {} found for {} tokens", media_files.len(), tokens.len()); + + let requested : HashSet<&MediaFileToken> = HashSet::from_iter(tokens.iter()); + let returned : HashSet<&MediaFileToken> = HashSet::from_iter(media_files.iter().map(|m| &m.token)); + + let diff = requested.difference(&returned) + .cloned() + .collect::>(); + + return Err(CommonWebError::BadInputWithSimpleMessage( + format!("Not all media files could be found. Media files found: {}, tokens provided: {}, in original: {:?}, req {:?}, ret {:?}", + media_files.len(), tokens.len(), diff, requested, returned))); + } + + let media_domain = get_media_domain(&http_request); + + let image_urls = media_files.into_iter() + .map(|file| { + let public_bucket_path = MediaFileBucketPath::from_object_hash( + &file.public_bucket_directory_hash, + file.maybe_public_bucket_prefix.as_deref(), + file.maybe_public_bucket_extension.as_deref()); + + let media_links = MediaLinksBuilder::from_media_path_and_env( + media_domain, + server_environment, + &public_bucket_path); + + (file.token, media_links.cdn_url.to_string()) + }) + .collect::>(); + + Ok(image_urls) +} diff --git a/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_optional_list.rs b/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_optional_list.rs new file mode 100644 index 0000000000..243a782309 --- /dev/null +++ b/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_optional_list.rs @@ -0,0 +1,75 @@ +use crate::http_server::common_responses::common_web_error::CommonWebError; +use crate::http_server::common_responses::media::media_links_builder::MediaLinksBuilder; +use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain; +use actix_web::HttpRequest; +use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath; +use log::{error, warn}; +use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::batch_get_media_files_by_tokens_with_connection; +use server_environment::ServerEnvironment; +use sqlx::pool::PoolConnection; +use sqlx::MySql; +use std::collections::HashSet; +use std::iter::FromIterator; +use tokens::tokens::media_files::MediaFileToken; + +pub async fn lookup_image_urls_as_optional_list( + http_request: &HttpRequest, + mysql_connection: &mut PoolConnection, + server_environment: ServerEnvironment, + tokens: &[MediaFileToken], +) -> Result>, CommonWebError> { + const CAN_SEE_DELETED: bool = false; + + let result = batch_get_media_files_by_tokens_with_connection( + mysql_connection, + tokens, + CAN_SEE_DELETED, + ).await; + + let media_files = match result { + Ok(files) => files, + Err(err) => { + error!("Error getting media files by tokens: {:?}", err); + return Err(CommonWebError::ServerError); + } + }; + + if media_files.len() != tokens.len() { + warn!("Wrong number of media files returned for tokens: {} found for {} tokens", media_files.len(), tokens.len()); + + let requested : HashSet<&MediaFileToken> = HashSet::from_iter(tokens.iter()); + let returned : HashSet<&MediaFileToken> = HashSet::from_iter(media_files.iter().map(|m| &m.token)); + + let diff = requested.difference(&returned) + .cloned() + .collect::>(); + + return Err(CommonWebError::BadInputWithSimpleMessage( + format!("Not all media files could be found. Media files found: {}, tokens provided: {}, in original: {:?}, req {:?}, ret {:?}", + media_files.len(), tokens.len(), diff, requested, returned))); + } + + let media_domain = get_media_domain(&http_request); + + let image_urls = media_files.iter() + .map(|file| { + let public_bucket_path = MediaFileBucketPath::from_object_hash( + &file.public_bucket_directory_hash, + file.maybe_public_bucket_prefix.as_deref(), + file.maybe_public_bucket_extension.as_deref()); + + let media_links = MediaLinksBuilder::from_media_path_and_env( + media_domain, + server_environment, + &public_bucket_path); + + media_links.cdn_url.to_string() + }) + .collect::>(); + + if image_urls.is_empty() { + return Ok(None) + } + + Ok(Some(image_urls)) +} diff --git a/crates/service/web/storyteller_web/src/util/lookup/mod.rs b/crates/service/web/storyteller_web/src/util/lookup/mod.rs index e31f48b51a..5b31bab895 100644 --- a/crates/service/web/storyteller_web/src/util/lookup/mod.rs +++ b/crates/service/web/storyteller_web/src/util/lookup/mod.rs @@ -1 +1,3 @@ -pub mod fetch_all_required_media_files; \ No newline at end of file +pub mod fetch_all_required_media_files; +pub mod lookup_image_urls_as_optional_list; +pub mod lookup_image_urls_as_map; \ No newline at end of file