From dd6e8d00a293adcce6c34f57a040e111655f3c77 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 30 Apr 2026 18:32:07 -0400 Subject: [PATCH] batch get prompts endpoint --- .../src/prompts/batch_get_prompts.rs | 68 ++++++++ .../artcraft_api_defs/src/prompts/mod.rs | 1 + .../src/queries/prompts/batch_get_prompts.rs | 92 ++++++++++ .../mysql_queries/src/queries/prompts/mod.rs | 1 + .../web/storyteller_web/src/docs/api_doc.rs | 7 + .../prompts/batch_get_prompts_handler.rs | 159 ++++++++++++++++++ .../src/http_server/endpoints/prompts/mod.rs | 1 + .../application_routes/prompts_routes.rs | 2 + 8 files changed, 331 insertions(+) create mode 100644 crates/api_clients/artcraft/artcraft_api_defs/src/prompts/batch_get_prompts.rs create mode 100644 crates/schema/database/mysql_queries/src/queries/prompts/batch_get_prompts.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/prompts/batch_get_prompts_handler.rs diff --git a/crates/api_clients/artcraft/artcraft_api_defs/src/prompts/batch_get_prompts.rs b/crates/api_clients/artcraft/artcraft_api_defs/src/prompts/batch_get_prompts.rs new file mode 100644 index 0000000000..e02f3cc739 --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_api_defs/src/prompts/batch_get_prompts.rs @@ -0,0 +1,68 @@ +use crate::prompts::get_prompt::GetPromptImageContextItem; +use chrono::{DateTime, Utc}; +use enums::by_table::prompts::prompt_type::PromptType; +use enums::common::generation::common_aspect_ratio::CommonAspectRatio; +use enums::common::generation::common_generation_mode::CommonGenerationMode; +use enums::common::generation::common_model_class::CommonModelClass; +use enums::common::generation::common_model_type::CommonModelType; +use enums::common::generation::common_resolution::CommonResolution; +use enums::common::generation_provider::GenerationProvider; +use serde_derive::{Deserialize, Serialize}; +use tokens::tokens::prompts::PromptToken; +use utoipa::ToSchema; + +pub const BATCH_GET_PROMPTS_PATH: &str = "/v1/prompts/batch"; + +#[derive(Deserialize, ToSchema)] +pub struct BatchGetPromptsRequest { + pub tokens: Vec, +} + +#[derive(Serialize, ToSchema)] +pub struct BatchGetPromptsResponse { + pub success: bool, + pub prompts: Vec, +} + +#[derive(Serialize, ToSchema)] +pub struct BatchPromptInfo { + pub token: PromptToken, + + /// The type of model used + pub maybe_model_type: Option, + + /// The class of model used + pub maybe_model_class: Option, + + /// The service provider used + pub maybe_generation_provider: Option, + + /// Positive prompt (technically optional, but usually present) + pub maybe_positive_prompt: Option, + + /// Negative prompt (optional) + pub maybe_negative_prompt: Option, + + /// OPTIONAL. The generation mode (eg. keyframe, reference, inpaint, etc.) + pub maybe_generation_mode: Option, + + /// OPTIONAL. The aspect ratio (eg. square, auto, wide_three_by_two, etc.) + pub maybe_aspect_ratio: Option, + + /// OPTIONAL. The resolution (eg. one_k, two_k, four_k, etc.) + pub maybe_resolution: Option, + + /// OPTIONAL. The number of outputs to generate (0-255). + pub maybe_batch_count: Option, + + /// OPTIONAL. Whether to generate audio. + pub maybe_generate_audio: Option, + + /// OPTIONAL. The duration in seconds. + pub maybe_duration_seconds: Option, + + /// Context images (optional) + pub maybe_context_images: Option>, + + pub created_at: DateTime, +} diff --git a/crates/api_clients/artcraft/artcraft_api_defs/src/prompts/mod.rs b/crates/api_clients/artcraft/artcraft_api_defs/src/prompts/mod.rs index 4a4e360b56..90ff797ace 100644 --- a/crates/api_clients/artcraft/artcraft_api_defs/src/prompts/mod.rs +++ b/crates/api_clients/artcraft/artcraft_api_defs/src/prompts/mod.rs @@ -1,2 +1,3 @@ +pub mod batch_get_prompts; pub mod create_prompt; pub mod get_prompt; \ No newline at end of file diff --git a/crates/schema/database/mysql_queries/src/queries/prompts/batch_get_prompts.rs b/crates/schema/database/mysql_queries/src/queries/prompts/batch_get_prompts.rs new file mode 100644 index 0000000000..1dc640bbe7 --- /dev/null +++ b/crates/schema/database/mysql_queries/src/queries/prompts/batch_get_prompts.rs @@ -0,0 +1,92 @@ +use chrono::{DateTime, Utc}; +use sqlx::mysql::MySqlRow; +use sqlx::pool::PoolConnection; +use sqlx::{FromRow, MySql, QueryBuilder, Row}; + +use enums::by_table::prompts::prompt_type::PromptType; +use enums::common::generation::common_aspect_ratio::CommonAspectRatio; +use enums::common::generation::common_generation_mode::CommonGenerationMode; +use enums::common::generation::common_model_type::CommonModelType; +use enums::common::generation::common_resolution::CommonResolution; +use enums::common::generation_provider::GenerationProvider; +use enums::traits::mysql_from_row::MySqlFromRow; +use tokens::tokens::prompts::PromptToken; +use tokens::traits::mysql_token_from_row::MySqlTokenFromRow; + +pub struct BatchPromptResult { + pub token: PromptToken, + pub maybe_model_type: Option, + pub maybe_generation_provider: Option, + pub maybe_positive_prompt: Option, + pub maybe_negative_prompt: Option, + pub maybe_generation_mode: Option, + pub maybe_aspect_ratio: Option, + pub maybe_resolution: Option, + pub maybe_batch_count: Option, + pub maybe_generate_audio: Option, + pub maybe_duration_seconds: Option, + pub created_at: DateTime, +} + +impl FromRow<'_, MySqlRow> for BatchPromptResult { + fn from_row(row: &MySqlRow) -> Result { + Ok(Self { + token: PromptToken::try_from_mysql_row(row, "token")?, + maybe_model_type: row.try_get::, _>("maybe_model_type")? + .and_then(|s| CommonModelType::from_str(&s).ok()), + maybe_generation_provider: row.try_get::, _>("maybe_generation_provider")? + .and_then(|s| GenerationProvider::from_str(&s).ok()), + maybe_positive_prompt: row.try_get("maybe_positive_prompt")?, + maybe_negative_prompt: row.try_get("maybe_negative_prompt")?, + maybe_generation_mode: row.try_get::, _>("maybe_generation_mode")? + .and_then(|s| CommonGenerationMode::from_str(&s).ok()), + maybe_aspect_ratio: row.try_get::, _>("maybe_aspect_ratio")? + .and_then(|s| CommonAspectRatio::from_str(&s).ok()), + maybe_resolution: row.try_get::, _>("maybe_resolution")? + .and_then(|s| CommonResolution::from_str(&s).ok()), + maybe_batch_count: row.try_get("maybe_batch_count")?, + maybe_generate_audio: row.try_get("maybe_generate_audio")?, + maybe_duration_seconds: row.try_get("maybe_duration_seconds")?, + created_at: row.try_get("created_at")?, + }) + } +} + +pub async fn batch_get_prompts( + tokens: &[PromptToken], + mysql_connection: &mut PoolConnection, +) -> Result, sqlx::Error> { + if tokens.is_empty() { + return Ok(Vec::new()); + } + + let mut query_builder: QueryBuilder = QueryBuilder::new(r#" +SELECT + p.token, + p.maybe_model_type, + p.maybe_generation_provider, + p.maybe_positive_prompt, + p.maybe_negative_prompt, + p.maybe_generation_mode, + p.maybe_aspect_ratio, + p.maybe_resolution, + p.maybe_batch_count, + p.maybe_generate_audio, + p.maybe_duration_seconds, + p.created_at +FROM prompts as p +WHERE p.token IN ( + "#); + + let mut separated = query_builder.separated(", "); + + for token in tokens { + separated.push_bind(token.to_string()); + } + + separated.push_unseparated(") "); + + let query = query_builder.build_query_as::(); + + query.fetch_all(&mut **mysql_connection).await +} diff --git a/crates/schema/database/mysql_queries/src/queries/prompts/mod.rs b/crates/schema/database/mysql_queries/src/queries/prompts/mod.rs index 30492efa68..c0752a002a 100644 --- a/crates/schema/database/mysql_queries/src/queries/prompts/mod.rs +++ b/crates/schema/database/mysql_queries/src/queries/prompts/mod.rs @@ -1,2 +1,3 @@ +pub mod batch_get_prompts; pub mod get_prompt; pub mod insert_prompt; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/docs/api_doc.rs b/crates/service/web/storyteller_web/src/docs/api_doc.rs index d96591c077..348aadafd9 100644 --- a/crates/service/web/storyteller_web/src/docs/api_doc.rs +++ b/crates/service/web/storyteller_web/src/docs/api_doc.rs @@ -155,6 +155,7 @@ use artcraft_api_defs::media_file::delete_media_file::DeleteMediaFilePathInfo; use artcraft_api_defs::media_file::delete_media_file::DeleteMediaFileRequest; use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest; use artcraft_api_defs::prompts::create_prompt::CreatePromptResponse; +use artcraft_api_defs::prompts::batch_get_prompts::*; use artcraft_api_defs::prompts::get_prompt::*; use artcraft_api_defs::users::change_password::{ChangePasswordRequest, ChangePasswordResponse}; use artcraft_api_defs::users::edit_email::{EditEmailRequest, EditEmailResponse}; @@ -220,6 +221,7 @@ use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEnti use enums::by_table::user_ratings::entity_type::UserRatingEntityType; use enums::by_table::user_ratings::rating_value::UserRatingValue; use enums::by_table::users::user_feature_flag::UserFeatureFlag; +use enums::common::generation::common_model_class::CommonModelClass; use enums::common::generation::common_model_type::CommonModelType; use enums::common::generation::common_aspect_ratio::CommonAspectRatio; use enums::common::generation::common_generation_mode::CommonGenerationMode; @@ -402,6 +404,7 @@ use crate::http_server::endpoints::media_files::list::list_batch_generated_redux crate::http_server::endpoints::moderation::user_feature_flags::moderator_edit_user_feature_flags_handler::moderator_edit_user_feature_flags_handler, crate::http_server::endpoints::moderation::user_feature_flags::moderator_list_all_available_user_feature_flags_handler::moderator_list_all_available_user_feature_flags_handler, crate::http_server::endpoints::moderation::user_feature_flags::moderator_list_user_feature_flags_handler::moderator_list_user_feature_flags_handler, + crate::http_server::endpoints::prompts::batch_get_prompts_handler::batch_get_prompts_handler, crate::http_server::endpoints::prompts::create_prompt_handler::create_prompt_handler, crate::http_server::endpoints::prompts::get_prompt_handler::get_prompt_handler, crate::http_server::endpoints::service::status_alert_handler::status_alert_handler, @@ -559,6 +562,7 @@ use crate::http_server::endpoints::media_files::list::list_batch_generated_redux MediaFileOriginProductCategory, MediaFileSubtype, MediaFileType, + CommonModelClass, CommonModelType, CommonAspectRatio, CommonGenerationMode, @@ -621,6 +625,9 @@ use crate::http_server::endpoints::media_files::list::list_batch_generated_redux BatchGetInferenceJobStatusQueryParams, BatchGetInferenceJobStatusSuccessResponse, BatchGetMediaFilesModelInfo, + BatchGetPromptsRequest, + BatchGetPromptsResponse, + BatchPromptInfo, CreatePromptRequest, CreatePromptResponse, BatchGetMediaFilesQueryParams, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/prompts/batch_get_prompts_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/prompts/batch_get_prompts_handler.rs new file mode 100644 index 0000000000..aeaa825526 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/prompts/batch_get_prompts_handler.rs @@ -0,0 +1,159 @@ +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use actix_web::web::Json; +use actix_web::{web, HttpRequest}; +use artcraft_api_defs::prompts::batch_get_prompts::{ + BatchGetPromptsRequest, BatchGetPromptsResponse, BatchPromptInfo, +}; +use artcraft_api_defs::prompts::get_prompt::GetPromptImageContextItem; +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 log::warn; +use mysql_queries::queries::prompt_context_items::list_prompt_context_items::list_prompt_context_items; +use mysql_queries::queries::prompts::batch_get_prompts::batch_get_prompts; +use tokens::tokens::prompts::PromptToken; + +use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; +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 crate::state::server_state::ServerState; + +const MAX_BATCH_SIZE: usize = 100; + +/// Batch get details on multiple prompts. +#[utoipa::path( + post, + tag = "Prompts", + path = "/v1/prompts/batch", + request_body = BatchGetPromptsRequest, + responses( + (status = 200, description = "Found", body = BatchGetPromptsResponse), + (status = 400, description = "Bad request"), + (status = 500, description = "Server error"), + ), +)] +pub async fn batch_get_prompts_handler( + http_request: HttpRequest, + body: Json, + server_state: web::Data>, +) -> Result, AdvancedCommonWebError> { + let request = body.into_inner(); + + if request.tokens.len() > MAX_BATCH_SIZE { + return Err(AdvancedCommonWebError::BadInputWithSimpleMessage( + format!("tokens must contain between 0 and {} items", MAX_BATCH_SIZE), + )); + } + + // Deduplicate and trim whitespace + let unique_tokens: Vec = request + .tokens + .into_iter() + .map(|t| PromptToken::new_from_str(t.as_str().trim())) + .filter(|t| !t.as_str().is_empty()) + .collect::>() + .into_iter() + .collect(); + + if unique_tokens.is_empty() { + return Ok(Json(BatchGetPromptsResponse { + success: true, + prompts: Vec::new(), + })); + } + + let mut mysql_connection = server_state.mysql_pool.acquire().await?; + + let results = batch_get_prompts(&unique_tokens, &mut mysql_connection) + .await + .map_err(|err| { + warn!("Batch get prompts query error: {:?}", err); + AdvancedCommonWebError::from_error(err) + })?; + + let media_domain = get_media_domain(&http_request); + + // Fetch context items for each prompt + let mut context_items_map: HashMap> = HashMap::new(); + + for result in &results { + let items_result = list_prompt_context_items( + &result.token, + &mut mysql_connection, + ).await; + + let items = items_result.unwrap_or_else(|e| { + warn!("Error listing prompt context items for {}: {:?}", result.token.as_str(), e); + Vec::new() + }); + + let filtered_items: Vec = items + .iter() + .filter_map(|item| { + match item.context_semantic_type { + PromptContextSemanticType::VidStartFrame => {}, + PromptContextSemanticType::VidEndFrame => {}, + PromptContextSemanticType::Imgref => {}, + PromptContextSemanticType::ImgrefCharacter => {}, + PromptContextSemanticType::ImgrefStyle => {}, + PromptContextSemanticType::ImgrefBg => {}, + _ => { + return None + }, + } + + let bucket_path = MediaFileBucketPath::from_object_hash( + &item.public_bucket_directory_hash, + item.maybe_public_bucket_prefix.as_deref(), + item.maybe_public_bucket_extension.as_deref(), + ); + + Some(GetPromptImageContextItem { + media_token: item.media_token.clone(), + semantic: item.context_semantic_type, + media_links: MediaLinksBuilder::from_media_path_and_env( + media_domain, + server_state.server_environment, + &bucket_path, + ), + }) + }) + .collect(); + + if !filtered_items.is_empty() { + context_items_map.insert(result.token.as_str().to_string(), filtered_items); + } + } + + let prompts: Vec = results + .into_iter() + .map(|result| { + let maybe_context_images = context_items_map + .remove(result.token.as_str()) + .filter(|items| !items.is_empty()); + + BatchPromptInfo { + token: result.token, + maybe_model_class: result.maybe_model_type.map(|ty| ty.get_model_class()), + maybe_model_type: result.maybe_model_type, + maybe_generation_provider: result.maybe_generation_provider, + maybe_positive_prompt: result.maybe_positive_prompt, + maybe_negative_prompt: result.maybe_negative_prompt, + maybe_generation_mode: result.maybe_generation_mode, + maybe_aspect_ratio: result.maybe_aspect_ratio, + maybe_resolution: result.maybe_resolution, + maybe_batch_count: result.maybe_batch_count, + maybe_generate_audio: result.maybe_generate_audio, + maybe_duration_seconds: result.maybe_duration_seconds, + maybe_context_images, + created_at: result.created_at, + } + }) + .collect(); + + Ok(Json(BatchGetPromptsResponse { + success: true, + prompts, + })) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/prompts/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/prompts/mod.rs index 3b028ec003..dbf92d70a4 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/prompts/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/prompts/mod.rs @@ -1,2 +1,3 @@ +pub mod batch_get_prompts_handler; pub mod create_prompt_handler; pub mod get_prompt_handler; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/routes/application_routes/prompts_routes.rs b/crates/service/web/storyteller_web/src/http_server/routes/application_routes/prompts_routes.rs index 99771ebf20..c54df1c9d4 100644 --- a/crates/service/web/storyteller_web/src/http_server/routes/application_routes/prompts_routes.rs +++ b/crates/service/web/storyteller_web/src/http_server/routes/application_routes/prompts_routes.rs @@ -1,3 +1,4 @@ +use crate::http_server::endpoints::prompts::batch_get_prompts_handler::batch_get_prompts_handler; use crate::http_server::endpoints::prompts::create_prompt_handler::create_prompt_handler; use crate::http_server::endpoints::prompts::get_prompt_handler::get_prompt_handler; use actix_helpers::route_builder::RouteBuilder; @@ -21,6 +22,7 @@ where RouteBuilder::from_app(app) // NB: This poor RouteBuilder utility requires that POST comes first, otherwise the GET glob // will capture it and force 504 Method Not Allowed for all POSTs. + .add_post("/v1/prompts/batch", batch_get_prompts_handler) .add_post("/v1/prompts/create", create_prompt_handler) .add_get("/v1/prompts/{token}", get_prompt_handler) .into_app()