diff --git a/crates/api_clients/grok_client/src/credentials/grok_full_credentials.rs b/crates/api_clients/grok_client/src/credentials/grok_full_credentials.rs index 6afdfa676e..0b37a18a2d 100644 --- a/crates/api_clients/grok_client/src/credentials/grok_full_credentials.rs +++ b/crates/api_clients/grok_client/src/credentials/grok_full_credentials.rs @@ -1,5 +1,6 @@ use crate::credentials::grok_client_secrets::GrokClientSecrets; use crate::credentials::grok_cookies::GrokCookies; +use crate::datatypes::api::user_id::UserId; /// Includes the cookies and all the magic other bits needed to call the API. #[derive(Clone)] @@ -22,4 +23,8 @@ impl GrokFullCredentials { client_secrets, } } + + pub fn get_user_id_ref(&self) -> &UserId { + &self.client_secrets.user_id + } } diff --git a/crates/api_clients/grok_client/src/error/grok_client_error.rs b/crates/api_clients/grok_client/src/error/grok_client_error.rs index 2efa87fdee..253396101f 100644 --- a/crates/api_clients/grok_client/src/error/grok_client_error.rs +++ b/crates/api_clients/grok_client/src/error/grok_client_error.rs @@ -26,6 +26,12 @@ pub enum GrokClientError { /// Can't read a local file for uploading. CannotReadLocalFileForUpload(std::io::Error), + /// Couldn't write to the file + CannotOpenLocalFileForWriting(std::io::Error), + + /// Couldn't write to the file + CannotWriteLocalFile(std::io::Error), + /// File for upload has an invalid path. FileForUploadHasInvalidPath, @@ -96,6 +102,8 @@ impl Display for GrokClientError { Self::WebsocketSendError(err) => write!(f, "Websocket send error: {}", err), Self::CannotOpenLocalFileForUpload(err) => write!(f, "Cannot open local file for upload: {}", err), Self::CannotReadLocalFileForUpload(err) => write!(f, "Cannot read local file for upload: {}", err), + Self::CannotOpenLocalFileForWriting(err) => write!(f, "Cannot open local file for writing: {}", err), + Self::CannotWriteLocalFile(err) => write!(f, "Cannot write local file: {}", err), Self::FileForUploadHasInvalidPath => write!(f, "File for upload has invalid path"), Self::HtmlParsingError => write!(f, "Html parsing error"), Self::FailedToDecodeVerificationToken(err) => write!(f, "Failed to decode verification token: {}", err), diff --git a/crates/api_clients/grok_client/src/requests/download_video_file/grok_download_video.rs b/crates/api_clients/grok_client/src/requests/download_video_file/grok_download_video.rs new file mode 100644 index 0000000000..ea415399e2 --- /dev/null +++ b/crates/api_clients/grok_client/src/requests/download_video_file/grok_download_video.rs @@ -0,0 +1,171 @@ +use crate::client::browser_user_agents::FIREFOX_143_MAC_USER_AGENT; +use crate::datatypes::api::file_id::FileId; +use crate::datatypes::api::user_id::UserId; +use crate::error::grok_client_error::GrokClientError; +use crate::error::grok_error::GrokError; +use crate::error::grok_generic_api_error::GrokGenericApiError; +use crate::utils::user_and_file_id_to_video_url::user_and_file_id_to_video_url; +use futures::{StreamExt, TryStreamExt}; +use log::{error, info, warn}; +use std::io::Write; +use std::path::Path; +use std::time::Duration; +use tokio::fs::File; +use tokio::io::AsyncWriteExt; +use wreq::header::{ACCEPT, ACCEPT_ENCODING, ACCEPT_LANGUAGE, CACHE_CONTROL, CONNECTION, CONTENT_TYPE, COOKIE, ORIGIN, REFERER, TE, UPGRADE_INSECURE_REQUESTS, USER_AGENT}; +use wreq::{Client, Request, Response}; +use wreq_util::Emulation; + +pub struct GrokDownloadVideo<'a> { + pub cookies: &'a str, + pub user_id: &'a UserId, + pub file_id: &'a FileId, + pub request_timeout: Option, +} + +pub struct DownloadedVideoBytes { + pub bytes: Vec, +} + +impl <'a> GrokDownloadVideo<'a> { + + pub async fn download_to_path>(&self, path: P) -> Result<(), GrokError> { + let response = self.send_request().await?; + + let mut stream = response.bytes_stream(); + + let mut file = File::create(&path) + .await + .map_err(|err| GrokClientError::CannotOpenLocalFileForWriting(err))?; + + //let _r = stream.for_each_concurrent(None, |chunk_result| async { + // match chunk_result { + // Ok(chunk) => { + // file; + // //if let Err(e) = file.write_all(&chunk).await { + // // //error!("Error writing to file: {}", e); + // //} + // } + // Err(e) => { + // //error!("Error in stream: {}", e); + // } + // } + //}).await; + + while let Some(chunk) = stream.try_next() + .await + .map_err(|err| GrokGenericApiError::WreqError(err))? + { + file.write_all(&chunk) + .await + .map_err(|err| GrokClientError::CannotOpenLocalFileForWriting(err))? + } + + file.flush().await.unwrap(); + + Ok(()) + } + + pub async fn download_bytes(&self) -> Result { + let response = self.send_request().await?; + + let download_bytes = response.bytes() + .await + .map_err(|err| { + error!("Error reading Grok create media response body: {:?}", err); + GrokGenericApiError::WreqError(err) + })?; + + Ok(DownloadedVideoBytes { + bytes: download_bytes.to_vec(), + }) + } + + async fn send_request(&self) -> Result { + let client = Client::builder() + .emulation(Emulation::Firefox143) + .build() + .map_err(|err| GrokClientError::WreqClientError(err))?; + + let video_url = user_and_file_id_to_video_url(&self.user_id, &self.file_id, false); + + info!("Video file URL: {}", video_url); + + let mut request_builder = client.get(&video_url) + .header(USER_AGENT, FIREFOX_143_MAC_USER_AGENT) + .header(ACCEPT, "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8") + .header(ACCEPT_LANGUAGE, "en-US,en;q=0.5") + .header(ACCEPT_ENCODING, "gzip, deflate, br, zstd") + .header(REFERER, "https://grok.com/") + .header("Sec-GPC", "1") + .header(CONNECTION, "keep-alive") + .header(COOKIE, self.cookies.to_string()) + .header(UPGRADE_INSECURE_REQUESTS, "1") + .header("sec-fetch-dest", "document") + .header("sec-fetch-mode", "navigate") + .header("sec-fetch-site", "same-site") + .header("sec-fetch-user", "?1") + .header("priority", "u=0, i"); + + if let Some(timeout) = self.request_timeout { + request_builder = request_builder.timeout(timeout); + } + + let response = request_builder.send() + .await + .map_err(|err| GrokGenericApiError::WreqError(err))?; + + let status = response.status(); + + // TODO: Classify errors + if !status.is_success() { + warn!("Not successful liking media (status: {:?})", status); + // error!("Upload file request returned an error (code {}) : {:?}", status.as_u16(), response_body); + // return Err(classify_general_http_status_code_and_body(status, response_body)); + } + + Ok(response) + } +} + +#[cfg(test)] +mod tests { + use crate::datatypes::api::file_id::FileId; + use crate::recipes::request_client_secrets::{request_client_secrets, RequestClientSecretsArgs}; + use crate::requests::download_video_file::download_video_file::{download_video_file, DownloadVideoFileArgs}; + use crate::test_utils::get_test_cookies::get_typed_test_cookies; + use errors::AnyhowResult; + + #[tokio::test] + #[ignore] // Client side tests only + async fn test_download() -> AnyhowResult<()> { + //setup_test_logging(LevelFilter::Trace); + + let cookies = get_typed_test_cookies()?; + + let file_id = FileId("d9b300c6-9562-4e24-a87a-3ede2a53f0bc".to_string()); // Ernest + + let secrets = request_client_secrets(RequestClientSecretsArgs { + cookies: &cookies, + }).await?; + + println!("Verification Token: {:?}", secrets.verification_token); + println!("Sentry Trace: {:?}", secrets.sentry_trace); + println!("Numbers: {:?}", secrets.numbers); + println!("Svg Path: {:?}", secrets.svg_path_data); + println!("Baggage: {:?}", secrets.baggage); + + // TDOO: This implementation is untested. + //let download = download_video_file(DownloadVideoFileArgs { + // cookies: cookies.as_str(), + // user_id: &secrets.user_id, + // file_id: &file_id, + // request_timeout: None, + //}).await?; + + //println!("Video byte length: {:?}", download.bytes.len()); + + assert_eq!(1, 2); + Ok(()) + } +} \ No newline at end of file diff --git a/crates/api_clients/grok_client/src/requests/download_video_file/mod.rs b/crates/api_clients/grok_client/src/requests/download_video_file/mod.rs index a078360084..bb83484e8a 100644 --- a/crates/api_clients/grok_client/src/requests/download_video_file/mod.rs +++ b/crates/api_clients/grok_client/src/requests/download_video_file/mod.rs @@ -1 +1,2 @@ -mod download_video_file; \ No newline at end of file +pub mod download_video_file; +pub mod grok_download_video; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/grok/handle_grok_video.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/grok/handle_grok_video.rs index 88e7e0d3af..ab798a0e77 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/grok/handle_grok_video.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_video/grok/handle_grok_video.rs @@ -113,6 +113,7 @@ async fn get_grok_creds(app: &AppHandle, grok_credential_manager: &GrokCredentia Ok(secrets) => { let full_creds = GrokFullCredentials::from_cookies_and_client_secrets(cookies, secrets); grok_credential_manager.replace_full_credentials(full_creds.clone())?; + grok_credential_manager.persist_to_disk()?; return Ok(full_creds) } } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs index 4ea87ff7c7..2d3297f03e 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs @@ -68,6 +68,7 @@ impl TaskEnqueueSuccess{ Some(GenerationModel::Gemini25Flash) => Some(TaskModelType::Gemini25Flash), Some(GenerationModel::GptImage1) => Some(TaskModelType::GptImage1), Some(GenerationModel::Recraft3) => Some(TaskModelType::Recraft3), + Some(GenerationModel::GrokVideo) => Some(TaskModelType::GrokVideo), Some(GenerationModel::Kling21Pro) => Some(TaskModelType::Kling21Pro), Some(GenerationModel::Kling21Master) => Some(TaskModelType::Kling21Master), Some(GenerationModel::Seedance10Lite) => Some(TaskModelType::Seedance10Lite), diff --git a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs b/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs index 5464035eb4..967c3c0b60 100644 --- a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs +++ b/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs @@ -10,6 +10,7 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::artcraft_platform_info::ArtcraftPlatformInfo; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::threads::grok_video_task_polling::grok_video_task_polling_thread::grok_video_task_polling_thread; use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; use crate::services::midjourney::threads::midjourney_long_polling_thread::midjourney_long_polling_thread; use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; @@ -67,6 +68,15 @@ pub async fn handle_tauri_startup( &sora_task_queue, )?; + tauri::async_runtime::spawn(grok_video_task_polling_thread( + app.clone(), + app_env_configs.clone(), + root.clone(), + task_database.clone(), + grok_creds_manager.clone(), + storyteller_creds_manager.clone(), + )); + tauri::async_runtime::spawn(midjourney_long_polling_thread( app.clone(), app_env_configs.clone(), diff --git a/crates/desktop/artcraft/src/services/grok/mod.rs b/crates/desktop/artcraft/src/services/grok/mod.rs index b94e311046..285ecaa7e1 100644 --- a/crates/desktop/artcraft/src/services/grok/mod.rs +++ b/crates/desktop/artcraft/src/services/grok/mod.rs @@ -1,4 +1,5 @@ pub mod commands; pub mod state; +pub mod threads; pub mod util; pub mod windows; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs b/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs index fad90684f6..4c2cf9818e 100644 --- a/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs +++ b/crates/desktop/artcraft/src/services/grok/state/grok_credential_manager.rs @@ -131,6 +131,11 @@ impl GrokCredentialManager { } } + // NB: This is just a heuristic. We'll add better checks later. + pub fn do_task_polling(&self) -> anyhow::Result { + self.session_appears_active() + } + // NB: This is just a heuristic. We'll add better checks later. pub fn session_appears_active(&self) -> anyhow::Result { let holder = match self.credential_data.read() { diff --git a/crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/grok_video_task_polling_thread.rs b/crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/grok_video_task_polling_thread.rs new file mode 100644 index 0000000000..839a22c78e --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/grok_video_task_polling_thread.rs @@ -0,0 +1,361 @@ +use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; +use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::data_dir::trait_data_subdir::DataSubdir; +use crate::core::state::task_database::TaskDatabase; +use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES; +use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +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 artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template; +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_media_file_class::TaskMediaFileClass; +use enums::tauri::tasks::task_status::TaskStatus; +use errors::AnyhowResult; +use grok_client::credentials::grok_full_credentials::GrokFullCredentials; +use grok_client::requests::download_video_file::download_video_file::{download_video_file, DownloadVideoFileArgs}; +use grok_client::requests::download_video_file::grok_download_video::GrokDownloadVideo; +use grok_client::requests::media_posts::list_media_posts::grok_list_media_posts::{GrokMediaPostList, GrokMediaPostListRequest, VideoData}; +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; +use midjourney_client::endpoints::imagine::{imagine, ImagineItem, ImagineRequest, MidjourneyJobType}; +use midjourney_client::utils::get_image_url::get_image_url; +use midjourney_client::utils::image_downloader_client::ImageDownloaderClient; +use once_cell::sync::Lazy; +use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList}; +use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs}; +use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; +use std::collections::{HashMap, HashSet}; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use std::time::Duration; +use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; +use storyteller_client::endpoints::media_files::get_media_file::get_media_file; +use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs}; +use storyteller_client::endpoints::media_files::upload_video_media_file_from_file::{upload_video_media_file_from_file, UploadVideoFromFileArgs}; +use storyteller_client::endpoints::prompts::create_prompt::create_prompt; +use storyteller_client::error::api_error::ApiError; +use storyteller_client::error::storyteller_error::StorytellerError; +use tauri::AppHandle; +use tokens::tokens::batch_generations::BatchGenerationToken; +use url::Url; +use grok_client::error::grok_error::GrokError; + +pub async fn grok_video_task_polling_thread( + app_handle: AppHandle, + app_env_configs: AppEnvConfigs, + app_data_root: AppDataRoot, + task_database: TaskDatabase, + creds: GrokCredentialManager, + storyteller_creds_manager: StorytellerCredentialManager, +) -> ! { + loop { + let res = polling_loop( + &app_handle, + &app_env_configs, + &app_data_root, + &task_database, + &creds, + &storyteller_creds_manager, + ).await; + if let Err(err) = res { + error!("An error occurred: {:?}", err); + } + // NB: Only sleep if an error occurs. + tokio::time::sleep(std::time::Duration::from_millis(30_000)).await; + } +} + +async fn polling_loop( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + grok_creds: &GrokCredentialManager, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> AnyhowResult<()> { + loop { + if !grok_creds.do_task_polling()? { + tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; + continue; + } + + // TODO: Graceful wait, fix this long function body + let storyteller_creds = match storyteller_creds_manager.get_credentials()? { + Some(creds) => creds, + None => { + error!("No Storyteller credentials found. Cannot proceed with Grok polling."); + tokio::time::sleep(std::time::Duration::from_millis(5_000)).await; + continue; + } + }; + + let grok_full_creds = match grok_creds.maybe_copy_full_credentials()? { + Some(full_creds) => full_creds, + None => { + tokio::time::sleep(std::time::Duration::from_millis(30_000)).await; + continue; + } + }; + + + let local_tasks = list_tasks_by_provider_and_status(ListTasksByProviderAndStatusArgs { + db: task_database.get_connection(), + provider: GenerationProvider::Grok, + task_statuses: &TASK_DATABASE_PENDING_STATUSES, + }).await?; + + poll_grok_tasks( + app_handle, + app_env_configs, + app_data_root, + task_database, + &grok_full_creds, + &storyteller_creds, + local_tasks, + ).await?; + + tokio::time::sleep(std::time::Duration::from_millis(2_000)).await; + } +} + +async fn poll_grok_tasks( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + grok_full_creds: &GrokFullCredentials, + storyteller_creds: &StorytellerCredentialSet, + local_tasks: TaskList, +) -> AnyhowResult<()> { + let local_tasks = local_tasks.tasks; + + if local_tasks.is_empty() { + return Ok(()) + } + + // Map of Grok Post ID to Local Task. + let local_tasks_by_grok_post_id = local_tasks.iter() + .filter_map(|task| { + if let Some(provider_job_id) = &task.provider_job_id { + Some((provider_job_id.clone(), task.clone())) + } else { + None + } + }) + .collect::>(); + + let list_media_request = GrokMediaPostListRequest { + cookie: grok_full_creds.cookies.as_str(), + cursor: None, + request_timeout: Some(Duration::from_millis(20_000)), + }; + + let list_result = list_media_request.send().await?; + + let grok_posts = list_result.posts; + + let grok_video_posts_by_id = { + let mut hash = HashMap::new(); + for post in grok_posts.iter() { + if let Some(video_data) = &post.video_data { + hash.insert(post.post_id.to_string(), video_data.clone()); + } + } + hash + }; + + for (grok_post_id, local_task) in local_tasks_by_grok_post_id.iter() { + // TODO: Copy prompt from this. + let grok_video_data = match grok_video_posts_by_id.get(grok_post_id) { + Some(video_data) => video_data, + None => continue, + }; + + upload_midjourney_batch( + &app_handle, + &app_env_configs, + app_data_root, + task_database, + &storyteller_creds, + grok_full_creds, + grok_post_id, + &local_task, + grok_video_data + ).await?; + + tokio::time::sleep(std::time::Duration::from_millis(2_000)).await; + } + + tokio::time::sleep(std::time::Duration::from_millis(60_000)).await; + + Ok(()) +} + +async fn upload_midjourney_batch( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds: &StorytellerCredentialSet, + grok_full_creds: &GrokFullCredentials, + grok_post_id: &String, + local_task: &Task, + grok_video_post: &VideoData +) -> AnyhowResult<()> { + + let request = CreatePromptRequest { + uuid_idempotency_token: generate_random_uuid(), + positive_prompt: grok_video_post.prompt.clone(), + negative_prompt: None, + model_type: Some(ModelType::Midjourney), // TODO: TEMP + generation_provider: Some(GenerationProvider::Midjourney), // TODO: TEMP + }; + + let prompt_response = create_prompt( + &app_env_configs.storyteller_host, + Some(storyteller_creds), + request + ).await?; + + info!("Created prompt: {:?}", &prompt_response.prompt_token); + + let mut maybe_primary_media_file_token = None; + + info!("Downloading generated Grok video ..."); + + let tempdir = app_data_root.temp_dir().path(); + let download_filename = format!("{}.mp4", grok_post_id); + let download_path = tempdir.join(download_filename); + + let download = GrokDownloadVideo { + cookies: grok_full_creds.cookies.as_str(), + user_id: grok_full_creds.get_user_id_ref(), + file_id: &grok_video_post.file_id, + request_timeout: None, + }; + + download.download_to_path(&download_path).await?; + + let mut wait_delay = 0; + + loop { + info!("Uploading to backend..."); + + // TODO: media_files.origin_category + // TODO: media_files.maybe_prompt_token + // TODO: media_files.maybe_generation_provider + // TODO: media_files.maybe_origin_model_type + // TODO: media_files.maybe_origin_model_token (sref?) + // TODO: media_files.is_batch_generated + // TODO: media_files.maybe_batch_token + // TODO: media_files.is_user_upload + + // TODO: batch_generations.token + // TODO: batch_generations.entity_type + // TODO: batch_generations.entity_token + + let result = upload_video_media_file_from_file(UploadVideoFromFileArgs { + api_host: &app_env_configs.storyteller_host, + maybe_creds: Some(&storyteller_creds), + path: &download_path, + maybe_prompt_token: Some(&prompt_response.prompt_token), + }).await; + + match result { + Ok(result) => { + info!("Successfully uploaded to backend: {:?}", result.media_file_token); + if maybe_primary_media_file_token.is_none() { + maybe_primary_media_file_token = Some(result.media_file_token); + } + break; + }, + Err(ApiError::TooManyRequests(_)) => { + error!("Too many requests, retrying upload after delay..."); + // If we hit a rate limit, we can retry after a short delay. + wait_delay += 10; + if wait_delay > 60 { + wait_delay = 60; + } + tokio::time::sleep(std::time::Duration::from_secs(wait_delay)).await; + continue; // Retry the upload. + } + Err(err) => { + error!("Failed to upload to backend: {:?}", err); + return Err(err.into()) + }, + } + } // End loop + + let mut maybe_cdn_url = None; + let mut maybe_thumbnail_url_template = None; + + if let Some(media_file_token) = maybe_primary_media_file_token.as_ref() { + info!("Looking up file to grab CDN and thumbnail URLs: {:?} ...", media_file_token); + + let lookup_result = get_media_file( + &app_env_configs.storyteller_host, + media_file_token, + ).await; + match lookup_result { + Ok(response) => { + maybe_cdn_url = Some(response.media_file.media_links.cdn_url.to_string()); + maybe_thumbnail_url_template = media_links_to_thumbnail_template(&response.media_file.media_links) + .map(|s| s.to_string()); + } + Err(err) => { + error!("Failed to look up media file after upload: {:?} (failing open)", err); + } + } + } + + let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs { + db: task_database.get_connection(), + task_id: &local_task.id, + maybe_batch_token: None, + maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(), + maybe_primary_media_file_class: Some(TaskMediaFileClass::Video), + maybe_primary_media_file_thumbnail_url_template: maybe_thumbnail_url_template.as_deref(), + maybe_primary_media_file_cdn_url: maybe_cdn_url.as_deref(), + }).await?; + + if !updated { + return Ok(()); // If anything breaks with queries, don't spam events. + } + + let event = GenerationCompleteEvent { + //media_file_token: result.media_file_token, + action: Some(GenerationAction::GenerateVideo), + service: GenerationServiceProvider::Midjourney, // TODO + model: None, + }; + + if let Err(err) = event.send(&app_handle) { + error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open + } + + //let result = maybe_handle_text_to_image_complete_event( + // app_handle, + // app_env_configs, + // Some(storyteller_creds), + // local_task, + // &batch_token, + //).await; + + //if let Err(err) = result { + // error!("Failed to send text-to-image complete event: {:?}", err); + //} + + Ok(()) +} + diff --git a/crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/mod.rs b/crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/mod.rs new file mode 100644 index 0000000000..8a7f640414 --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/mod.rs @@ -0,0 +1 @@ +pub mod grok_video_task_polling_thread; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/grok/threads/mod.rs b/crates/desktop/artcraft/src/services/grok/threads/mod.rs new file mode 100644 index 0000000000..5239c8592a --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/threads/mod.rs @@ -0,0 +1 @@ +pub mod grok_video_task_polling; \ No newline at end of file diff --git a/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs b/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs index c7cd1cda3b..93b00b6c9b 100644 --- a/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs +++ b/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs @@ -39,6 +39,8 @@ pub enum TaskModelType { Midjourney, // Video models + #[serde(rename = "grok_video")] + GrokVideo, // Video version unspecified/unknown #[serde(rename = "kling_1.6_pro")] Kling16Pro, #[serde(rename = "kling_2.1_pro")] @@ -85,6 +87,7 @@ impl TaskModelType { Self::Recraft3 => "recraft_3", Self::Midjourney => "midjourney", // Video models + Self::GrokVideo => "grok_video", Self::Kling16Pro => "kling_1.6_pro", Self::Kling21Pro => "kling_2.1_pro", Self::Kling21Master => "kling_2.1_master", @@ -114,6 +117,7 @@ impl TaskModelType { "recraft_3" => Ok(Self::Recraft3), "midjourney" => Ok(Self::Midjourney), // Video models + "grok_video" => Ok(Self::GrokVideo), "kling_1.6_pro" => Ok(Self::Kling16Pro), "kling_2.1_pro" => Ok(Self::Kling21Pro), "kling_2.1_master" => Ok(Self::Kling21Master), @@ -146,6 +150,7 @@ impl TaskModelType { Self::Recraft3, Self::Midjourney, // Video models + Self::GrokVideo, Self::Kling16Pro, Self::Kling21Pro, Self::Kling21Master, @@ -185,6 +190,7 @@ mod tests { assert_serialization(TaskModelType::Recraft3, "recraft_3"); assert_serialization(TaskModelType::Midjourney, "midjourney"); // Video models + assert_serialization(TaskModelType::GrokVideo, "grok_video"); assert_serialization(TaskModelType::Kling16Pro, "kling_1.6_pro"); assert_serialization(TaskModelType::Kling21Pro, "kling_2.1_pro"); assert_serialization(TaskModelType::Kling21Master, "kling_2.1_master"); @@ -213,6 +219,7 @@ mod tests { assert_eq!(TaskModelType::Recraft3.to_str(), "recraft_3"); assert_eq!(TaskModelType::Midjourney.to_str(), "midjourney"); // Video models + assert_eq!(TaskModelType::GrokVideo.to_str(), "grok_video"); assert_eq!(TaskModelType::Kling16Pro.to_str(), "kling_1.6_pro"); assert_eq!(TaskModelType::Kling21Pro.to_str(), "kling_2.1_pro"); assert_eq!(TaskModelType::Kling21Master.to_str(), "kling_2.1_master"); @@ -241,6 +248,7 @@ mod tests { assert_eq!(TaskModelType::from_str("recraft_3").unwrap(), TaskModelType::Recraft3); assert_eq!(TaskModelType::from_str("midjourney").unwrap(), TaskModelType::Midjourney); // Video models + assert_eq!(TaskModelType::from_str("grok_video").unwrap(), TaskModelType::GrokVideo); assert_eq!(TaskModelType::from_str("kling_1.6_pro").unwrap(), TaskModelType::Kling16Pro); assert_eq!(TaskModelType::from_str("kling_2.1_pro").unwrap(), TaskModelType::Kling21Pro); assert_eq!(TaskModelType::from_str("kling_2.1_master").unwrap(), TaskModelType::Kling21Master); @@ -268,7 +276,7 @@ mod tests { #[test] fn all_variants() { let mut variants = TaskModelType::all_variants(); - assert_eq!(variants.len(), 21); + assert_eq!(variants.len(), 22); // Image models assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Dev)); assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Schnell)); @@ -282,6 +290,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(TaskModelType::Recraft3)); assert_eq!(variants.pop_first(), Some(TaskModelType::Midjourney)); // Video models + assert_eq!(variants.pop_first(), Some(TaskModelType::GrokVideo)); assert_eq!(variants.pop_first(), Some(TaskModelType::Kling16Pro)); assert_eq!(variants.pop_first(), Some(TaskModelType::Kling21Pro)); assert_eq!(variants.pop_first(), Some(TaskModelType::Kling21Master));