provide prompt token in media file upload

This commit is contained in:
Brandon Thomas committed 2025-08-15 21:11:49 -04:00
1 parent 457ae45af9
commit 905282b27b
22 files changed
+138 -25

No files matched your search

@@ -7,6 +7,9 @@ use crate::error::api_error::ApiError;
/// An error arising from the client, either before or after the request is sent.
#[derive(Debug)]
pub enum ClientError {
/// An invalid request was detected before it was sent.
InvalidPreflightRequest(String),
/// Could not determine the filetype for the file.
FileTypeNotKnown(PathBuf),
@@ -28,6 +31,7 @@ impl Error for ClientError {}
impl Display for ClientError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::InvalidPreflightRequest(msg) => write!(f, "Invalid preflight request: {}", msg),
Self::FileTypeNotKnown(path) => write!(f, "Could not determine the filetype for the file: {:?}", path),
Self::FileTypeNotHandled(path) => write!(f, "Could not handle the file type for the file: {:?}", path),
Self::IoError(err) => write!(f, "IO error: {}", err),
@@ -5,6 +5,7 @@ pub mod error;
pub mod generate;
pub mod jobs;
pub mod media_files;
pub mod prompts;
pub mod recipes;
pub mod shared_response_types;
pub mod utils;
@@ -16,6 +16,7 @@ use enums::by_table::media_files::media_file_class::MediaFileClass;
use enums::by_table::media_files::media_file_engine_category::MediaFileEngineCategory;
use enums::by_table::media_files::media_file_type::MediaFileType;
use enums::by_table::model_weights::weights_category::WeightsCategory;
use enums::common::generation_provider::GenerationProvider;
use enums::common::visibility::Visibility;
use enums::no_table::style_transfer::style_transfer_name::StyleTransferName;
use enums_public::by_table::model_weights::public_weights_types::PublicWeightsType;
@@ -42,6 +43,14 @@ pub struct UploadImageFromFileArgs<'a, P: AsRef<Path>> {
/// If true, we should hide the image from the user's gallery.
pub is_intermediate_system_file: bool,
/// If provided, this is the prompt that this image is associated with.
/// NOTE: Cannot set `is_intermediate_system_file = true` if this is set.
pub maybe_prompt_token: Option<PromptToken>,
// /// If provided, this is the service provider that created the image.
// /// NOTE: Cannot set `is_intermediate_system_file = true` if this is set.
// pub maybe_generation_provider: Option<GenerationProvider>,
}
@@ -49,6 +58,8 @@ pub struct UploadImageFromFileArgs<'a, P: AsRef<Path>> {
pub async fn upload_image_media_file_from_file<P: AsRef<Path>>(
args: UploadImageFromFileArgs<'_, P>,
) -> Result<UploadImageMediaFileSuccessResponse, StorytellerError> {
validate_args(&args)?;
let url = get_route(args.api_host);
@@ -70,6 +81,10 @@ pub async fn upload_image_media_file_from_file<P: AsRef<Path>>(
if args.is_intermediate_system_file {
form = form.text("is_intermediate_system_file", "true");
}
if let Some(prompt_token) = &args.maybe_prompt_token {
form = form.text("maybe_prompt_token", prompt_token.to_string());
}
let mut request_builder = client.post(url)
.header("User-Agent", USER_AGENT)
@@ -109,3 +124,15 @@ pub struct UploadImageMediaFileSuccessResponse {
pub success: bool,
pub media_file_token: MediaFileToken,
}
fn validate_args<P: AsRef<Path>>(
args: &UploadImageFromFileArgs<'_, P>,
) -> Result<(), StorytellerError> {
if args.is_intermediate_system_file && args.maybe_prompt_token.is_some() {
return Err(StorytellerError::Client(ClientError::InvalidPreflightRequest(
"Cannot set `is_intermediate_system_file` to true if `maybe_prompt_token` is provided."
.to_string())));
}
Ok(())
}
@@ -0,0 +1,18 @@
use crate::credentials::storyteller_credential_set::StorytellerCredentialSet;
use crate::error::storyteller_error::StorytellerError;
use crate::utils::api_host::ApiHost;
use crate::utils::basic_json_post_request::basic_json_post_request;
use artcraft_api_defs::prompts::create_prompt::{CreatePromptRequest, CreatePromptResponse, CREATE_PROMPT_PATH};
pub async fn create_prompt(
api_host: &ApiHost,
maybe_creds: Option<&StorytellerCredentialSet>,
request: CreatePromptRequest,
) -> Result<CreatePromptResponse, StorytellerError> {
Ok(basic_json_post_request(
api_host,
CREATE_PROMPT_PATH,
maybe_creds,
request,
).await?)
}
@@ -0,0 +1 @@
pub mod create_prompt;
@@ -8,10 +8,14 @@ use crate::core::state::task_database::TaskDatabase;
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
use crate::services::midjourney::utils::download_midjourney_image::download_midjourney_image;
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
use cookie_store::cookie_store::CookieStore;
use enums::by_table::prompts::prompt_type::PromptType;
use enums::common::generation_provider::GenerationProvider;
use enums::common::model_type::ModelType;
use enums::tauri::tasks::task_status::TaskStatus;
use errors::AnyhowResult;
use idempotency::uuid::generate_random_uuid;
use log::{error, info};
use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId;
@@ -29,6 +33,7 @@ use storyteller_client::credentials::storyteller_credential_set::StorytellerCred
use storyteller_client::error::api_error::ApiError;
use storyteller_client::error::storyteller_error::StorytellerError;
use storyteller_client::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs};
use storyteller_client::prompts::create_prompt::create_prompt;
use tauri::AppHandle;
use url::Url;
@@ -196,6 +201,22 @@ async fn check_midjourney_tasks(
None => continue,
};
let request = CreatePromptRequest {
uuid_idempotency_token: generate_random_uuid(),
positive_prompt: midjourney_item.full_command.clone(),
negative_prompt: None,
model_type: Some(ModelType::Midjourney),
generation_provider: Some(GenerationProvider::Midjourney),
};
let prompt_response = create_prompt(
&app_env_configs.storyteller_host,
Some(storyteller_creds),
request
).await?;
info!("Created prompt: {:?}", prompt_response.prompt_token);
for index in 0..4 {
info!("Downloading generated Midjourney file...");
@@ -13,6 +13,7 @@ use enums::common::visibility::Visibility;
use errors::AnyhowResult;
use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken;
use tokens::tokens::media_files::MediaFileToken;
use tokens::tokens::prompts::PromptToken;
use tokens::tokens::users::UserToken;
use crate::queries::generic_synthetic_ids::transactional_increment_generic_synthetic_id::transactional_increment_generic_synthetic_id;
@@ -43,6 +44,8 @@ pub struct InsertMediaFileFromUploadArgs<'a> {
pub file_size_bytes: u64,
pub sha256_checksum: &'a str,
pub maybe_prompt_token: Option<&'a PromptToken>,
pub maybe_title: Option<&'a str>,
pub maybe_duration_millis: Option<u64>,
@@ -112,6 +115,8 @@ SET
checksum_sha2 = ?,
maybe_prompt_token = ?,
maybe_title = ?,
maybe_duration_millis = ?,
@@ -155,6 +160,8 @@ SET
args.sha256_checksum,
args.maybe_prompt_token.map(|t| t.as_str()),
args.maybe_title,
args.maybe_duration_millis,
@@ -22,9 +22,6 @@ pub struct CreatePromptRequest {
/// OPTIONAL. The service used.
pub generation_provider: Option<GenerationProvider>,
/// OPTIONAL. The type of prompt.
pub prompt_type: Option<String>,
}
#[derive(Serialize, Deserialize, ToSchema)]
@@ -30,6 +30,10 @@ pub enum ModelType {
#[serde(rename = "recraft_3")]
Recraft3,
/// Midjourney without distinguishing a model type or version
#[serde(rename = "midjourney")]
Midjourney,
//// Image Infill models
//#[serde(rename = "flux_pro_1_infill")]
//FluxPro1Infill,
@@ -71,6 +75,7 @@ impl ModelType {
Self::FluxProKontextMax => "flux_pro_kontext_max",
Self::GptImage1 => "gpt_image_1",
Self::Recraft3 => "recraft_3",
Self::Midjourney => "midjourney",
// Video models
Self::Kling16Pro => "kling_1p6_pro",
@@ -96,6 +101,7 @@ impl ModelType {
"flux_pro_kontext_max" => Ok(Self::FluxProKontextMax),
"gpt_image_1" => Ok(Self::GptImage1),
"recraft_3" => Ok(Self::Recraft3),
"midjourney" => Ok(Self::Midjourney),
// Video models
"kling_1p6_pro" => Ok(Self::Kling16Pro),
@@ -125,6 +131,7 @@ impl ModelType {
Self::FluxProKontextMax,
Self::GptImage1,
Self::Recraft3,
Self::Midjourney,
// Video models
Self::Kling16Pro,
@@ -159,6 +166,7 @@ mod tests {
assert_serialization(ModelType::FluxProKontextMax, "flux_pro_kontext_max");
assert_serialization(ModelType::GptImage1, "gpt_image_1");
assert_serialization(ModelType::Recraft3, "recraft_3");
assert_serialization(ModelType::Midjourney, "midjourney");
// Video models
assert_serialization(ModelType::Kling16Pro, "kling_1p6_pro");
assert_serialization(ModelType::Kling21Pro, "kling_2p1_pro");
@@ -181,6 +189,7 @@ mod tests {
assert_eq!(ModelType::FluxProKontextMax.to_str(), "flux_pro_kontext_max");
assert_eq!(ModelType::GptImage1.to_str(), "gpt_image_1");
assert_eq!(ModelType::Recraft3.to_str(), "recraft_3");
assert_eq!(ModelType::Midjourney.to_str(), "midjourney");
// Video models
assert_eq!(ModelType::Kling16Pro.to_str(), "kling_1p6_pro");
@@ -205,6 +214,7 @@ mod tests {
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("recraft_3").unwrap(), ModelType::Recraft3);
assert_eq!(ModelType::from_str("midjourney").unwrap(), ModelType::Midjourney);
// Video models
assert_eq!(ModelType::from_str("kling_1p6_pro").unwrap(), ModelType::Kling16Pro);
assert_eq!(ModelType::from_str("kling_2p1_pro").unwrap(), ModelType::Kling21Pro);
@@ -219,7 +229,7 @@ mod tests {
#[test]
fn all_variants() {
let mut variants = ModelType::all_variants();
assert_eq!(variants.len(), 15);
assert_eq!(variants.len(), 16);
// Image models
assert_eq!(variants.pop_first(), Some(ModelType::Flux1Dev));
assert_eq!(variants.pop_first(), Some(ModelType::Flux1Schnell));
@@ -229,6 +239,7 @@ mod tests {
assert_eq!(variants.pop_first(), Some(ModelType::FluxProKontextMax));
assert_eq!(variants.pop_first(), Some(ModelType::GptImage1));
assert_eq!(variants.pop_first(), Some(ModelType::Recraft3));
assert_eq!(variants.pop_first(), Some(ModelType::Midjourney));
// Video models
assert_eq!(variants.pop_first(), Some(ModelType::Kling16Pro));
assert_eq!(variants.pop_first(), Some(ModelType::Kling21Pro));
@@ -205,6 +205,7 @@ pub async fn create_scene_handler(
maybe_engine_category: Some(MediaFileEngineCategory::Scene),
maybe_animation_type: None,
maybe_mime_type: Some(mime_type),
maybe_prompt_token: None,
file_size_bytes: file_size_bytes as u64,
maybe_duration_millis: None,
sha256_checksum: &hash,
@@ -332,6 +332,7 @@ pub async fn process_upload_media_file(
upload_type,
maybe_engine_category: None,
maybe_animation_type: None,
maybe_prompt_token: None,
maybe_mime_type: Some(mime_type),
file_size_bytes: file_size_bytes as u64,
maybe_duration_millis: None,
@@ -246,6 +246,7 @@ pub async fn upload_snapshot_media_file_handler(
upload_type: UploadType::Filesystem, // TODO(bt,2024-05-02): This should be a parameter and a well-known enum.
maybe_engine_category: None,
maybe_animation_type: None,
maybe_prompt_token: None,
maybe_mime_type: Some(&mimetype),
file_size_bytes: file_size_bytes as u64,
maybe_duration_millis: None,
@@ -324,6 +324,7 @@ pub async fn upload_audio_media_file_handler(
upload_type: UploadType::Filesystem, // TODO(bt,2024-05-02): This should be a parameter and a well-known enum.
maybe_engine_category: None,
maybe_animation_type: None,
maybe_prompt_token: None,
maybe_mime_type: Some(&mimetype),
file_size_bytes: file_size_bytes as u64,
maybe_duration_millis,
@@ -209,6 +209,7 @@ pub async fn upload_engine_asset_media_file_handler(
maybe_creator_anonymous_visitor_token: maybe_avt_token.as_ref(),
creator_ip_address: &ip_address,
creator_set_visibility,
maybe_prompt_token: None,
upload_type: UploadType::Filesystem,
maybe_engine_category: upload_media_request.maybe_engine_category,
maybe_animation_type: upload_media_request.maybe_animation_type,
@@ -364,6 +364,7 @@ pub async fn process_upload_media_file(
creator_set_visibility,
upload_type,
maybe_engine_category: None,
maybe_prompt_token: None,
maybe_animation_type: None,
maybe_mime_type: Some(mime_type),
file_size_bytes: file_size_bytes as u64,
@@ -12,9 +12,13 @@ use log::{error, info, warn};
use once_cell::sync::Lazy;
use utoipa::ToSchema;
use crate::http_server::endpoints::media_files::upload::upload_error::MediaFileUploadError;
use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format;
use crate::state::server_state::ServerState;
use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
use enums::by_table::media_files::media_file_class::MediaFileClass;
use enums::by_table::media_files::media_file_type::MediaFileType;
use enums::common::generation_provider::GenerationProvider;
use enums::common::visibility::Visibility;
use hashing::sha256::sha256_hash_bytes::sha256_hash_bytes;
use http_server_common::request::get_request_ip::get_request_ip;
@@ -24,9 +28,6 @@ use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_
use mysql_queries::queries::media_files::create::specialized_insert::insert_media_file_from_file_upload::{insert_media_file_from_file_upload, InsertMediaFileFromUploadArgs, UploadType};
use tokens::tokens::media_files::MediaFileToken;
use tokens::tokens::prompts::PromptToken;
use crate::http_server::endpoints::media_files::upload::upload_error::MediaFileUploadError;
use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format;
use crate::state::server_state::ServerState;
/// Form-multipart request fields.
///
@@ -55,24 +56,11 @@ pub struct UploadImageMediaFileForm {
#[schema(value_type = Option<Visibility>, format = Binary)]
maybe_visibility: Option<Text<Visibility>>,
// TODO: These fields --
// /// Optional: Prompt associated with this image
// /// NOTE: Cannot set `is_intermediate_system_file = true` if this is set.
// #[multipart(limit = "2 KiB")]
// #[schema(value_type = Option<String>, format = Binary)]
// maybe_prompt_token: Option<Text<PromptToken>>,
// /// Optional: Provider of image generation service
// /// NOTE: Cannot set `is_intermediate_system_file = true` if this is set.
// #[multipart(limit = "2 KiB")]
// #[schema(value_type = Option<String>, format = Binary)]
// maybe_generation_provider: Option<Text<String>>,
// /// Optional: Type of model used
// /// NOTE: Cannot set `is_intermediate_system_file = true` if this is set.
// #[multipart(limit = "2 KiB")]
// #[schema(value_type = Option<String>, format = Binary)]
// maybe_origin_model_type: Option<Text<String>>,
/// Optional: Prompt associated with this image
/// NOTE: Cannot set `is_intermediate_system_file = true` if this is set.
#[multipart(limit = "2 KiB")]
#[schema(value_type = Option<String>, format = Binary)]
maybe_prompt_token: Option<Text<PromptToken>>,
/// Optional: Whether this is a system file (eg. cover files we should hide)
#[multipart(limit = "2 KiB")]
@@ -120,6 +108,8 @@ pub async fn upload_image_media_file_handler(
server_state: web::Data<Arc<ServerState>>,
MultipartForm(mut form): MultipartForm<UploadImageMediaFileForm>,
) -> Result<Json<UploadImageMediaFileSuccessResponse>, MediaFileUploadError> {
validate_form(&form)?;
let mut mysql_connection = server_state.mysql_pool
.acquire()
@@ -275,6 +265,10 @@ pub async fn upload_image_media_file_handler(
let is_intermediate_system_file = form.is_intermediate_system_file
.map(|b| b.0)
.unwrap_or(false);
let maybe_prompt_token = form.maybe_prompt_token
.as_ref()
.map(|token| token.0.clone());
let (token, record_id) = insert_media_file_from_file_upload(InsertMediaFileFromUploadArgs {
maybe_media_class: Some(MediaFileClass::Image),
@@ -287,6 +281,7 @@ pub async fn upload_image_media_file_handler(
maybe_engine_category: None,
maybe_animation_type: None,
maybe_mime_type: Some(&mimetype),
maybe_prompt_token: maybe_prompt_token.as_ref(),
file_size_bytes: file_size_bytes as u64,
maybe_duration_millis: None,
sha256_checksum: &hash,
@@ -311,3 +306,22 @@ pub async fn upload_image_media_file_handler(
media_file_token: token,
}))
}
fn validate_form(form: &UploadImageMediaFileForm) -> Result<(), MediaFileUploadError> {
if form.uuid_idempotency_token.is_empty() {
return Err(MediaFileUploadError::BadInput("Idempotency token is required".to_string()));
}
let is_intermediate_system_file = form.is_intermediate_system_file
.as_ref()
.map(|b| b.0)
.unwrap_or(false);
if is_intermediate_system_file && form.maybe_prompt_token.is_some() {
return Err(MediaFileUploadError::BadInput(
"Cannot set `is_intermediate_system_file` to true if `maybe_prompt_token` is provided."
.to_string()));
}
Ok(())
}
@@ -206,6 +206,7 @@ pub async fn upload_new_engine_asset_media_file_handler(
creator_ip_address: &ip_address,
creator_set_visibility: file_info.creator_set_visibility,
upload_type: UploadType::Filesystem,
maybe_prompt_token: None,
maybe_engine_category: Some(file_info.engine_category),
maybe_animation_type: file_info.maybe_animation_type,
maybe_mime_type: Some(file_info.mimetype),
@@ -230,6 +230,7 @@ pub async fn upload_new_scene_media_file_handler(
maybe_creator_anonymous_visitor_token: maybe_avt_token.as_ref(),
creator_ip_address: &ip_address,
creator_set_visibility,
maybe_prompt_token: None,
upload_type: UploadType::Filesystem,
maybe_engine_category: Some(MediaFileEngineCategory::Scene),
maybe_animation_type: None,
@@ -264,6 +264,7 @@ pub async fn upload_pmx_media_file_handler(
upload_type: UploadType::Filesystem,
maybe_engine_category: Some(engine_category),
maybe_animation_type,
maybe_prompt_token: None,
maybe_mime_type: Some("application/octet-stream"),
file_size_bytes: pmx_details.file_size_bytes,
maybe_duration_millis,
@@ -224,6 +224,7 @@ pub async fn upload_scene_snapshot_media_file_handler(
creator_ip_address: &ip_address,
creator_set_visibility: Visibility::Hidden,
upload_type: UploadType::Filesystem,
maybe_prompt_token: None,
maybe_engine_category: Some(MediaFileEngineCategory::Scene),
maybe_animation_type: None,
maybe_mime_type: Some(MIMETYPE),
@@ -371,6 +371,7 @@ pub async fn upload_new_video_media_file_handler(
upload_type: UploadType::Filesystem, // TODO(bt,2024-05-02): This should be a parameter and a well-known enum.
maybe_engine_category: None,
maybe_animation_type: None,
maybe_prompt_token: None,
maybe_mime_type: Some(&mimetype),
file_size_bytes: file_size_bytes as u64,
maybe_duration_millis,
@@ -227,6 +227,7 @@ pub async fn upload_video_media_file_handler(
creator_set_visibility,
upload_type,
maybe_engine_category: None,
maybe_prompt_token: None,
maybe_animation_type: None,
maybe_mime_type: Some(mimetype),
file_size_bytes: file_size_bytes as u64,