From ff6ae55a4ea2c1be09f2f5e9c74ce2951243ce95 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 20 Aug 2025 17:57:16 -0400 Subject: [PATCH] insert batch tokens on batch generation --- ...ae00fc4e085d68b064e370f0e3d92b82049b2.json | 12 +++++++ ...6c96299001e143c745ee30df99b72d1235d35.json | 12 ------- ...generic_inference_job_successfully_done.rs | 30 ++++++++++------ .../media_file_insert_builder.rs | 13 +++++-- .../endpoints/webhooks/fal_webhook_handler.rs | 19 +++++----- .../webhooks/handle_images_payload.rs | 36 ++++++++++++++----- 6 files changed, 81 insertions(+), 41 deletions(-) create mode 100644 .sqlx/query-7e223e284b992b6eaf31bec057bae00fc4e085d68b064e370f0e3d92b82049b2.json delete mode 100644 .sqlx/query-f68a8917e639418a0bb952bef606c96299001e143c745ee30df99b72d1235d35.json diff --git a/.sqlx/query-7e223e284b992b6eaf31bec057bae00fc4e085d68b064e370f0e3d92b82049b2.json b/.sqlx/query-7e223e284b992b6eaf31bec057bae00fc4e085d68b064e370f0e3d92b82049b2.json new file mode 100644 index 0000000000..4266e89a6c --- /dev/null +++ b/.sqlx/query-7e223e284b992b6eaf31bec057bae00fc4e085d68b064e370f0e3d92b82049b2.json @@ -0,0 +1,12 @@ +{ + "db_name": "MySQL", + "query": "\nUPDATE generic_inference_jobs\nSET\n status = ?,\n on_success_result_entity_type = ?,\n on_success_result_entity_token = ?,\n on_success_result_batch_token = ?,\n failure_reason = NULL,\n internal_debugging_failure_reason = NULL,\n success_execution_millis = NULL,\n success_inference_execution_millis = NULL,\n retry_at = NULL,\n successfully_completed_at = NOW()\nWHERE token = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 5 + }, + "nullable": [] + }, + "hash": "7e223e284b992b6eaf31bec057bae00fc4e085d68b064e370f0e3d92b82049b2" +} diff --git a/.sqlx/query-f68a8917e639418a0bb952bef606c96299001e143c745ee30df99b72d1235d35.json b/.sqlx/query-f68a8917e639418a0bb952bef606c96299001e143c745ee30df99b72d1235d35.json deleted file mode 100644 index 75f2574c3f..0000000000 --- a/.sqlx/query-f68a8917e639418a0bb952bef606c96299001e143c745ee30df99b72d1235d35.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "db_name": "MySQL", - "query": "\nUPDATE generic_inference_jobs\nSET\n status = ?,\n on_success_result_entity_type = ?,\n on_success_result_entity_token = ?,\n failure_reason = NULL,\n internal_debugging_failure_reason = NULL,\n success_execution_millis = NULL,\n success_inference_execution_millis = NULL,\n retry_at = NULL,\n successfully_completed_at = NOW()\nWHERE token = ?\n ", - "describe": { - "columns": [], - "parameters": { - "Right": 4 - }, - "nullable": [] - }, - "hash": "f68a8917e639418a0bb952bef606c96299001e143c745ee30df99b72d1235d35" -} diff --git a/crates/schema/database/mysql_queries/src/queries/generic_inference/fal/mark_fal_generic_inference_job_successfully_done.rs b/crates/schema/database/mysql_queries/src/queries/generic_inference/fal/mark_fal_generic_inference_job_successfully_done.rs index 835f94853c..7a522e3102 100644 --- a/crates/schema/database/mysql_queries/src/queries/generic_inference/fal/mark_fal_generic_inference_job_successfully_done.rs +++ b/crates/schema/database/mysql_queries/src/queries/generic_inference/fal/mark_fal_generic_inference_job_successfully_done.rs @@ -1,19 +1,27 @@ use anyhow::anyhow; -use sqlx::MySqlPool; +use sqlx::{Executor, MySql}; +use std::marker::PhantomData; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; use enums::common::job_status_plus::JobStatusPlus; use errors::AnyhowResult; +use tokens::tokens::batch_generations::BatchGenerationToken; use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::media_files::MediaFileToken; -// TODO: Maybe in the future we set a batch_token instead of media token - -pub async fn mark_fal_generic_inference_job_successfully_done( - pool: &MySqlPool, - job_token: &InferenceJobToken, - media_file_token: MediaFileToken, +pub struct MarkJobArgs<'e, 'c, E> + where E: 'e + Executor<'c, Database = MySql> +{ + pub job_token: &'e InferenceJobToken, + pub media_file_token: &'e MediaFileToken, + pub maybe_batch_token: Option<&'e BatchGenerationToken>, + pub mysql_executor: E, + pub phantom: PhantomData<&'c E>, +} +pub async fn mark_fal_generic_inference_job_successfully_done<'e, 'c : 'e, E>( + args: MarkJobArgs<'e, 'c, E> ) -> AnyhowResult<()> + where E: 'e + Executor<'c, Database = MySql> { const STATUS : JobStatusPlus = JobStatusPlus::CompleteSuccess; const RESULT_TYPE : InferenceResultType = InferenceResultType::MediaFile; @@ -25,6 +33,7 @@ SET status = ?, on_success_result_entity_type = ?, on_success_result_entity_token = ?, + on_success_result_batch_token = ?, failure_reason = NULL, internal_debugging_failure_reason = NULL, success_execution_millis = NULL, @@ -35,10 +44,11 @@ WHERE token = ? "#, STATUS.to_str(), RESULT_TYPE.to_str(), - media_file_token, - job_token, + args.media_file_token, + args.maybe_batch_token.map(|t| t.as_str()), + args.job_token, ) - .execute(pool) + .execute(args.mysql_executor) .await; match query_result { diff --git a/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_builder/media_file_insert_builder.rs b/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_builder/media_file_insert_builder.rs index 98b4cc4a53..43de51f58d 100644 --- a/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_builder/media_file_insert_builder.rs +++ b/crates/schema/database/mysql_queries/src/queries/media_files/create/insert_builder/media_file_insert_builder.rs @@ -10,6 +10,7 @@ use enums::by_table::media_files::media_file_type::MediaFileType; use enums::common::visibility::Visibility; use sqlx::MySqlPool; use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken; +use tokens::tokens::batch_generations::BatchGenerationToken; use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::prompts::PromptToken; use tokens::tokens::users::UserToken; @@ -67,7 +68,7 @@ pub struct MediaFileInsertBuilder { // maybe_prompt_token: Option<&'a PromptToken>, // // If batch generated, this is the batch token. - // maybe_batch_token: Option<&'a BatchGenerationToken>, + maybe_batch_generation_token: Option, // Storage details public_bucket_directory_hash: Option, // NB: Non-nullable field(s) @@ -100,6 +101,7 @@ impl MediaFileInsertBuilder { maybe_frame_height: None, checksum_sha2: None, maybe_engine_category: None, + maybe_batch_generation_token: None, public_bucket_directory_hash: None, maybe_prompt_token: None, } @@ -228,7 +230,12 @@ impl MediaFileInsertBuilder { // TODO: maybe_text_transcript // TODO: maybe_scene_source_media_file_token // TODO: maybe_prompt_token - // TODO: maybe_batch_token + + pub fn maybe_batch_generation_token(mut self, maybe_batch_token: Option<&BatchGenerationToken>) -> Self { + self.maybe_batch_generation_token = maybe_batch_token + .map(|token| token.clone()); + self + } pub fn public_bucket_directory_hash(mut self, public_bucket_directory_hash: &MediaFileBucketPath) -> Self { self.public_bucket_directory_hash = Some(public_bucket_directory_hash.clone()); @@ -281,7 +288,7 @@ impl MediaFileInsertBuilder { maybe_text_transcript: None, // TODO maybe_scene_source_media_file_token: None, // TODO maybe_prompt_token: self.maybe_prompt_token.as_ref(), - maybe_batch_token: None, // TODO + maybe_batch_token: self.maybe_batch_generation_token.as_ref(), public_bucket_directory_hash: bucket_path.get_object_hash(), maybe_public_bucket_prefix: bucket_path.get_optional_prefix(), maybe_public_bucket_extension: bucket_path.get_optional_extension(), diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/webhooks/fal_webhook_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/webhooks/fal_webhook_handler.rs index 0ed6ddfe07..44a2cda25d 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/webhooks/fal_webhook_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/webhooks/fal_webhook_handler.rs @@ -15,7 +15,7 @@ use http_server_common::response::response_success_helpers::SimpleGenericJsonSuc use http_server_common::response::serialize_as_json_error::serialize_as_json_error; use log::{info, warn}; use mysql_queries::queries::generic_inference::fal::get_inference_job_by_fal_id::get_inference_job_by_fal_id; -use mysql_queries::queries::generic_inference::fal::mark_fal_generic_inference_job_successfully_done::mark_fal_generic_inference_job_successfully_done; +use mysql_queries::queries::generic_inference::fal::mark_fal_generic_inference_job_successfully_done::{mark_fal_generic_inference_job_successfully_done, MarkJobArgs}; use serde_json::Value; use utoipa::ToSchema; // 1. tauri --> hit endpoint to enqueue @@ -95,7 +95,7 @@ impl From for FalWebhookError { /// Fal webhook #[utoipa::path( post, - tag = "Webooks", + tag = "Webhooks", path = "/v1/webhooks/fal", responses( (status = 200, description = "Success", body = SimpleGenericJsonSuccess), @@ -143,6 +143,7 @@ pub async fn fal_webhook_handler( info!("Fal webhook job record: {:?}", job); let mut maybe_media_token = None; + let mut maybe_batch_token = None; if let Some(payload_obj) = payload.as_object() { if payload_obj.contains_key("image") { @@ -150,7 +151,7 @@ pub async fn fal_webhook_handler( let token = handle_image_payload(payload_obj, &job, &server_state).await?; maybe_media_token = Some(token); } else if payload_obj.contains_key("images") { - maybe_media_token = handle_images_payload(payload_obj, &job, &server_state).await?; + (maybe_media_token, maybe_batch_token) = handle_images_payload(payload_obj, &job, &server_state).await?; } else if payload_obj.contains_key("video") { info!("Handling video payload for job: {:?}", job.job_token); let token = handle_video_payload(payload_obj, &job, &server_state).await?; @@ -165,11 +166,13 @@ pub async fn fal_webhook_handler( if let Some(media_token) = maybe_media_token { info!("Media file token: {:?}", media_token); // TODO: Update job metadata. - mark_fal_generic_inference_job_successfully_done( - &server_state.mysql_pool, - &job.job_token, - media_token, - ).await.map_err(|err| { + mark_fal_generic_inference_job_successfully_done(MarkJobArgs { + job_token: &job.job_token, + media_file_token: &media_token, + maybe_batch_token: maybe_batch_token.as_ref(), + mysql_executor: &server_state.mysql_pool, + phantom: Default::default(), + }).await.map_err(|err| { warn!("Error marking job as successfully done: {:?}", err); FalWebhookError::ServerError })?; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/webhooks/handle_images_payload.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/webhooks/handle_images_payload.rs index 0aa57c1c80..e338803b6d 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/webhooks/handle_images_payload.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/webhooks/handle_images_payload.rs @@ -15,6 +15,7 @@ use mimetypes::mimetype_info::mimetype_info::MimetypeInfo; use mysql_queries::queries::generic_inference::fal::get_inference_job_by_fal_id::FalJobDetails; use mysql_queries::queries::media_files::create::insert_builder::media_file_insert_builder::MediaFileInsertBuilder; use serde_json::{Map, Value}; +use tokens::tokens::batch_generations::BatchGenerationToken; use tokens::tokens::media_files::MediaFileToken; const PREFIX : Option<&str> = Some("artcraft_"); @@ -33,8 +34,8 @@ pub async fn handle_images_payload( payload: &Map, job: &FalJobDetails, server_state: &ServerState, -) -> AnyhowResult> { - +) -> AnyhowResult<(Option, Option)> { + let images_value = payload.get("images") .ok_or_else(|| anyhow!("no `images` key in payload"))?; @@ -43,22 +44,31 @@ pub async fn handle_images_payload( let images: Vec = serde_json::from_value(images_value.clone())?; let mut maybe_media_token = None; - + + // NB: We are not going to create `batch_generations` table records. We don't need them. + // The foreign key in `media_files` is enough to look up the rest of the batch. + let mut maybe_batch_token = None; + + if images.len() > 1 { + maybe_batch_token = Some(BatchGenerationToken::generate()); + } + for image in images { - let media_token = upload_image(job, server_state, &image).await?; + let media_token = upload_image(job, server_state, &image, maybe_batch_token.as_ref()).await?; if maybe_media_token.is_none() { maybe_media_token = Some(media_token); // Set the first media token } } - Ok(maybe_media_token) + Ok((maybe_media_token, maybe_batch_token)) } async fn upload_image( job: &FalJobDetails, server_state: &ServerState, image: &FalWebhookImage, + maybe_batch_token: Option<&BatchGenerationToken>, ) -> AnyhowResult { let image_url = image.url .as_deref() @@ -98,6 +108,7 @@ async fn upload_image( image, &file_bytes, mimetype_info, + maybe_batch_token, ).await } _ => { @@ -107,6 +118,7 @@ async fn upload_image( image, &file_bytes, mimetype_info, + maybe_batch_token, ).await } } @@ -118,6 +130,7 @@ async fn upload_image_bytes( image: &FalWebhookImage, file_bytes: &[u8], mimetype_info: MimetypeInfo, + maybe_batch_token: Option<&BatchGenerationToken>, ) -> AnyhowResult { let mime_type = mimetype_info.mime_type(); @@ -145,7 +158,7 @@ async fn upload_image_bytes( &mime_type) .await?; - let media_token = MediaFileInsertBuilder::new() + let mut query_builder = MediaFileInsertBuilder::new() .maybe_creator_user(job.maybe_creator_user_token.as_ref()) .maybe_creator_anonymous_visitor(job.maybe_creator_anonymous_visitor_token.as_ref()) .creator_ip_address(&job.creator_ip_address) @@ -159,11 +172,18 @@ async fn upload_image_bytes( .frame_width(image_info.width()) .frame_height(image_info.height()) .maybe_prompt_token(job.maybe_prompt_token.as_ref()) - .checksum_sha2(&file_hash) + .checksum_sha2(&file_hash); + + if let Some(batch_token) = maybe_batch_token { + query_builder = query_builder.maybe_batch_generation_token(Some(batch_token)); + } + + let media_token = query_builder .insert_pool(&server_state.mysql_pool) .await?; - info!("Image media file uploaded with token: {}", media_token); + info!("Image media file uploaded with token: {} ; possible batch token: {:?}", + media_token, maybe_batch_token); Ok(media_token) }