diff --git a/Cargo.lock b/Cargo.lock index 249c9b1f47..e28a311f3c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -694,6 +694,7 @@ dependencies = [ "tokio", "tokio-util", "url", + "url_utils", "web_base64", "workspace-build-acceleration", ] @@ -12457,6 +12458,14 @@ dependencies = [ "workspace-build-acceleration", ] +[[package]] +name = "url_utils" +version = "0.0.1" +dependencies = [ + "url", + "workspace-build-acceleration", +] + [[package]] name = "urlencoding" version = "2.1.3" diff --git a/Cargo.toml b/Cargo.toml index 6b7115c0ef..fffb8137b8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -84,6 +84,7 @@ members = [ "crates/lib/testing", "crates/lib/tts_common", "crates/lib/url_config", + "crates/lib/url_utils", "crates/lib/user_input_common", "crates/lib/web_scrapers", #"crates/lib/api_clients/fakeyou_client", @@ -210,6 +211,7 @@ subprocess_common = { path = "crates/lib/subprocess_common" } test_data = { path = "crates/lib/test_data" } testing = { path = "crates/lib/testing" } tokens = { path = "crates/schema/public/tokens" } +url_utils = { path = "crates/lib/url_utils" } web_base64 = { path = "crates/lib/data/web_base64" } # External libs diff --git a/crates/desktop/artcraft/Cargo.toml b/crates/desktop/artcraft/Cargo.toml index 37a51880ac..81025533cc 100644 --- a/crates/desktop/artcraft/Cargo.toml +++ b/crates/desktop/artcraft/Cargo.toml @@ -42,6 +42,7 @@ primitives.workspace = true sqlite_tasks.workspace = true storyteller_client.workspace = true tokens.workspace = true +url_utils.workspace = true web_base64.workspace = true # External diff --git a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/download_extension.rs b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/download_extension.rs new file mode 100644 index 0000000000..a6841377a0 --- /dev/null +++ b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/download_extension.rs @@ -0,0 +1,15 @@ + +#[derive(Clone, Copy)] +pub enum DownloadExtension { + Png, + Mp4, +} + +impl DownloadExtension { + pub fn as_extension_without_period(&self) -> &'static str { + match self { + DownloadExtension::Png => "png", + DownloadExtension::Mp4 => "mp4", + } + } +} diff --git a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs index a8af5529ff..0695d39ada 100644 --- a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs +++ b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs @@ -9,6 +9,7 @@ use crate::core::state::task_database::TaskDatabase; use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES; use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; use crate::services::sora::state::sora_task_queue::SoraTaskQueue; +use crate::services::sora::threads::sora_task_polling::helpers::download_extension::DownloadExtension; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest; use enums::common::generation_provider::GenerationProvider; @@ -35,6 +36,7 @@ use storyteller_client::endpoints::media_files::upload_video_media_file_from_fil use storyteller_client::endpoints::prompts::create_prompt::create_prompt; use tauri::AppHandle; use tempdir::TempDir; +use url_utils::extension::extract_download_extension_from_url::{extract_download_extension_from_url, extract_download_extension_from_url_str}; pub struct SuccessfulGeneration { pub prompt: Option, @@ -61,6 +63,7 @@ pub async fn handle_classic_successful_generations( storyteller_creds: &StorytellerCredentialSet, succeeded_tasks_by_id: &HashMap, sqlite_tasks_by_sora_task_id: &HashMap, + recommended_download_extension: DownloadExtension, ) -> AnyhowResult<()> { for (task_id, generation) in succeeded_tasks_by_id.iter() { @@ -98,7 +101,7 @@ pub async fn handle_classic_successful_generations( for (_i, item) in generation.items.iter().enumerate() { info!("Downloading generated file..."); - let download_path = download_generation_item(item, &app_data_root).await?; + let download_path = download_generation_item(item, &app_data_root, recommended_download_extension).await?; info!("Uploading to backend..."); @@ -158,18 +161,26 @@ pub async fn handle_classic_successful_generations( } -async fn download_generation_item(generation: &GenerationItem, app_data_root: &AppDataRoot) -> AnyhowResult { - let url = Url::parse(&generation.url)?; +async fn download_generation_item( + generation: &GenerationItem, + app_data_root: &AppDataRoot, + recommended_download_extension: DownloadExtension +) -> AnyhowResult { + info!("Downloading generation item from URL: {}", generation.url.as_str()); let response = reqwest::get(&generation.url).await?; let image_bytes = response.bytes().await?; - let ext = url.path().split(".").last().unwrap_or("png"); - + let extension = extract_download_extension_from_url_str(&generation.url) + .map(|ext| ext.as_extension_without_period()) + .unwrap_or_else(|| recommended_download_extension.as_extension_without_period()); + let tempdir = app_data_root.temp_dir().path(); - let download_filename = format!("{}.{}", generation.item_id, ext); + let download_filename = format!("{}.{}", generation.item_id, extension); let download_path = tempdir.join(download_filename); + info!("Writing to path: {:?}", download_path); + let mut file = File::create(&download_path)?; file.write_all(&image_bytes)?; diff --git a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/mod.rs b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/mod.rs index 09a7aef4fc..d147b37680 100644 --- a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/mod.rs +++ b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/mod.rs @@ -2,3 +2,4 @@ pub (super) mod handle_successful_generations; pub (super) mod handle_failed_generations; pub (super) mod poll_classic_sora_tasks; pub (super) mod poll_sora_2_tasks; +mod download_extension; diff --git a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_classic_sora_tasks.rs b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_classic_sora_tasks.rs index e0e20167dc..94e846218a 100644 --- a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_classic_sora_tasks.rs +++ b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_classic_sora_tasks.rs @@ -36,6 +36,7 @@ use storyteller_client::endpoints::media_files::upload_image_media_file_from_fil use storyteller_client::endpoints::prompts::create_prompt::create_prompt; use tauri::AppHandle; use tempdir::TempDir; +use crate::services::sora::threads::sora_task_polling::helpers::download_extension::DownloadExtension; pub async fn poll_classic_sora_tasks( app_handle: &AppHandle, @@ -117,6 +118,7 @@ pub async fn poll_classic_sora_tasks( &storyteller_creds, &sora_succeeded_tasks_by_id, &local_sqlite_tasks_by_sora_task_id, + DownloadExtension::Png, ).await?; Ok(()) diff --git a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_sora_2_tasks.rs b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_sora_2_tasks.rs index 2c515d831c..79d7c29fd3 100644 --- a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_sora_2_tasks.rs +++ b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_sora_2_tasks.rs @@ -38,6 +38,7 @@ use storyteller_client::endpoints::media_files::upload_image_media_file_from_fil use storyteller_client::endpoints::prompts::create_prompt::create_prompt; use tauri::AppHandle; use tempdir::TempDir; +use crate::services::sora::threads::sora_task_polling::helpers::download_extension::DownloadExtension; pub async fn poll_sora_2_tasks( app_handle: &AppHandle, @@ -110,6 +111,7 @@ pub async fn poll_sora_2_tasks( &storyteller_creds, &sora_succeeded_drafts_by_id, &local_sqlite_tasks_by_sora_task_id, + DownloadExtension::Mp4, ).await?; Ok(()) diff --git a/crates/lib/url_utils/Cargo.toml b/crates/lib/url_utils/Cargo.toml new file mode 100644 index 0000000000..2c34d2b1f7 --- /dev/null +++ b/crates/lib/url_utils/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "url_utils" +edition = "2024" +version = "0.0.1" +publish = false + +[package.metadata.cargo-udeps.ignore] +normal = ["workspace-build-acceleration"] + +[lib] +name = "url_utils" +path = "src/lib.rs" + +[dependencies] + +# Workspace hack (faster compile times) +workspace-build-acceleration.workspace = true + +# Internal +# (None) + +# External +url.workspace = true + +[dev-dependencies] +# (None) diff --git a/crates/lib/url_utils/src/extension/download_extension.rs b/crates/lib/url_utils/src/extension/download_extension.rs new file mode 100644 index 0000000000..c24057b04b --- /dev/null +++ b/crates/lib/url_utils/src/extension/download_extension.rs @@ -0,0 +1,71 @@ + +/// When downloading a remote file, it makes sense to try to match the same file extension (with reasonable safeguards). +#[derive(Debug, Copy, Clone, Eq, PartialEq)] +pub enum DownloadExtension { + // Image + Jpg, + Png, + // Video + Mp4, + Webm, + // Audio + Mp3, + Wav, + // Text + Json, + Txt, +} + +impl DownloadExtension { + pub fn as_extension_without_period(&self) -> &'static str { + match self { + // Image + Self::Jpg => "jpg", + Self::Png => "png", + // Video + Self::Mp4 => "mp4", + Self::Webm => "webm", + // Audio + Self::Mp3 => "mp3", + Self::Wav => "wav", + // Text + Self::Json => "json", + Self::Txt => "txt", + } + } + + pub fn as_extension_with_period(&self) -> &'static str { + match self { + // Image + Self::Jpg => ".jpg", + Self::Png => ".png", + // Video + Self::Mp4 => ".mp4", + Self::Webm => ".webm", + // Audio + Self::Mp3 => ".mp3", + Self::Wav => ".wav", + // Text + Self::Json => ".json", + Self::Txt => ".txt", + } + } + + pub fn try_from_str(extension: &str) -> Option { + match extension.to_lowercase().as_str() { + // Image + "jpg" | "jpeg" => Some(Self::Jpg), + "png" => Some(Self::Png), + // Video + "mp4" => Some(Self::Mp4), + "webm" => Some(Self::Webm), + // Audio + "mp3" => Some(Self::Mp3), + "wav" => Some(Self::Wav), + // Text + "json" => Some(Self::Json), + "txt" => Some(Self::Txt), + _ => None, + } + } +} diff --git a/crates/lib/url_utils/src/extension/extract_download_extension_from_url.rs b/crates/lib/url_utils/src/extension/extract_download_extension_from_url.rs new file mode 100644 index 0000000000..a78cdba2cd --- /dev/null +++ b/crates/lib/url_utils/src/extension/extract_download_extension_from_url.rs @@ -0,0 +1,37 @@ +use crate::extension::download_extension::DownloadExtension; +use url::Url; + +/// This should be used for naming files on the local file system based on their URL. +pub fn extract_download_extension_from_url_str(url: &str) -> Option { + let parsed_url = Url::parse(url).ok()?; + extract_download_extension_from_url(&parsed_url) +} + +/// This should be used for naming files on the local file system based on their URL. +pub fn extract_download_extension_from_url(url: &Url) -> Option { + let path = url.path(); + let extension = std::path::Path::new(path) + .extension() + .and_then(|ext| ext.to_str())?; + DownloadExtension::try_from_str(extension) +} + +#[cfg(test)] +mod tests { + use crate::extension::download_extension::DownloadExtension; + use crate::extension::extract_download_extension_from_url::extract_download_extension_from_url_str; + + #[test] + fn success_case() { + let url = "https://storyteller.ai/example.jpg"; + let maybe_extension= extract_download_extension_from_url_str(url); + assert_eq!(maybe_extension, Some(DownloadExtension::Jpg)); + } + + #[test] + fn difficult_url() { + let url = "https://videos.youtube.com/foo/bar/files/00000000-f19a-1275-a092-caa60b38f8d2%2Fraw?se=2025-11-21T22%3A26%3A53Z&sp=r&sv=2024-08-09&sr=b&skoid=8cfff87a-01fc-27c1-9091-32999dcd1389&sktid=aa8ccaf6-e6da-484e-ab11-9c7496c21c53&skt=2025-11-17T21%3A09%3A08Z&ske=2025-10-23T21%3A14%3A08Z&sks=b&skv=2024-08-04&sig=T2U6%2Brw9WJdJagBkCL4/GlFr72KhZGP3QPlobIO7JlU%3D&ac=oaisdsorprnorthcentralus"; + let maybe_extension= extract_download_extension_from_url_str(url); + assert!(maybe_extension.is_none()); + } +} diff --git a/crates/lib/url_utils/src/extension/mod.rs b/crates/lib/url_utils/src/extension/mod.rs new file mode 100644 index 0000000000..e847b17dfd --- /dev/null +++ b/crates/lib/url_utils/src/extension/mod.rs @@ -0,0 +1,2 @@ +pub mod download_extension; +pub mod extract_download_extension_from_url; \ No newline at end of file diff --git a/crates/lib/url_utils/src/lib.rs b/crates/lib/url_utils/src/lib.rs new file mode 100644 index 0000000000..f339dde0fb --- /dev/null +++ b/crates/lib/url_utils/src/lib.rs @@ -0,0 +1 @@ +pub mod extension; \ No newline at end of file