gpt image 1.5 handler plus helpers

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<String>,
+  pub mask_image_token: Option<MediaFileToken>,
 
   /// Number of images to generate. Default is one.
   pub num_images: Option<GptImage1p5MultiFunctionImageGenNumImages>,

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));

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<GptImage1p5MultiFunctionImageGenRequest>,
+  server_state: web::Data<Arc<ServerState>>
+) -> Result<Json<GptImage1p5MultiFunctionImageGenResponse>, 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::<Vec<_>>();
+      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,
+  }))
+}

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;

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<MySql>,
-  server_environment: ServerEnvironment,
-  tokens: &[MediaFileToken],
-) -> Result<Option<Vec<String>>, 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::<Vec<&MediaFileToken>>();
-
-    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::<Vec<_>>();
-
-  if image_urls.is_empty() {
-    return Ok(None)
-  }
-  
-  Ok(Some(image_urls))
-}

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<MySql>,
+  server_environment: ServerEnvironment,
+  tokens: &[MediaFileToken],
+) -> Result<HashMap<MediaFileToken, String>, 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::<Vec<&MediaFileToken>>();
+
+    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::<HashMap<_, _>>();
+
+  Ok(image_urls)
+}

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<MySql>,
+  server_environment: ServerEnvironment,
+  tokens: &[MediaFileToken],
+) -> Result<Option<Vec<String>>, 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::<Vec<&MediaFileToken>>();
+
+    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::<Vec<_>>();
+
+  if image_urls.is_empty() {
+    return Ok(None)
+  }
+
+  Ok(Some(image_urls))
+}

crates/service/web/storyteller_web/src/util/lookup/mod.rs

@@ -1 +1,3 @@
-pub mod fetch_all_required_media_files;
+pub mod fetch_all_required_media_files;
+pub mod lookup_image_urls_as_optional_list;
+pub mod lookup_image_urls_as_map;