From f00dc3a006bb1d889bc12795487166780790ccd9 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 30 Apr 2026 19:07:03 -0400 Subject: [PATCH] batch get prompts endpoint (#1457) --- .../src/prompts/batch_get_prompts.rs | 73 ++++++++ .../src/prompts/get_prompt.rs | 14 +- .../artcraft_api_defs/src/prompts/mod.rs | 1 + .../batch_list_prompt_context_items.rs | 70 ++++++++ .../src/queries/prompt_context_items/mod.rs | 1 + .../src/queries/prompts/batch_get_prompts.rs | 95 +++++++++++ .../mysql_queries/src/queries/prompts/mod.rs | 1 + .../web/storyteller_web/src/docs/api_doc.rs | 8 +- .../prompts/batch_get_prompts_handler.rs | 156 ++++++++++++++++++ .../endpoints/prompts/get_prompt_handler.rs | 120 ++++---------- .../src/http_server/endpoints/prompts/mod.rs | 1 + .../application_routes/prompts_routes.rs | 2 + 12 files changed, 446 insertions(+), 96 deletions(-) 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/prompt_context_items/batch_list_prompt_context_items.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..5ab9aa7487 --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_api_defs/src/prompts/batch_get_prompts.rs @@ -0,0 +1,73 @@ +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 std::collections::HashSet; + +use serde_derive::{Deserialize, Serialize}; +use tokens::tokens::prompts::PromptToken; +use utoipa::{IntoParams, ToSchema}; + +pub const BATCH_GET_PROMPTS_PATH: &str = "/v1/prompt/batch"; + +#[derive(Deserialize, ToSchema, IntoParams)] +pub struct BatchGetPromptsQuery { + /// Prompt tokens to look up. + pub tokens: HashSet, +} + +#[derive(Serialize, ToSchema)] +pub struct BatchGetPromptsResponse { + pub success: bool, + pub prompts: Vec, +} + +#[derive(Serialize, ToSchema)] +pub struct BatchPromptInfo { + pub token: PromptToken, + + pub prompt_type: PromptType, + + /// 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/get_prompt.rs b/crates/api_clients/artcraft/artcraft_api_defs/src/prompts/get_prompt.rs index 8eedbce967..77369ec369 100644 --- a/crates/api_clients/artcraft/artcraft_api_defs/src/prompts/get_prompt.rs +++ b/crates/api_clients/artcraft/artcraft_api_defs/src/prompts/get_prompt.rs @@ -27,13 +27,6 @@ pub struct GetPromptSuccessResponse { pub prompt: PromptInfo, } -#[derive(Serialize, ToSchema)] -pub struct GetPromptImageContextItem { - pub media_token: MediaFileToken, - pub semantic: PromptContextSemanticType, - pub media_links: MediaLinks, -} - #[derive(Serialize, ToSchema)] pub struct PromptInfo { pub token: PromptToken, @@ -130,6 +123,13 @@ pub struct PromptInfo { pub created_at: DateTime, } +#[derive(Serialize, ToSchema)] +pub struct GetPromptImageContextItem { + pub media_token: MediaFileToken, + pub semantic: PromptContextSemanticType, + pub media_links: MediaLinks, +} + #[derive(Serialize, ToSchema)] pub struct PromptInfoModeratorFields { /// How many milliseconds it took to run generation. 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/prompt_context_items/batch_list_prompt_context_items.rs b/crates/schema/database/mysql_queries/src/queries/prompt_context_items/batch_list_prompt_context_items.rs new file mode 100644 index 0000000000..bfe8d50799 --- /dev/null +++ b/crates/schema/database/mysql_queries/src/queries/prompt_context_items/batch_list_prompt_context_items.rs @@ -0,0 +1,70 @@ +use sqlx::mysql::MySqlRow; +use sqlx::pool::PoolConnection; +use sqlx::{FromRow, MySql, QueryBuilder, Row}; + +use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType; +use enums::traits::mysql_from_row::MySqlFromRow; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::prompts::PromptToken; +use tokens::traits::mysql_token_from_row::MySqlTokenFromRow; + +/// A prompt context item with its owning prompt token, for batch queries. +pub struct BatchPromptContextItem { + pub prompt_token: PromptToken, + pub media_token: MediaFileToken, + pub context_semantic_type: PromptContextSemanticType, + pub public_bucket_directory_hash: String, + pub maybe_public_bucket_prefix: Option, + pub maybe_public_bucket_extension: Option, +} + +impl FromRow<'_, MySqlRow> for BatchPromptContextItem { + fn from_row(row: &MySqlRow) -> Result { + Ok(Self { + prompt_token: PromptToken::try_from_mysql_row(row, "prompt_token")?, + media_token: MediaFileToken::try_from_mysql_row(row, "media_token")?, + context_semantic_type: PromptContextSemanticType::try_from_mysql_row(row, "context_semantic_type")?, + public_bucket_directory_hash: row.try_get("public_bucket_directory_hash")?, + maybe_public_bucket_prefix: row.try_get("maybe_public_bucket_prefix")?, + maybe_public_bucket_extension: row.try_get("maybe_public_bucket_extension")?, + }) + } +} + +/// Batch fetch all context items for multiple prompts in a single query. +/// Returns items grouped by their prompt_token (via the prompt_token field). +pub async fn batch_list_prompt_context_items( + prompt_tokens: &[PromptToken], + mysql_connection: &mut PoolConnection, +) -> Result, sqlx::Error> { + if prompt_tokens.is_empty() { + return Ok(Vec::new()); + } + + let mut query_builder: QueryBuilder = QueryBuilder::new(r#" +SELECT + pci.prompt_token, + pci.media_token, + pci.context_semantic_type, + + m.public_bucket_directory_hash, + m.maybe_public_bucket_prefix, + m.maybe_public_bucket_extension + +FROM prompt_context_items pci +JOIN media_files m ON pci.media_token = m.token +WHERE pci.prompt_token IN ( +"#); + + let mut separated = query_builder.separated(", "); + for token in prompt_tokens { + separated.push_bind(token.to_string()); + } + separated.push_unseparated(") "); + + let query = query_builder.build_query_as::(); + + let results = query.fetch_all(&mut **mysql_connection).await?; + + Ok(results) +} diff --git a/crates/schema/database/mysql_queries/src/queries/prompt_context_items/mod.rs b/crates/schema/database/mysql_queries/src/queries/prompt_context_items/mod.rs index 7611a2e790..57380a3d52 100644 --- a/crates/schema/database/mysql_queries/src/queries/prompt_context_items/mod.rs +++ b/crates/schema/database/mysql_queries/src/queries/prompt_context_items/mod.rs @@ -1,2 +1,3 @@ +pub mod batch_list_prompt_context_items; pub mod insert_batch_prompt_context_items; pub mod list_prompt_context_items; \ 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..d0e2bc7590 --- /dev/null +++ b/crates/schema/database/mysql_queries/src/queries/prompts/batch_get_prompts.rs @@ -0,0 +1,95 @@ +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 prompt_type: PromptType, + 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")?, + prompt_type: PromptType::try_from_mysql_row(row, "prompt_type")?, + 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.prompt_type, + 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 fba8bfe7a2..35210bc1d7 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, + BatchGetPromptsQuery, + BatchGetPromptsResponse, + BatchPromptInfo, CreatePromptRequest, CreatePromptResponse, BatchGetMediaFilesQueryParams, @@ -862,7 +869,6 @@ use crate::http_server::endpoints::media_files::list::list_batch_generated_redux GetMediaFilePathInfo, GetMediaFileSuccessResponse, GetProfilePathInfo, - GetPromptError, GetPromptPathInfo, GetPromptSuccessResponse, GetUnifiedQueueStatsError, 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..2626fb02b5 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/prompts/batch_get_prompts_handler.rs @@ -0,0 +1,156 @@ +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use actix_web::web::Json; +use actix_web::{web, HttpRequest}; +use actix_web_lab::extract::Query; +use artcraft_api_defs::prompts::batch_get_prompts::{ + BatchGetPromptsQuery, 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::batch_list_prompt_context_items::batch_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. +/// +/// Called, eg. {{api_host}}/v1/prompt/batch?tokens={val}&tokens={val}&tokens=... +#[utoipa::path( + get, + tag = "Prompts", + path = "/v1/prompt/batch", + params(BatchGetPromptsQuery), + 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, + query: Query, + server_state: web::Data>, +) -> Result, AdvancedCommonWebError> { + if query.tokens.len() > MAX_BATCH_SIZE { + return Err(AdvancedCommonWebError::BadInputWithSimpleMessage( + format!("tokens must contain at most {} items", MAX_BATCH_SIZE), + )); + } + + // Deduplicate and trim whitespace + let unique_tokens: Vec = query.tokens + .iter() + .map(|t| PromptToken::new_from_str(t.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?; + + // Batch query: fetch all prompts in one query + 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) + })?; + + // Batch query: fetch all context items for all prompts in one query + let all_context_items = batch_list_prompt_context_items(&unique_tokens, &mut mysql_connection) + .await + .map_err(|err| { + warn!("Batch list prompt context items query error: {:?}", err); + AdvancedCommonWebError::from_error(err) + })?; + + let media_domain = get_media_domain(&http_request); + + // Group context items by prompt token + let mut context_items_map: HashMap> = HashMap::new(); + + for item in all_context_items { + // Filter to only image-like context types + match item.context_semantic_type { + PromptContextSemanticType::VidStartFrame + | PromptContextSemanticType::VidEndFrame + | PromptContextSemanticType::Imgref + | PromptContextSemanticType::ImgrefCharacter + | PromptContextSemanticType::ImgrefStyle + | PromptContextSemanticType::ImgrefBg + | PromptContextSemanticType::VidRef => {} // Include + _ => continue, // Ignore + } + + 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(), + ); + + let context_item = GetPromptImageContextItem { + media_token: item.media_token, + semantic: item.context_semantic_type, + media_links: MediaLinksBuilder::from_media_path_and_env( + media_domain, + server_state.server_environment, + &bucket_path, + ), + }; + + context_items_map + .entry(item.prompt_token.as_str().to_string()) + .or_default() + .push(context_item); + } + + // Assemble final response + 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, + prompt_type: result.prompt_type, + 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/get_prompt_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/prompts/get_prompt_handler.rs index 7ab383cbfa..0a3e770e51 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/prompts/get_prompt_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/prompts/get_prompt_handler.rs @@ -1,14 +1,7 @@ -use std::fmt; use std::sync::Arc; -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::http_server::web_utils::response_error_helpers::to_simple_json_error; -use crate::state::server_state::ServerState; -use actix_web::error::ResponseError; -use actix_web::http::StatusCode; -use actix_web::web::Path; -use actix_web::{web, HttpMessage, HttpRequest, HttpResponse}; +use actix_web::web::{Json, Path}; +use actix_web::{web, HttpMessage, HttpRequest}; use artcraft_api_defs::prompts::get_prompt::{ GetPromptImageContextItem, GetPromptPathInfo, GetPromptSuccessResponse, PromptInfo, PromptInfoModeratorFields, @@ -18,38 +11,11 @@ use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptC use log::{error, warn}; use mysql_queries::queries::prompt_context_items::list_prompt_context_items::list_prompt_context_items; use mysql_queries::queries::prompts::get_prompt::get_prompt_from_connection; -use utoipa::ToSchema; -#[derive(Debug, ToSchema)] -pub enum GetPromptError { - ServerError, - NotFound, -} - -impl ResponseError for GetPromptError { - fn status_code(&self) -> StatusCode { - match *self { - GetPromptError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, - GetPromptError::NotFound => StatusCode::NOT_FOUND, - } - } - - fn error_response(&self) -> HttpResponse { - let error_reason = match self { - GetPromptError::ServerError => "server error".to_string(), - GetPromptError::NotFound => "not found".to_string(), - }; - - to_simple_json_error(&error_reason, self.status_code()) - } -} - -// NB: Not using derive_more::Display since Clion doesn't understand it. -impl fmt::Display for GetPromptError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{:?}", self) - } -} +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; /// Get details on a prompt. #[utoipa::path( @@ -58,8 +24,8 @@ impl fmt::Display for GetPromptError { path = "/v1/prompts/{token}", responses( (status = 200, description = "Found", body = GetPromptSuccessResponse), - (status = 404, description = "Not found", body = GetPromptError), - (status = 500, description = "Server error", body = GetPromptError), + (status = 404, description = "Not found"), + (status = 500, description = "Server error"), ), params( ("path" = GetPromptPathInfo, description = "Path for Request") @@ -68,44 +34,32 @@ impl fmt::Display for GetPromptError { pub async fn get_prompt_handler( http_request: HttpRequest, path: Path, - server_state: web::Data>) -> Result -{ - let mut mysql_connection = server_state.mysql_pool - .acquire() - .await - .map_err(|err| { - error!("Error acquiring MySQL connection: {:?}", err); - GetPromptError::ServerError - })?; + server_state: web::Data>, +) -> Result, AdvancedCommonWebError> { + 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); - GetPromptError::ServerError - })?; + .session_checker + .maybe_get_user_session_from_connection(&http_request, &mut mysql_connection) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + AdvancedCommonWebError::from(e) + })?; let is_moderator = maybe_user_session - .map(|session| session.can_ban_users) - .unwrap_or(false); + .map(|session| session.can_ban_users) + .unwrap_or(false); let prompt_token = path.into_inner().token; - let result = get_prompt_from_connection( - &prompt_token, - &mut mysql_connection - ).await; - - let result = match result { - Err(e) => { - warn!("query error: {:?}", e); - return Err(GetPromptError::ServerError); - } - Ok(None) => return Err(GetPromptError::NotFound), - Ok(Some(result)) => result, - }; + let result = get_prompt_from_connection(&prompt_token, &mut mysql_connection) + .await + .map_err(|err| { + warn!("query error: {:?}", err); + AdvancedCommonWebError::from(err) + })? + .ok_or(AdvancedCommonWebError::NotFound)?; let mut maybe_style_name = None; let mut maybe_strength = None; @@ -114,14 +68,12 @@ pub async fn get_prompt_handler( let mut maybe_travel_prompt = None; let mut maybe_frame_skip = None; - // Flags let mut used_face_detailer = false; let mut used_upscaler = false; let mut lipsync_enabled = false; let mut lcm_disabled = false; let mut use_cinematic = false; - // Moderator fields let mut main_ipa_workflow = None; let mut face_detailer_workflow = None; let mut upscaler_workflow = None; @@ -136,14 +88,12 @@ pub async fn get_prompt_handler( maybe_travel_prompt = inner_payload.travel_prompt.clone(); maybe_frame_skip = inner_payload.frame_skip; - // Flags used_face_detailer = inner_payload.used_face_detailer.unwrap_or(false); used_upscaler = inner_payload.used_upscaler.unwrap_or(false); lipsync_enabled = inner_payload.lipsync_enabled.unwrap_or(false); lcm_disabled = inner_payload.disable_lcm.unwrap_or(false); use_cinematic = inner_payload.use_cinematic.unwrap_or(false); - // Moderator fields main_ipa_workflow = inner_payload.main_ipa_workflow.clone(); face_detailer_workflow = inner_payload.face_detailer_workflow.clone(); upscaler_workflow = inner_payload.upscaler_workflow.clone(); @@ -164,7 +114,7 @@ pub async fn get_prompt_handler( let items_result = list_prompt_context_items( &result.token, - &mut mysql_connection + &mut mysql_connection, ).await; let items = items_result.unwrap_or_else(|e| { @@ -189,7 +139,8 @@ pub async fn get_prompt_handler( 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()); + item.maybe_public_bucket_extension.as_deref(), + ); Some(GetPromptImageContextItem { media_token: item.media_token.clone(), @@ -208,7 +159,7 @@ pub async fn get_prompt_handler( Some(items) }; - let response = GetPromptSuccessResponse { + Ok(Json(GetPromptSuccessResponse { success: true, prompt: PromptInfo { token: result.token, @@ -239,12 +190,5 @@ pub async fn get_prompt_handler( maybe_global_ipa_image_token, maybe_frame_skip, }, - }; - - let body = serde_json::to_string(&response) - .map_err(|e| GetPromptError::ServerError)?; - - Ok(HttpResponse::Ok() - .content_type("application/json") - .body(body)) + })) } 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..fa34a7dbfd 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_get("/v1/prompt/batch", batch_get_prompts_handler) .add_post("/v1/prompts/create", create_prompt_handler) .add_get("/v1/prompts/{token}", get_prompt_handler) .into_app()