From a511e0e0a4eb127a35a053f5cc3d7cc6cfdbf9c0 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 31 Oct 2025 21:55:45 -0400 Subject: [PATCH] grok images (#893) --- .env | 1 + ...09b73bdb81de4a08dcaa96a86afb5f9be7a68.json | 12 + ...261de1b5fa05a9787679ce9ccd4ac576b1069.json | 68 ++++ .../src/datatypes/api/request_id.rs | 6 + .../image_websocket/grok_websocket.rs | 108 ++++++ ...websocket.rs => grok_wrapped_websocket.rs} | 5 +- .../listen_for_websocket_images.rs | 14 +- .../listen_for_websocket_request_id.rs | 15 +- .../src/requests/image_websocket/mod.rs | 5 +- .../image_websocket/prompt_websocket_image.rs | 14 +- .../artcraft/src/core/artcraft_error.rs | 7 + .../core/commands/enqueue/generate_error.rs | 6 + .../commands/enqueue/task_enqueue_success.rs | 1 + .../enqueue_text_to_image_command.rs | 24 ++ .../generic/handle_image_artcraft.rs | 1 + .../enqueue/text_to_image/grok/handle_grok.rs | 138 +++++++ .../enqueue/text_to_image/grok/mod.rs | 1 + .../commands/enqueue/text_to_image/mod.rs | 3 +- .../core/events/generation_events/common.rs | 2 + .../lifecycle/startup/handle_tauri_startup.rs | 13 + crates/desktop/artcraft/src/lib.rs | 6 + .../grok/state/grok_image_prompt_queue.rs | 76 ++++ .../grok/state/grok_websocket_manager.rs | 62 +++ .../artcraft/src/services/grok/state/mod.rs | 6 +- .../grok_image_websocket_thread.rs | 364 ++++++++++++++++++ .../grok_image_websocket_thread/mod.rs | 1 + .../grok_video_task_polling_thread.rs | 3 +- .../artcraft/src/services/grok/threads/mod.rs | 3 +- ...ybe_handle_text_to_image_complete_event.rs | 3 +- .../threads/midjourney_long_polling_thread.rs | 5 +- .../helpers/handle_failed_generations.rs | 2 +- .../helpers/handle_successful_generations.rs | 3 +- .../helpers/poll_classic_sora_tasks.rs | 5 +- .../helpers/poll_sora_2_tasks.rs | 5 +- .../sora_task_polling_thread.rs | 7 +- .../maybe_handle_inpainting_complete_event.rs | 3 +- ...ybe_handle_text_to_image_complete_event.rs | 3 +- ..._send_background_removal_complete_event.rs | 3 +- .../threads/storyteller_activity_thread.rs | 8 +- .../storyteller_task_polling_thread.rs | 3 +- ...et_task_by_provider_and_provider_job_id.rs | 78 ++++ .../list_tasks_by_provider_and_status.rs | 28 +- .../list_tasks_by_provider_and_tokens.rs | 28 +- .../database/sqlite_tasks/src/queries/mod.rs | 3 + .../database/sqlite_tasks/src/queries/task.rs | 33 ++ ...l_task_status_with_metadata_by_provider.rs | 64 +++ .../enums/src/tauri/tasks/task_model_type.rs | 11 +- .../model-list/src/lib/lists/ImageModels.ts | 15 + 48 files changed, 1172 insertions(+), 103 deletions(-) create mode 100644 .sqlx/query-1e3758851fc8451670f86a9432009b73bdb81de4a08dcaa96a86afb5f9be7a68.json create mode 100644 .sqlx/query-38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069.json create mode 100644 crates/api_clients/grok_client/src/requests/image_websocket/grok_websocket.rs rename crates/api_clients/grok_client/src/requests/image_websocket/{clonable_websocket.rs => grok_wrapped_websocket.rs} (98%) create mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs create mode 100644 crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/mod.rs create mode 100644 crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs create mode 100644 crates/desktop/artcraft/src/services/grok/state/grok_websocket_manager.rs create mode 100644 crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs create mode 100644 crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/mod.rs create mode 100644 crates/schema/database/sqlite_tasks/src/queries/get_task_by_provider_and_provider_job_id.rs create mode 100644 crates/schema/database/sqlite_tasks/src/queries/task.rs create mode 100644 crates/schema/database/sqlite_tasks/src/queries/update_successful_task_status_with_metadata_by_provider.rs diff --git a/.env b/.env index f3c793b84a..4096df8647 100644 --- a/.env +++ b/.env @@ -1,6 +1,7 @@ SERVER_ENVIRONMENT=development DATABASE_URL=mysql://storyteller:password@localhost/storyteller +#DATABASE_URL=sqlite:/tmp/tasks.sqlite #DATABASE_URL=mysql://root:root@localhost/storyteller #DATABASE_URL=mysql://root@localhost/storyteller diff --git a/.sqlx/query-1e3758851fc8451670f86a9432009b73bdb81de4a08dcaa96a86afb5f9be7a68.json b/.sqlx/query-1e3758851fc8451670f86a9432009b73bdb81de4a08dcaa96a86afb5f9be7a68.json new file mode 100644 index 0000000000..5c1b1e3f1a --- /dev/null +++ b/.sqlx/query-1e3758851fc8451670f86a9432009b73bdb81de4a08dcaa96a86afb5f9be7a68.json @@ -0,0 +1,12 @@ +{ + "db_name": "SQLite", + "query": "\n UPDATE tasks\n SET\n task_status = ?,\n on_complete_batch_token = ?,\n on_complete_primary_media_file_token = ?,\n on_complete_primary_media_file_class = ?,\n on_complete_primary_media_file_cdn_url = ?,\n on_complete_primary_media_file_thumbnail_url_template = ?\n WHERE\n provider = ?\n AND provider_job_id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 8 + }, + "nullable": [] + }, + "hash": "1e3758851fc8451670f86a9432009b73bdb81de4a08dcaa96a86afb5f9be7a68" +} diff --git a/.sqlx/query-38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069.json b/.sqlx/query-38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069.json new file mode 100644 index 0000000000..b400c3bf00 --- /dev/null +++ b/.sqlx/query-38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069.json @@ -0,0 +1,68 @@ +{ + "db_name": "SQLite", + "query": "\n SELECT\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_caller,\n frontend_subscriber_id,\n frontend_subscriber_payload\n FROM tasks\n WHERE\n provider = ?\n AND provider_job_id = ?\n ", + "describe": { + "columns": [ + { + "name": "id", + "ordinal": 0, + "type_info": "Text" + }, + { + "name": "task_status", + "ordinal": 1, + "type_info": "Text" + }, + { + "name": "task_type", + "ordinal": 2, + "type_info": "Text" + }, + { + "name": "model_type", + "ordinal": 3, + "type_info": "Text" + }, + { + "name": "provider", + "ordinal": 4, + "type_info": "Text" + }, + { + "name": "provider_job_id", + "ordinal": 5, + "type_info": "Text" + }, + { + "name": "frontend_caller", + "ordinal": 6, + "type_info": "Text" + }, + { + "name": "frontend_subscriber_id", + "ordinal": 7, + "type_info": "Text" + }, + { + "name": "frontend_subscriber_payload", + "ordinal": 8, + "type_info": "Text" + } + ], + "parameters": { + "Right": 2 + }, + "nullable": [ + false, + false, + false, + true, + true, + true, + true, + true, + true + ] + }, + "hash": "38ff2df237b16f077865711dd81261de1b5fa05a9787679ce9ccd4ac576b1069" +} diff --git a/crates/api_clients/grok_client/src/datatypes/api/request_id.rs b/crates/api_clients/grok_client/src/datatypes/api/request_id.rs index ff0f796166..1e06ece469 100644 --- a/crates/api_clients/grok_client/src/datatypes/api/request_id.rs +++ b/crates/api_clients/grok_client/src/datatypes/api/request_id.rs @@ -5,3 +5,9 @@ /// Request IDs are UUIds. #[derive(Clone, Debug)] pub struct RequestId(pub String); + +impl RequestId { + pub fn to_string(&self) -> String { + self.0.clone() + } +} diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/grok_websocket.rs b/crates/api_clients/grok_client/src/requests/image_websocket/grok_websocket.rs new file mode 100644 index 0000000000..9780099bb2 --- /dev/null +++ b/crates/api_clients/grok_client/src/requests/image_websocket/grok_websocket.rs @@ -0,0 +1,108 @@ +use crate::error::grok_client_error::GrokClientError; +use crate::error::grok_error::GrokError; +use crate::requests::image_websocket::messages::websocket_server_message::WebsocketServerMessage; +use futures::stream::FusedStream; +use futures::{Stream, TryStreamExt}; +use log::{info, warn}; +use serde::Serialize; +use std::sync::{Arc, RwLock}; +use std::time::Duration; +use wreq::ws::message::Message; +use wreq::ws::WebSocket; + +pub struct GrokWebsocket { + pub(crate) websocket: WebSocket, +} + +impl GrokWebsocket { + pub fn new(websocket: WebSocket) -> Self { + Self { + websocket, + } + } + + pub async fn send(&mut self, message: String) -> Result<(), GrokError> { + let message = Message::text(message); + self.websocket.send(message) + .await + .map_err(|err| { + GrokClientError::WreqClientError(err) + })?; + Ok(()) + } + + pub async fn send_serializable(&mut self, message: T) -> Result<(), GrokError> { + let message_json = serde_json::to_string(&message) + .map_err(|err| { + warn!("Failed to serialize prompt websocket message: {}", err); + GrokClientError::WebsocketRequestSerializationError(err) + })?; + self.send(message_json).await + } + + pub async fn get_response_with_timeout(&mut self, duration: Duration) -> Result, GrokError> { + let result = tokio::time::timeout(duration, self.websocket.try_next()).await; + + match result { + Err(elapsed) => { + info!("Websocket try_next() elapsed without receiving a message: {:?}", elapsed); + Ok(None) // Timeout elapsed. + } + Ok(inner) => { + let maybe_message = inner.map_err(GrokClientError::WebsocketReadError)?; + + match maybe_message { + None => Ok(None), + Some(Message::Text(text)) => { + let maybe_message = WebsocketServerMessage::from_json_str(&text)?; + Ok(Some(maybe_message)) + }, + Some(_) => { + warn!("Received non-text websocket message."); + Ok(None) + }, + } + } + } + } + + pub async fn try_next_timeout(&mut self, duration: Duration) -> Result, GrokClientError> { + let result = tokio::time::timeout(duration, self.websocket.try_next()).await; + + match result { + Err(elapsed) => { + info!("Websocket try_next() elapsed without receiving a message: {:?}", elapsed); + Ok(None) // Timeout elapsed. + } + Ok(inner) => { + let maybe_message = inner.map_err(GrokClientError::WebsocketReadError)?; + Ok(maybe_message) + } + } + } + + // pub async fn is_terminated(&self) -> Result { + // match self.websocket.read() { + // Err(_) => Err(GrokClientError::WebsocketLockError.into()), + // Ok(websocket) => Ok(websocket.is_terminated()), + // } + // } + + // /// NB: This will block forever if there is no item. + // pub async fn try_next(&self) -> Result, GrokClientError> { + // match self.websocket.write() { + // Err(_) => Err(GrokClientError::WebsocketLockError), + // Ok(mut websocket) => websocket + // .try_next() + // .await + // .map_err(|err| GrokClientError::WebsocketReadError(err)), + // } + // } + + // pub fn size_hint(&self) -> Result<(usize, Option), GrokClientError> { + // match self.websocket.read() { + // Err(_) => Err(GrokClientError::WebsocketLockError), + // Ok(websocket) => Ok(websocket.size_hint()) + // } + // } +} diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/clonable_websocket.rs b/crates/api_clients/grok_client/src/requests/image_websocket/grok_wrapped_websocket.rs similarity index 98% rename from crates/api_clients/grok_client/src/requests/image_websocket/clonable_websocket.rs rename to crates/api_clients/grok_client/src/requests/image_websocket/grok_wrapped_websocket.rs index d9e5cf2b11..712d1878f9 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/clonable_websocket.rs +++ b/crates/api_clients/grok_client/src/requests/image_websocket/grok_wrapped_websocket.rs @@ -11,11 +11,11 @@ use wreq::ws::message::Message; use wreq::ws::WebSocket; #[derive(Clone)] -pub struct ClonableWebsocket { +pub struct GrokWrappedWebsocket { pub(crate) websocket: Arc>, } -impl ClonableWebsocket { +impl GrokWrappedWebsocket { pub fn new(websocket: WebSocket) -> Self { Self { websocket: Arc::new(RwLock::new(websocket)), @@ -77,7 +77,6 @@ impl ClonableWebsocket { } } - pub async fn try_next_timeout(&self, duration: Duration) -> Result, GrokClientError> { match self.websocket.write() { Err(_) => Err(GrokClientError::WebsocketLockError), diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_images.rs b/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_images.rs index 5354b322e9..360107db13 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_images.rs +++ b/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_images.rs @@ -1,7 +1,8 @@ use crate::datatypes::api::request_id::RequestId; use crate::error::grok_client_error::GrokClientError; use crate::error::grok_error::GrokError; -use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket; +use crate::requests::image_websocket::grok_websocket::GrokWebsocket; +use crate::requests::image_websocket::grok_wrapped_websocket::GrokWrappedWebsocket; use crate::requests::image_websocket::messages::message_image_data::MessageImageData; use crate::utils::scrub_blobs_for_debug_logging::scrub_blobs_for_debug_logging; use log::info; @@ -12,7 +13,7 @@ use wreq::ws::message::Message; const DEFAULT_IMAGE_COUNT: usize = 6; pub struct ListenForWebsocketImagesArgs<'a> { - pub websocket: &'a ClonableWebsocket, + pub websocket: &'a mut GrokWebsocket, pub timeout: Duration, } @@ -113,8 +114,8 @@ pub async fn listen_for_websocket_images(args: ListenForWebsocketImagesArgs<'_>) #[cfg(test)] mod tests { - use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket; use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; + use crate::requests::image_websocket::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::listen_for_websocket_images::{listen_for_websocket_images, ListenForWebsocketImagesArgs}; use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; use crate::test_utils::get_test_cookies::get_test_cookies; @@ -137,13 +138,14 @@ mod tests { cookies: &cookies, }).await?; - let websocket = ClonableWebsocket::new(websocket); + //let websocket = GrokWrappedWebsocket::new(websocket); + let mut websocket = GrokWebsocket::new(websocket); println!("Sending..."); std::io::stdout().flush()?; let _result = prompt_websocket_image(PromptWebsocketImageArgs { - websocket_wrapped: websocket.clone(), + websocket: &mut websocket, prompt, }).await?; @@ -153,7 +155,7 @@ mod tests { println!("Polling."); let images = listen_for_websocket_images(ListenForWebsocketImagesArgs { - websocket: &websocket, + websocket: &mut websocket, timeout: Duration::from_millis(10_000), }).await?; diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs b/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs index d5c15b083a..a0d0a38999 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs +++ b/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs @@ -1,13 +1,13 @@ use crate::datatypes::api::request_id::RequestId; use crate::error::grok_client_error::GrokClientError; use crate::error::grok_error::GrokError; -use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket; +use crate::requests::image_websocket::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::messages::message_request_id::MessageRequestId; use std::time::{Duration, Instant}; use wreq::ws::message::Message; pub struct ListenForWebsocketRequestIdArgs<'a> { - pub websocket: &'a ClonableWebsocket, + pub websocket: &'a mut GrokWebsocket, pub timeout: Duration, } @@ -70,8 +70,8 @@ pub async fn listen_for_websocket_request_id(args: ListenForWebsocketRequestIdAr #[cfg(test)] mod tests { - use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket; use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; + use crate::requests::image_websocket::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::listen_for_websocket_request_id::{listen_for_websocket_request_id, ListenForWebsocketRequestIdArgs}; use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; use crate::test_utils::get_test_cookies::get_test_cookies; @@ -90,17 +90,18 @@ mod tests { let cookies = get_test_cookies()?; - let websocket = create_listen_websocket(CreateListenWebsocketArgs { + let mut websocket = create_listen_websocket(CreateListenWebsocketArgs { cookies: &cookies, }).await?; - let websocket = ClonableWebsocket::new(websocket); + //let websocket = GrokWrappedWebsocket::new(websocket); + let mut websocket = GrokWebsocket::new(websocket); println!("Sending..."); std::io::stdout().flush()?; let _result = prompt_websocket_image(PromptWebsocketImageArgs { - websocket_wrapped: websocket.clone(), + websocket: &mut websocket, prompt, }).await?; @@ -110,7 +111,7 @@ mod tests { println!("Polling."); let request_data = listen_for_websocket_request_id(ListenForWebsocketRequestIdArgs { - websocket: &websocket, + websocket: &mut websocket, timeout: Duration::from_millis(10_000), }).await?; diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/mod.rs b/crates/api_clients/grok_client/src/requests/image_websocket/mod.rs index 23b5af662c..f6c3676adb 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/mod.rs +++ b/crates/api_clients/grok_client/src/requests/image_websocket/mod.rs @@ -1,6 +1,7 @@ -pub mod clonable_websocket; +pub mod grok_wrapped_websocket; pub mod create_listen_websocket; pub mod listen_for_websocket_images; pub mod listen_for_websocket_request_id; pub mod messages; -pub mod prompt_websocket_image; \ No newline at end of file +pub mod prompt_websocket_image; +pub mod grok_websocket; \ No newline at end of file diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/prompt_websocket_image.rs b/crates/api_clients/grok_client/src/requests/image_websocket/prompt_websocket_image.rs index 113a061225..7b223f5e5e 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/prompt_websocket_image.rs +++ b/crates/api_clients/grok_client/src/requests/image_websocket/prompt_websocket_image.rs @@ -1,9 +1,9 @@ use crate::error::grok_error::GrokError; -use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket; +use crate::requests::image_websocket::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::messages::websocket_client_message::WebsocketClientMessage; pub struct PromptWebsocketImageArgs<'a> { - pub websocket_wrapped: ClonableWebsocket, + pub websocket: &'a mut GrokWebsocket, pub prompt: &'a str, } @@ -11,15 +11,16 @@ pub struct PromptWebsocketImageArgs<'a> { pub async fn prompt_websocket_image(args: PromptWebsocketImageArgs<'_>) -> Result<(), GrokError> { let message = WebsocketClientMessage::new_image_prompt(args.prompt); - args.websocket_wrapped.send_serializable(&message).await?; + args.websocket.send_serializable(&message).await?; Ok(()) } #[cfg(test)] mod tests { - use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket; use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; + use crate::requests::image_websocket::grok_websocket::GrokWebsocket; + use crate::requests::image_websocket::grok_wrapped_websocket::GrokWrappedWebsocket; use crate::requests::image_websocket::messages::websocket_server_message::WebsocketServerMessage; use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; use crate::test_utils::get_test_cookies::get_test_cookies; @@ -40,13 +41,14 @@ mod tests { cookies: &cookies, }).await?; - let websocket = ClonableWebsocket::new(websocket); + //let websocket = GrokWrappedWebsocket::new(websocket); + let mut websocket = GrokWebsocket::new(websocket); println!("Sending..."); std::io::stdout().flush()?; let result = prompt_websocket_image(PromptWebsocketImageArgs { - websocket_wrapped: websocket.clone(), + websocket: &mut websocket, prompt: "a dog riding a motorcycle", }).await?; diff --git a/crates/desktop/artcraft/src/core/artcraft_error.rs b/crates/desktop/artcraft/src/core/artcraft_error.rs index 8830a77973..72ef8847b2 100644 --- a/crates/desktop/artcraft/src/core/artcraft_error.rs +++ b/crates/desktop/artcraft/src/core/artcraft_error.rs @@ -11,6 +11,10 @@ pub enum ArtcraftError { // Service errors GrokError(GrokError), StorytellerError(StorytellerError), + // Lock errors + RwLockReadError, + RwLockWriteError, + MutexLockError, } impl Error for ArtcraftError {} @@ -23,6 +27,9 @@ impl Display for ArtcraftError { Self::IoError(e) => write!(f, "IoError: {:?}", e), Self::GrokError(e) => write!(f, "GrokError: {:?}", e), Self::StorytellerError(e) => write!(f, "StorytellerError: {:?}", e), + Self::RwLockReadError => write!(f, "RwLockReadError"), + Self::RwLockWriteError => write!(f, "RwLockWriteError"), + Self::MutexLockError => write!(f, "MutexLockError"), } } } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs b/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs index 5f16dd9b5e..28b2e501e5 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs @@ -1,4 +1,5 @@ use crate::core::artcraft_error::ArtcraftError; +use anyhow::anyhow; use base64::DecodeError; use errors::AnyhowError; use grok_client::error::grok_error::GrokError; @@ -79,6 +80,8 @@ pub enum BillingProvider { #[derive(Debug)] pub enum ProviderFailureReason { GrokError(GrokError), + /// NB: The Grok client doesn't say why certain errors (eg. missing fields) happen, so we synthesize this. + GrokJobEnqueueFailed, //Fal(FalErrorPlus), MidjourneyError(MidjourneyError), /// NB: The midjourney client doesn't categorize all errors, so we have to do so on our end. @@ -148,6 +151,9 @@ impl From for GenerateError { ArtcraftError::IoError(e) => Self::IoError(e), ArtcraftError::GrokError(e) => Self::ProviderFailure(ProviderFailureReason::GrokError(e)), ArtcraftError::StorytellerError(e) => Self::ProviderFailure(ProviderFailureReason::StorytellerError(e)), + ArtcraftError::RwLockReadError => Self::AnyhowError(anyhow!("Lock read error")), + ArtcraftError::RwLockWriteError => Self::AnyhowError(anyhow!("Lock write error")), + ArtcraftError::MutexLockError => Self::AnyhowError(anyhow!("Mutex lock error")), } } } 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 2d3297f03e..848dc484e6 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 @@ -67,6 +67,7 @@ impl TaskEnqueueSuccess{ Some(GenerationModel::FluxProKontextMax) => Some(TaskModelType::FluxProKontextMax), Some(GenerationModel::Gemini25Flash) => Some(TaskModelType::Gemini25Flash), Some(GenerationModel::GptImage1) => Some(TaskModelType::GptImage1), + Some(GenerationModel::GrokImage) => Some(TaskModelType::GrokImage), Some(GenerationModel::Recraft3) => Some(TaskModelType::Recraft3), Some(GenerationModel::GrokVideo) => Some(TaskModelType::GrokVideo), Some(GenerationModel::Kling21Pro) => Some(TaskModelType::Kling21Pro), diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs index 3eb152ea63..87a3915294 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs @@ -6,6 +6,7 @@ use crate::core::commands::enqueue::text_to_image::generic::handle_image_artcraf use crate::core::commands::enqueue::text_to_image::generic::handle_image_fal::handle_image_fal; use crate::core::commands::enqueue::text_to_image::gpt_image_1::handle_gpt_image_1::handle_gpt_image_1; use crate::core::commands::enqueue::text_to_image::gpt_image_1::handle_gpt_image_1_sora::handle_gpt_image_1_sora; +use crate::core::commands::enqueue::text_to_image::grok::handle_grok::handle_grok; use crate::core::commands::enqueue::text_to_image::midjourney::handle_midjourney::handle_midjourney; use crate::core::commands::response::failure_response_wrapper::{CommandErrorResponseWrapper, CommandErrorStatus}; use crate::core::commands::response::shorthand::Response; @@ -19,6 +20,8 @@ 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::provider_priority::{Provider, ProviderPriorityStore}; use crate::core::state::task_database::TaskDatabase; +use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::state::grok_image_prompt_queue::GrokImagePromptQueue; use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; use crate::services::sora::state::sora_task_queue::SoraTaskQueue; @@ -48,6 +51,8 @@ pub enum ImageModel { FluxPro11Ultra, #[serde(rename = "gpt_image_1")] GptImage1, + #[serde(rename = "grok_image")] + GrokImage, #[serde(rename = "recraft_3")] Recraft3, @@ -144,6 +149,8 @@ pub async fn enqueue_text_to_image_command( provider_priority_store: State<'_, ProviderPriorityStore>, task_database: State<'_, TaskDatabase>, mj_creds_manager: State<'_, MidjourneyCredentialManager>, + grok_creds_manager: State<'_, GrokCredentialManager>, + grok_image_prompt_queue: State<'_, GrokImagePromptQueue>, storyteller_creds_manager: State<'_, StorytellerCredentialManager>, sora_creds_manager: State<'_, SoraCredentialManager>, sora_task_queue: State<'_, SoraTaskQueue>, @@ -158,6 +165,8 @@ pub async fn enqueue_text_to_image_command( &provider_priority_store, &task_database, &mj_creds_manager, + &grok_creds_manager, + &grok_image_prompt_queue, &storyteller_creds_manager, &app_env_configs, &sora_creds_manager, @@ -226,6 +235,8 @@ pub async fn handle_request( provider_priority_store: &ProviderPriorityStore, task_database: &TaskDatabase, mj_creds_manager: &MidjourneyCredentialManager, + grok_creds_manager: &GrokCredentialManager, + grok_image_prompt_queue: &GrokImagePromptQueue, storyteller_creds_manager: &StorytellerCredentialManager, app_env_configs: &AppEnvConfigs, sora_creds_manager: &SoraCredentialManager, @@ -240,6 +251,8 @@ pub async fn handle_request( &storyteller_creds_manager, &app_env_configs, &mj_creds_manager, + &grok_creds_manager, + &grok_image_prompt_queue, &sora_creds_manager, &sora_task_queue, ).await; @@ -274,6 +287,8 @@ pub async fn dispatch_request( storyteller_creds_manager: &StorytellerCredentialManager, app_env_configs: &AppEnvConfigs, mj_creds_manager: &MidjourneyCredentialManager, + grok_creds_manager: &GrokCredentialManager, + grok_image_prompt_queue: &GrokImagePromptQueue, sora_creds_manager: &SoraCredentialManager, sora_task_queue: &SoraTaskQueue, ) -> Result { @@ -306,6 +321,15 @@ pub async fn dispatch_request( sora_task_queue, ).await; } + Some(ImageModel::GrokImage) => { + return handle_grok( + app, + request, + app_env_configs, + grok_creds_manager, + grok_image_prompt_queue, + ).await; + } Some(ImageModel::Midjourney) => { return handle_midjourney( app, diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/generic/handle_image_artcraft.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/generic/handle_image_artcraft.rs index a9225d2fb4..aa175a3e4a 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/generic/handle_image_artcraft.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/generic/handle_image_artcraft.rs @@ -58,6 +58,7 @@ pub async fn handle_image_artcraft( Some( ImageModel::Gemini25Flash | ImageModel::GptImage1 | + ImageModel::GrokImage | ImageModel::Midjourney ) => { return Err(GenerateError::AnyhowError(anyhow!("wrong logic: another branch should handle this: {:?}", request.model))); diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs new file mode 100644 index 0000000000..c00c03eac6 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs @@ -0,0 +1,138 @@ +use crate::core::commands::enqueue::generate_error::GenerateError; +use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; +use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::EnqueueTextToImageRequest; +use crate::core::events::generation_events::common::GenerationModel; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::state::grok_image_prompt_queue::{GrokImagePromptQueue, PromptItem}; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_type::TaskType; +use idempotency::uuid::generate_random_uuid; +use tauri::AppHandle; + +pub async fn handle_grok( + app: &AppHandle, + request: &EnqueueTextToImageRequest, + app_env_configs: &AppEnvConfigs, + creds_manager: &GrokCredentialManager, + grok_image_prompt_queue: &GrokImagePromptQueue, +) -> Result { + + + //// TODO: We can request population of the user info if absent or expired. + //let websocket = get_websocket( + // app, + // creds_manager, + // websocket_manager, + //).await?; + // + //let _result = prompt_websocket_image(PromptWebsocketImageArgs { + // websocket_wrapped: &websocket, + // prompt: request.prompt.as_deref().unwrap_or_else(||""), + //}).await?; + // + //let request_id = listen_for_websocket_request_id(ListenForWebsocketRequestIdArgs { + // websocket: &websocket, + // timeout: Duration::from_millis(10_000), + //}).await?; + // + //let request_id = match request_id.request_id { + // Some(request_id) => request_id.to_string(), + // None => { + // warn!("No request id from Grok image generation (websocket)"); + // return Err(GenerateError::ProviderFailure(ProviderFailureReason::GrokJobEnqueueFailed)); + // } + //}; + // + //info!("Successfully enqueued MidJourney. Job token: {}", request_id); + + let job_id = generate_random_uuid(); + + let prompt = request.prompt + .as_deref() + .map(|prompt| prompt.trim().to_string()) + .unwrap_or_else(|| "".to_string()); + + grok_image_prompt_queue.enqueue(PromptItem { + task_id: job_id.clone(), + prompt, + })?; + + Ok(TaskEnqueueSuccess { + provider: GenerationProvider::Grok, + model: Some(GenerationModel::GrokImage), + provider_job_id: Some(job_id), + task_type: TaskType::ImageGeneration, + }) +} + +// TODO: +// fn handle_midjourney_errors( +// app: &AppHandle, +// maybe_errors: Option> +// ) -> Result { +// if let Some(errors) = maybe_errors { +// if !errors.is_empty() { +// let messages: Vec = errors.iter() +// .map(|e| format!("{:?}", e)) +// .collect(); +// +// let combined_message = messages.join("; "); +// +// let event = FlashUserInputErrorEvent { +// message: format!("Midjourney Error: {}", combined_message), +// }; +// +// if let Err(err) = event.send(&app) { +// error!("Failed to send FlashUserInputErrorEvent: {:?}", err); // Fail open +// } +// } +// } +// +// Err(GenerateError::ProviderFailure(ProviderFailureReason::MidjourneyJobEnqueueFailed)) +// } + +//async fn get_websocket(app: &AppHandle, creds_manager: &GrokCredentialManager, websocket_manager: &GrokWebsocketManager) -> Result { +// let maybe_websocket = websocket_manager.grab_websocket()?; +// +// let websocket = match maybe_websocket { +// Some(websocket) => websocket, +// None => { +// let cookies = match creds_manager.maybe_copy_cookie_store() { +// Ok(Some(cookies)) => cookies.to_cookie_string(), +// Ok(None) => { +// error!("Grok credentials not found."); +// ShowProviderLoginModalEvent::send_for_provider(GenerationProvider::Grok, &app); +// return Err(GenerateError::needs_grok_credentials()); +// } +// Err(err) => { +// error!("Error reading Midjourney credentials: {:?}", err); +// ShowProviderLoginModalEvent::send_for_provider(GenerationProvider::Grok, &app); +// return Err(GenerateError::needs_grok_credentials()); +// }, +// }; +// +// info!("Creating Grok websocket..."); +// +// let create_websocket_result = create_listen_websocket(CreateListenWebsocketArgs { +// cookies: &cookies, +// }).await; +// +// let websocket = match create_websocket_result { +// Ok(websocket) => GrokWebsocket::new(websocket), +// Err(err) => { +// error!("Error creating websocket: {:?}", err); +// ShowProviderLoginModalEvent::send_for_provider(GenerationProvider::Grok, &app); +// return Err(GenerateError::needs_grok_credentials()); +// } +// }; +// +// info!("Setting managed Grok websocket..."); +// websocket_manager.set_websocket(websocket.clone())?; +// +// websocket +// } +// }; +// +// Ok(websocket) +//} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/mod.rs new file mode 100644 index 0000000000..181e756f1a --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/mod.rs @@ -0,0 +1 @@ +pub (super) mod handle_grok; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/mod.rs index 073f828a2a..7f981744f8 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/mod.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/mod.rs @@ -2,4 +2,5 @@ pub mod enqueue_text_to_image_command; pub mod generic; mod gemini_25_flash; mod gpt_image_1; -mod midjourney; \ No newline at end of file +mod midjourney; +mod grok; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/events/generation_events/common.rs b/crates/desktop/artcraft/src/core/events/generation_events/common.rs index 688164d404..1ee08c3772 100644 --- a/crates/desktop/artcraft/src/core/events/generation_events/common.rs +++ b/crates/desktop/artcraft/src/core/events/generation_events/common.rs @@ -30,6 +30,8 @@ pub enum GenerationModel { FluxProKontextMax, #[serde(rename = "gpt_image_1")] GptImage1, + #[serde(rename = "grok_image")] + GrokImage, #[serde(rename = "gemini_25_flash")] Gemini25Flash, 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 967c3c0b60..3c33ef011b 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,8 @@ 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::state::grok_image_prompt_queue::GrokImagePromptQueue; +use crate::services::grok::threads::grok_image_websocket_thread::grok_image_websocket_thread::grok_image_websocket_thread; 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; @@ -29,6 +31,7 @@ pub async fn handle_tauri_startup( sora_task_queue: SoraTaskQueue, mj_creds_manager: MidjourneyCredentialManager, grok_creds_manager: GrokCredentialManager, + grok_image_prompt_queue: GrokImagePromptQueue, ) -> AnyhowResult<()> { set_app_log_level( @@ -77,6 +80,16 @@ pub async fn handle_tauri_startup( storyteller_creds_manager.clone(), )); + tauri::async_runtime::spawn(grok_image_websocket_thread( + app.clone(), + app_env_configs.clone(), + root.clone(), + task_database.clone(), + grok_creds_manager.clone(), + grok_image_prompt_queue.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/lib.rs b/crates/desktop/artcraft/src/lib.rs index 7635aac599..511c694205 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -32,6 +32,7 @@ use crate::core::threads::main_window_thread::main_window_thread::main_window_th use crate::services::grok::commands::grok_clear_credentials_command::grok_clear_credentials_command; use crate::services::grok::commands::grok_get_credential_info_command::grok_get_credential_info_command; use crate::services::grok::commands::grok_open_login_command::grok_open_login_command; +use crate::services::grok::state::grok_image_prompt_queue::GrokImagePromptQueue; use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; use crate::services::midjourney::commands::midjourney_clear_credentials_command::midjourney_clear_credentials_command; use crate::services::midjourney::commands::midjourney_get_credential_info_command::midjourney_get_credential_info_command; @@ -104,6 +105,9 @@ pub fn run() { let grok_creds_manager = GrokCredentialManager::initialize_from_disk_infallible(&app_data_root); let grok_creds_manager_2 = grok_creds_manager.clone(); + let grok_prompt_queue = GrokImagePromptQueue::new(); + let grok_prompt_queue_2 = grok_prompt_queue.clone(); + println!("Initializing backend runtime..."); let builder = tauri::Builder::default() @@ -142,6 +146,7 @@ pub fn run() { sora_tasks, midjourney_creds_manager_2, grok_creds_manager_2, + grok_prompt_queue_2, ).await; if let Err(err) = result { @@ -157,6 +162,7 @@ pub fn run() { .manage(app_preferences) .manage(artcraft_platform_info) .manage(grok_creds_manager) + .manage(grok_prompt_queue) .manage(midjourney_creds_manager) .manage(sora_creds_manager) .manage(sora_task_queue) diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs b/crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs new file mode 100644 index 0000000000..6bf31824f4 --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs @@ -0,0 +1,76 @@ +use crate::core::artcraft_error::ArtcraftError; +use log::error; +use std::collections::VecDeque; +use std::sync::{Arc, Mutex}; + +#[derive(Clone)] +pub struct GrokImagePromptQueue { + pub prompt_queue: Arc>>, +} + + +#[derive(Clone, Debug)] +pub struct PromptItem { + /// Local database task ID + pub task_id: String, + + /// Text prompt + pub prompt: String, +} + +impl GrokImagePromptQueue { + pub fn new() -> Self { + Self { + prompt_queue: Arc::new(Mutex::new(VecDeque::new())), + } + } + + pub fn enqueue(&self, prompt_item: PromptItem) -> Result<(), ArtcraftError> { + match self.prompt_queue.lock() { + Ok(mut queue) => { + queue.push_back(prompt_item); + Ok(()) + }, + Err(err) => { + error!("Error locking prompt queue: {:?}", err); + Err(ArtcraftError::MutexLockError) + }, + } + } + + pub fn dequeue(&self) -> Result, ArtcraftError> { + match self.prompt_queue.lock() { + Ok(mut queue) => { + Ok(queue.pop_front()) + }, + Err(err) => { + error!("Error locking prompt queue: {:?}", err); + Err(ArtcraftError::MutexLockError) + }, + } + } + + pub fn is_empty(&self) -> Result { + match self.prompt_queue.lock() { + Ok(mut queue) => { + Ok(queue.is_empty()) + }, + Err(err) => { + error!("Error locking prompt queue: {:?}", err); + Err(ArtcraftError::MutexLockError) + }, + } + } + + pub fn len(&self) -> Result { + match self.prompt_queue.lock() { + Ok(mut queue) => { + Ok(queue.len()) + }, + Err(err) => { + error!("Error locking prompt queue: {:?}", err); + Err(ArtcraftError::MutexLockError) + }, + } + } +} diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_websocket_manager.rs b/crates/desktop/artcraft/src/services/grok/state/grok_websocket_manager.rs new file mode 100644 index 0000000000..9e3f5c8e97 --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/state/grok_websocket_manager.rs @@ -0,0 +1,62 @@ +use crate::core::artcraft_error::ArtcraftError; +use errors::AnyhowResult; +use grok_client::requests::image_websocket::grok_wrapped_websocket::GrokWrappedWebsocket; +use log::error; +use std::sync::{Arc, RwLock}; + + +/// NB: This is inefficient because the websockets are locked at two layers. +/// Should be fine for our performance needs, though. +#[derive(Clone)] +pub struct GrokWebsocketManager { + //websocket: Arc>>, + websocket: Arc>>, +} + +impl GrokWebsocketManager { + pub fn new() -> Self { + Self { + websocket: Arc::new(RwLock::new(None)), + } + } + + pub fn set_websocket(&self, websocket: GrokWrappedWebsocket) -> Result<(), ArtcraftError> { + match self.websocket.write() { + Ok(mut guard) => { + //*guard = Some(websocket); + *guard = Some(()); + Ok(()) + } + Err(err) => { + error!("Error writing locked websocket: {}", err); + Err(ArtcraftError::RwLockWriteError) + } + } + } + + pub fn clear_websocket(&self) -> Result<(), ArtcraftError> { + match self.websocket.write() { + Ok(mut guard) => { + *guard = None; + Ok(()) + } + Err(err) => { + error!("Error writing locked websocket: {}", err); + Err(ArtcraftError::RwLockWriteError) + } + } + } + + pub fn grab_websocket(&self) -> Result, ArtcraftError> { + match self.websocket.read() { + Ok(guard) => { + //Ok(guard.clone()) + Ok(None) + } + Err(err) => { + error!("Error reading locked websocket: {}", err); + Err(ArtcraftError::RwLockReadError) + } + } + } +} diff --git a/crates/desktop/artcraft/src/services/grok/state/mod.rs b/crates/desktop/artcraft/src/services/grok/state/mod.rs index 9ea2b969bd..5b557e4567 100644 --- a/crates/desktop/artcraft/src/services/grok/state/mod.rs +++ b/crates/desktop/artcraft/src/services/grok/state/mod.rs @@ -1,3 +1,7 @@ pub mod grok_credential_manager; +pub mod grok_image_prompt_queue; pub mod grok_serializable_state; -mod grok_credential_holder; \ No newline at end of file +pub mod grok_websocket_manager; + +// Private +mod grok_credential_holder; diff --git a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs new file mode 100644 index 0000000000..fdf88e1f1e --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs @@ -0,0 +1,364 @@ +use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::functional_events::text_to_image_generation_complete_event::{GeneratedImage, TextToImageGenerationCompleteEvent}; +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::task_database::TaskDatabase; +use crate::core::utils::download_url_to_temp_dir::download_url_to_temp_dir; +use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES; +use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::state::grok_image_prompt_queue::{GrokImagePromptQueue, PromptItem}; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use anyhow::anyhow; +use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest; +use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template; +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_model_type::TaskModelType; +use errors::AnyhowResult; +use grok_client::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; +use grok_client::requests::image_websocket::grok_websocket::GrokWebsocket; +use grok_client::requests::image_websocket::grok_wrapped_websocket::GrokWrappedWebsocket; +use grok_client::requests::image_websocket::listen_for_websocket_images::{listen_for_websocket_images, ImageResults, ListenForWebsocketImagesArgs}; +use grok_client::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; +use idempotency::uuid::generate_random_uuid; +use log::{error, info, warn}; +use sqlite_tasks::queries::get_task_by_provider_and_provider_job_id::{get_task_by_provider_and_provider_job_id, GetTaskByProviderAndProviderJobIdArgs}; +use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs}; +use sqlite_tasks::queries::task::Task; +use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs}; +use sqlite_tasks::queries::update_successful_task_status_with_metadata_by_provider::{update_successful_task_status_with_metadata_by_provider, UpdateSuccessfulTaskByProviderArgs}; +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::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files; +use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs, UploadImageMediaFileSuccessResponse}; +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 tokens::tokens::sqlite::tasks::TaskId; + +pub async fn grok_image_websocket_thread( + app_handle: AppHandle, + app_env_configs: AppEnvConfigs, + app_data_root: AppDataRoot, + task_database: TaskDatabase, + creds: GrokCredentialManager, + prompt_queue: GrokImagePromptQueue, + storyteller_creds_manager: StorytellerCredentialManager, +) -> ! { + loop { + let res = inner_loop( + &app_handle, + &app_env_configs, + &app_data_root, + &task_database, + &creds, + &prompt_queue, + &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 inner_loop( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + grok_creds: &GrokCredentialManager, + prompt_queue: &GrokImagePromptQueue, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> AnyhowResult<()> { + + loop { + if !grok_creds.do_task_polling()? { + tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; + continue; + } + + let maybe_cookie_header = grok_creds.maybe_copy_cookie_header_string()?; + + let cookie_header = match maybe_cookie_header { + Some(cookie_header) => cookie_header, + None => { + tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; + continue; + } + }; + + let websocket = create_listen_websocket(CreateListenWebsocketArgs { + cookies: &cookie_header, + }).await?; + + let mut websocket = GrokWebsocket::new(websocket); + + //poll_task_loop( + // &app_handle, + // &app_env_configs, + // &app_data_root, + // &task_database, + // &grok_creds, + // websocket, + // prompt_queue, + // &storyteller_creds_manager, + //).await?; + + loop { + // 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?; + // // Only Grok images + // let local_tasks = local_tasks.tasks.into_iter() + // .filter(|task| match task.model_type { + // Some(TaskModelType::GrokImage) => true, + // _ => false, + // }) + // .collect::>(); + + let maybe_prompt = prompt_queue.dequeue()?; + + let prompt_item = match maybe_prompt { + Some(prompt_item) => prompt_item, + None => { + tokio::time::sleep(std::time::Duration::from_millis(1_000)).await; + continue; + } + }; + + info!("Prompting Grok websocket: {}", prompt_item.prompt); + + let _result = prompt_websocket_image(PromptWebsocketImageArgs { + websocket: &mut websocket, + prompt: &prompt_item.prompt, + }).await?; + + let images = listen_for_websocket_images(ListenForWebsocketImagesArgs { + websocket: &mut websocket, + timeout: Duration::from_millis(30_000), + }).await?; + + upload_images_to_storyteller( + &app_handle, + &app_env_configs, + &app_data_root, + &task_database, + &storyteller_creds_manager, + &prompt_item, + &images, + ).await? + } + + } +} + +//async fn poll_task_loop( +// app_handle: &AppHandle, +// app_env_configs: &AppEnvConfigs, +// app_data_root: &AppDataRoot, +// task_database: &TaskDatabase, +// grok_creds: &GrokCredentialManager, +// grok_websocket: GrokWebsocket, +// prompt_queue: &GrokImagePromptQueue, +// storyteller_creds_manager: &StorytellerCredentialManager, +//) -> AnyhowResult<()> { +// +//} + +async fn upload_images_to_storyteller( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds_manager: &StorytellerCredentialManager, + prompt_item: &PromptItem, + images: &ImageResults, +) -> AnyhowResult<()> { + + loop { + 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 prompt = prompt_item.prompt.trim().to_string(); + + let request = CreatePromptRequest { + uuid_idempotency_token: generate_random_uuid(), + positive_prompt: Some(prompt), + negative_prompt: None, + model_type: Some(ModelType::GrokImage), + generation_provider: Some(GenerationProvider::Grok), + }; + + let prompt_response = create_prompt( + &app_env_configs.storyteller_host, + Some(&storyteller_creds), + request + ).await?; + + info!("Created prompt: {:?}", &prompt_response.prompt_token); + + // TODO: Move this from clientside to the backend. + // The first upload should produce a batch token that we can reuse. + let batch_token = BatchGenerationToken::generate(); + + let mut maybe_primary_media_file_token = None; + + for (i, image) in images.images.iter().enumerate() { + loop { + let url = image.url.to_string(); + let file = download_url_to_temp_dir(&url, app_data_root).await?; + + info!("Uploading image {} of {} ...", i, images.images.len()); + + let result = upload_image_media_file_from_file(UploadImageFromFileArgs { + api_host: &app_env_configs.storyteller_host, + maybe_creds: Some(&storyteller_creds), + path: file.path(), + is_intermediate_system_file: false, + maybe_prompt_token: Some(&prompt_response.prompt_token), + maybe_batch_token: Some(&batch_token), + }).await; + + match result { + Ok(uploaded) => { + if maybe_primary_media_file_token.is_none() { + maybe_primary_media_file_token = Some(uploaded.media_file_token); + } + break; // Break retry loop. + } + Err(StorytellerError::Api(ApiError::TooManyRequests(_))) => { + tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; + continue; + } + Err(err) => { + error!("Error uploading image: {}", err); + break; + } + } + } // Retry on 429s, etc. + } + + let event = GenerationCompleteEvent { + //media_file_token: result.media_file_token, + action: Some(GenerationAction::GenerateImage), + service: GenerationServiceProvider::Grok, + model: None, + }; + + event.send_infallible(&app_handle); + + let task = get_task_by_provider_and_provider_job_id(GetTaskByProviderAndProviderJobIdArgs { + db: task_database.get_connection(), + provider: GenerationProvider::Grok, + provider_job_id: &prompt_item.task_id, + }).await?; + + let task = match task { + Some(task) => task, + None => { + warn!("Couldn't find local sqlite task by id: {}", &prompt_item.task_id); + return Ok(()); + } + }; + + 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: &task.id, + maybe_batch_token: Some(&batch_token), + maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(), + maybe_primary_media_file_class: Some(TaskMediaFileClass::Image), + 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. + } + + send_frontend_ui_update( + app_handle, + app_env_configs, + Some(&storyteller_creds), + &task, + &batch_token, + ).await?; + + return Ok(()) + } +} + +async fn send_frontend_ui_update( + app: &AppHandle, + app_env_configs: &AppEnvConfigs, + maybe_creds: Option<&StorytellerCredentialSet>, + task: &Task, + batch_token: &BatchGenerationToken, +) -> AnyhowResult<()> { + + let result = list_batch_generated_redux_media_files( + &app_env_configs.storyteller_host, + maybe_creds, + batch_token, + ).await?; + + if result.media_files.is_empty() { + return Err(anyhow!("No media files found for batch token: {}", batch_token)); + } + + let media_files = result.media_files + .into_iter() + .map(|file| GeneratedImage { + media_token: file.token, + cdn_url: file.media_links.cdn_url, + maybe_thumbnail_template: file.media_links.maybe_thumbnail_template, + }) + .collect(); + + let event = TextToImageGenerationCompleteEvent { + generated_images: media_files, + maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(), + maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(), + }; + + event.send_infallible(app); + + Ok(()) +} diff --git a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/mod.rs b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/mod.rs new file mode 100644 index 0000000000..5021bcae7c --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/mod.rs @@ -0,0 +1 @@ +pub mod grok_image_websocket_thread; \ No newline at end of file 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 index 278424f636..10e38f7841 100644 --- 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 @@ -33,7 +33,8 @@ use midjourney_client::endpoints::imagine::{imagine, ImagineItem, ImagineRequest 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::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList}; +use sqlite_tasks::queries::task::Task; 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}; diff --git a/crates/desktop/artcraft/src/services/grok/threads/mod.rs b/crates/desktop/artcraft/src/services/grok/threads/mod.rs index 5239c8592a..07dab82921 100644 --- a/crates/desktop/artcraft/src/services/grok/threads/mod.rs +++ b/crates/desktop/artcraft/src/services/grok/threads/mod.rs @@ -1 +1,2 @@ -pub mod grok_video_task_polling; \ No newline at end of file +pub mod grok_video_task_polling; +pub mod grok_image_websocket_thread; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/midjourney/threads/events/maybe_handle_text_to_image_complete_event.rs b/crates/desktop/artcraft/src/services/midjourney/threads/events/maybe_handle_text_to_image_complete_event.rs index 8b275f8a5e..3d8fcfb6cc 100644 --- a/crates/desktop/artcraft/src/services/midjourney/threads/events/maybe_handle_text_to_image_complete_event.rs +++ b/crates/desktop/artcraft/src/services/midjourney/threads/events/maybe_handle_text_to_image_complete_event.rs @@ -7,7 +7,8 @@ use enums::tauri::tasks::task_type::TaskType; use enums::tauri::ux::tauri_command_caller::TauriCommandCaller; use errors::AnyhowResult; use log::{error, warn}; -use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList}; +use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList}; +use sqlite_tasks::queries::task::Task; use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; use storyteller_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files; use tauri::AppHandle; diff --git a/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs b/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs index 7d04f2272e..bfeaa0a04c 100644 --- a/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs +++ b/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs @@ -27,7 +27,8 @@ use midjourney_client::endpoints::imagine::{imagine, ImagineItem, ImagineRequest 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::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList}; +use sqlite_tasks::queries::task::Task; 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}; @@ -362,7 +363,7 @@ async fn upload_midjourney_batch( task_id: &local_task.id, maybe_batch_token: Some(&batch_token), maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(), - maybe_primary_media_file_class: Some(TaskMediaFileClass::Image), // TODO: Just images for now. + maybe_primary_media_file_class: Some(TaskMediaFileClass::Image), 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?; diff --git a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_failed_generations.rs b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_failed_generations.rs index 63b5cde618..f126be4c23 100644 --- a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_failed_generations.rs +++ b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_failed_generations.rs @@ -11,7 +11,7 @@ use errors::AnyhowResult; use log::info; use openai_sora_client::requests::common::task_id::TaskId; use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::PartialTaskResponse; -use sqlite_tasks::queries::list_tasks_by_provider_and_status::Task; +use sqlite_tasks::queries::task::Task; use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; use std::collections::HashMap; use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; 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 559dcb6152..e538f1a3fc 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 @@ -28,7 +28,8 @@ use openai_sora_client::recipes::list_classic_sora_tasks_with_session_auto_renew use openai_sora_client::requests::common::task_id::TaskId; use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::{PartialGeneration, PartialTaskResponse}; use reqwest::Url; -use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList}; +use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList}; +use sqlite_tasks::queries::task::Task; 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}; 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 94e846218a..3a4169a6b1 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 @@ -10,6 +10,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::sora::threads::sora_task_polling::helpers::handle_failed_generations::{handle_classic_failed_generations, FailedGeneration}; use crate::services::sora::threads::sora_task_polling::helpers::handle_successful_generations::{handle_classic_successful_generations, GenerationItem, SuccessfulGeneration}; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; @@ -26,7 +27,8 @@ use openai_sora_client::recipes::list_classic_sora_tasks_with_session_auto_renew use openai_sora_client::requests::common::task_id::TaskId; use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::TaskStatus; use reqwest::Url; -use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList}; +use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList}; +use sqlite_tasks::queries::task::Task; use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; use std::collections::{HashMap, HashSet}; use std::fs::File; @@ -36,7 +38,6 @@ 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, 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 79d7c29fd3..f0b19c5e69 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 @@ -10,6 +10,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::sora::threads::sora_task_polling::helpers::handle_failed_generations::{handle_classic_failed_generations, FailedGeneration}; use crate::services::sora::threads::sora_task_polling::helpers::handle_successful_generations::{handle_classic_successful_generations, GenerationItem, SuccessfulGeneration}; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; @@ -28,7 +29,8 @@ use openai_sora_client::requests::common::task_id::TaskId; use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::TaskStatus; use openai_sora_client::requests::list_sora2_drafts::list_sora2_drafts::Draft; use reqwest::Url; -use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList}; +use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList}; +use sqlite_tasks::queries::task::Task; use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; use std::collections::{HashMap, HashSet}; use std::fs::File; @@ -38,7 +40,6 @@ 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, diff --git a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/sora_task_polling_thread.rs b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/sora_task_polling_thread.rs index 16ea483cee..4c5279f415 100644 --- a/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/sora_task_polling_thread.rs +++ b/crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/sora_task_polling_thread.rs @@ -12,6 +12,8 @@ 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::handle_failed_generations::handle_classic_failed_generations; use crate::services::sora::threads::sora_task_polling::helpers::handle_successful_generations::handle_classic_successful_generations; +use crate::services::sora::threads::sora_task_polling::helpers::poll_classic_sora_tasks::poll_classic_sora_tasks; +use crate::services::sora::threads::sora_task_polling::helpers::poll_sora_2_tasks::poll_sora_2_tasks; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest; use enums::common::generation_provider::GenerationProvider; @@ -26,7 +28,8 @@ use openai_sora_client::recipes::list_classic_sora_tasks_with_session_auto_renew use openai_sora_client::requests::common::task_id::TaskId; use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::TaskStatus; use reqwest::Url; -use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList}; +use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList}; +use sqlite_tasks::queries::task::Task; use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; use std::collections::{HashMap, HashSet}; use std::fs::File; @@ -36,8 +39,6 @@ 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::poll_classic_sora_tasks::poll_classic_sora_tasks; -use crate::services::sora::threads::sora_task_polling::helpers::poll_sora_2_tasks::poll_sora_2_tasks; pub async fn sora_task_polling_thread( app_handle: AppHandle, diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_inpainting_complete_event.rs b/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_inpainting_complete_event.rs index 4b881c317c..122742ac29 100644 --- a/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_inpainting_complete_event.rs +++ b/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_inpainting_complete_event.rs @@ -9,7 +9,8 @@ use enums::tauri::tasks::task_type::TaskType; use enums::tauri::ux::tauri_command_caller::TauriCommandCaller; use errors::AnyhowResult; use log::{error, warn}; -use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task}; +use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs}; +use sqlite_tasks::queries::task::Task; use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; use storyteller_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files; use tauri::AppHandle; diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_text_to_image_complete_event.rs b/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_text_to_image_complete_event.rs index 04722ae90b..f386f47e45 100644 --- a/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_text_to_image_complete_event.rs +++ b/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_text_to_image_complete_event.rs @@ -7,7 +7,8 @@ use enums::tauri::tasks::task_type::TaskType; use enums::tauri::ux::tauri_command_caller::TauriCommandCaller; use errors::AnyhowResult; use log::{error, warn}; -use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task}; +use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs}; +use sqlite_tasks::queries::task::Task; use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; use storyteller_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files; use tauri::AppHandle; diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_send_background_removal_complete_event.rs b/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_send_background_removal_complete_event.rs index 130f19c498..f5fa250e7d 100644 --- a/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_send_background_removal_complete_event.rs +++ b/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_send_background_removal_complete_event.rs @@ -6,7 +6,8 @@ use enums::tauri::tasks::task_type::TaskType; use enums::tauri::ux::tauri_command_caller::TauriCommandCaller; use errors::AnyhowResult; use log::error; -use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task}; +use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs}; +use sqlite_tasks::queries::task::Task; use tokens::tokens::media_files::MediaFileToken; pub async fn maybe_send_background_removal_complete_event( diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_activity_thread.rs b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_activity_thread.rs index be57c5e5a1..2ee52c44c3 100644 --- a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_activity_thread.rs +++ b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_activity_thread.rs @@ -17,15 +17,15 @@ use errors::AnyhowResult; use log::{debug, error, info}; use os_info::Info; use reqwest::Url; -use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task}; +use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs}; use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; use std::time::Instant; -use storyteller_client::endpoints::analytics::log_active_user::log_active_user; use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; -use storyteller_client::error::api_error::ApiError; -use storyteller_client::error::storyteller_error::StorytellerError; +use storyteller_client::endpoints::analytics::log_active_user::log_active_user; use storyteller_client::endpoints::jobs::list_session_jobs::{list_session_jobs, States}; use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file; +use storyteller_client::error::api_error::ApiError; +use storyteller_client::error::storyteller_error::StorytellerError; use tauri::AppHandle; use tokens::tokens::app_session::AppSessionToken; diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs index 02a6196222..3f9c8f7f63 100644 --- a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs +++ b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs @@ -27,7 +27,8 @@ use enums::tauri::tasks::task_type::TaskType; use errors::AnyhowResult; use log::{error, info}; use reqwest::Url; -use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task}; +use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs}; +use sqlite_tasks::queries::task::Task; 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; diff --git a/crates/schema/database/sqlite_tasks/src/queries/get_task_by_provider_and_provider_job_id.rs b/crates/schema/database/sqlite_tasks/src/queries/get_task_by_provider_and_provider_job_id.rs new file mode 100644 index 0000000000..3d1e3aaf80 --- /dev/null +++ b/crates/schema/database/sqlite_tasks/src/queries/get_task_by_provider_and_provider_job_id.rs @@ -0,0 +1,78 @@ +use crate::connection::TaskDbConnection; +use crate::error::SqliteTasksError; +use crate::queries::task::Task; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_model_type::TaskModelType; +use enums::tauri::tasks::task_status::TaskStatus; +use enums::tauri::tasks::task_type::TaskType; +use enums::tauri::ux::tauri_command_caller::TauriCommandCaller; +use sqlx::Error; +use tokens::tokens::sqlite::tasks::TaskId; + +pub struct GetTaskByProviderAndProviderJobIdArgs<'a> { + pub db: &'a TaskDbConnection, + pub provider: GenerationProvider, + pub provider_job_id: &'a str, +} + + +pub async fn get_task_by_provider_and_provider_job_id( + args: GetTaskByProviderAndProviderJobIdArgs<'_>, +) -> Result, SqliteTasksError> { + + // TODO(bt,2025-07-12): Fix this. The sqlx mysql queries never required temporaries + let temp_provider = args.provider.to_str(); + let temp_provider_job_id = args.provider_job_id; + + let query = sqlx::query!(r#" + SELECT + id, + task_status, + task_type, + model_type, + provider, + provider_job_id, + frontend_caller, + frontend_subscriber_id, + frontend_subscriber_payload + FROM tasks + WHERE + provider = ? + AND provider_job_id = ? + "#, + temp_provider, + temp_provider_job_id, + ); + + // info!("Query: {:?}", query.sql()); + + let result = query.fetch_one(args.db.get_pool()).await; + + let record = match result { + Ok(record) => record, + Err(err) => return match err { + Error::RowNotFound => Ok(None), + _ => Err(err.into()), + }, + }; + + // NB: Not sure why query can't figure out this isn't nullable + let provider = record.provider.as_deref() + .unwrap_or_else(|| args.provider.to_str()); + + Ok(Some(Task { + id: TaskId::new_from_str(&record.id), + status: TaskStatus::from_str(&record.task_status)?, + task_type: TaskType::from_str(&record.task_type)?, + model_type: record.model_type + .map(|model| TaskModelType::from_str(&model)) + .transpose()?, + provider: GenerationProvider::from_str(provider)?, + provider_job_id: record.provider_job_id, + frontend_caller: record.frontend_caller + .map(|caller| TauriCommandCaller::from_str(&caller)) + .transpose()?, + frontend_subscriber_id: record.frontend_subscriber_id, + frontend_subscriber_payload: record.frontend_subscriber_payload, + })) +} diff --git a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_status.rs b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_status.rs index 0d16ae4629..20e5f2003b 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_status.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_status.rs @@ -1,5 +1,6 @@ use crate::connection::TaskDbConnection; use crate::error::SqliteTasksError; +use crate::queries::task::{RawTask, Task}; use enums::common::generation_provider::GenerationProvider; use enums::tauri::tasks::task_model_type::TaskModelType; use enums::tauri::tasks::task_status::TaskStatus; @@ -19,33 +20,6 @@ pub struct TaskList { pub tasks: Vec, } -#[derive(Debug, Clone)] -pub struct Task { - pub id: TaskId, - pub status: TaskStatus, - pub task_type: TaskType, - pub model_type: Option, - pub provider: GenerationProvider, - pub provider_job_id: Option, - pub frontend_caller: Option, - pub frontend_subscriber_id: Option, - pub frontend_subscriber_payload: Option, -} - -#[derive(Debug)] -#[derive(sqlx::FromRow)] -struct RawTask { - id: String, - task_status: String, - task_type: String, - model_type: Option, - provider: String, - provider_job_id: Option, - frontend_caller: Option, - frontend_subscriber_id: Option, - frontend_subscriber_payload: Option, -} - pub async fn list_tasks_by_provider_and_status( args: ListTasksByProviderAndStatusArgs<'_>, ) -> Result { diff --git a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs index a6abe5086b..1aff07d89c 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs @@ -1,5 +1,6 @@ use crate::connection::TaskDbConnection; use crate::error::SqliteTasksError; +use crate::queries::task::{RawTask, Task}; use enums::common::generation_provider::GenerationProvider; use enums::tauri::tasks::task_model_type::TaskModelType; use enums::tauri::tasks::task_status::TaskStatus; @@ -18,33 +19,6 @@ pub struct TaskList { pub tasks: Vec, } -#[derive(Debug, Clone)] -pub struct Task { - pub id: TaskId, - pub status: TaskStatus, - pub task_type: TaskType, - pub model_type: Option, - pub provider: GenerationProvider, - pub provider_job_id: Option, - pub frontend_caller: Option, - pub frontend_subscriber_id: Option, - pub frontend_subscriber_payload: Option, -} - -#[derive(Debug)] -#[derive(sqlx::FromRow)] -struct RawTask { - id: String, - task_status: String, - task_type: String, - model_type: Option, - provider: String, - provider_job_id: Option, - frontend_caller: Option, - frontend_subscriber_id: Option, - frontend_subscriber_payload: Option, -} - pub async fn list_tasks_by_provider_and_tokens( args: ListTasksArgs<'_>, ) -> Result { diff --git a/crates/schema/database/sqlite_tasks/src/queries/mod.rs b/crates/schema/database/sqlite_tasks/src/queries/mod.rs index a318edea1c..e5938771a7 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/mod.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/mod.rs @@ -1,7 +1,10 @@ pub mod create_task; +pub mod get_task_by_provider_and_provider_job_id; pub mod list_tasks_by_provider_and_status; pub mod list_tasks_by_provider_and_tokens; pub mod list_tasks_for_frontend; pub mod mark_task_as_dismissed; +pub mod task; pub mod update_successful_task_status_with_metadata; +pub mod update_successful_task_status_with_metadata_by_provider; pub mod update_task_status; \ No newline at end of file diff --git a/crates/schema/database/sqlite_tasks/src/queries/task.rs b/crates/schema/database/sqlite_tasks/src/queries/task.rs new file mode 100644 index 0000000000..27ebe63db1 --- /dev/null +++ b/crates/schema/database/sqlite_tasks/src/queries/task.rs @@ -0,0 +1,33 @@ +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_model_type::TaskModelType; +use enums::tauri::tasks::task_status::TaskStatus; +use enums::tauri::tasks::task_type::TaskType; +use enums::tauri::ux::tauri_command_caller::TauriCommandCaller; +use tokens::tokens::sqlite::tasks::TaskId; + +#[derive(Debug, Clone)] +pub struct Task { + pub id: TaskId, + pub status: TaskStatus, + pub task_type: TaskType, + pub model_type: Option, + pub provider: GenerationProvider, + pub provider_job_id: Option, + pub frontend_caller: Option, + pub frontend_subscriber_id: Option, + pub frontend_subscriber_payload: Option, +} + +#[derive(Debug)] +#[derive(sqlx::FromRow)] +pub (crate) struct RawTask { + pub (crate) id: String, + pub (crate) task_status: String, + pub (crate) task_type: String, + pub (crate) model_type: Option, + pub (crate) provider: String, + pub (crate) provider_job_id: Option, + pub (crate) frontend_caller: Option, + pub (crate) frontend_subscriber_id: Option, + pub (crate) frontend_subscriber_payload: Option, +} diff --git a/crates/schema/database/sqlite_tasks/src/queries/update_successful_task_status_with_metadata_by_provider.rs b/crates/schema/database/sqlite_tasks/src/queries/update_successful_task_status_with_metadata_by_provider.rs new file mode 100644 index 0000000000..12b25cc63f --- /dev/null +++ b/crates/schema/database/sqlite_tasks/src/queries/update_successful_task_status_with_metadata_by_provider.rs @@ -0,0 +1,64 @@ +use enums::common::generation_provider::GenerationProvider; +use crate::connection::TaskDbConnection; +use crate::error::SqliteTasksError; +use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass; +use enums::tauri::tasks::task_status::TaskStatus; +use tokens::tokens::batch_generations::BatchGenerationToken; +use tokens::tokens::media_files::MediaFileToken; + +const SUCCESSFUL_STATUS : &str = TaskStatus::CompleteSuccess.to_str(); + +pub struct UpdateSuccessfulTaskByProviderArgs<'a> { + pub db: &'a TaskDbConnection, + pub provider: GenerationProvider, + pub provider_job_id: &'a str, + pub maybe_batch_token: Option<&'a BatchGenerationToken>, + pub maybe_primary_media_file_token: Option<&'a MediaFileToken>, + pub maybe_primary_media_file_class: Option, + pub maybe_primary_media_file_cdn_url: Option<&'a str>, + pub maybe_primary_media_file_thumbnail_url_template: Option<&'a str>, +} + +/// Returns true if rows were updated. +pub async fn update_successful_task_status_with_metadata_by_provider( + args: UpdateSuccessfulTaskByProviderArgs<'_>, +) -> Result { + + // TODO(bt,2025-07-12): Fix this. The sqlx mysql queries never required temporaries + let provider_temp = args.provider.to_str(); + let provider_job_id_temp = args.provider_job_id; + let maybe_batch_token_temp = args.maybe_batch_token.map(|t| t.as_str()); + let maybe_primary_media_token_temp = args.maybe_primary_media_file_token.map(|t| t.as_str()); + let maybe_primary_media_class = args.maybe_primary_media_file_class.map(|c| c.to_str()); + + // TODO(bt,2025-07-15): We can't set a LIMIT without a certain compiler flag for SQLite ? + let query = sqlx::query!(r#" + UPDATE tasks + SET + task_status = ?, + on_complete_batch_token = ?, + on_complete_primary_media_file_token = ?, + on_complete_primary_media_file_class = ?, + on_complete_primary_media_file_cdn_url = ?, + on_complete_primary_media_file_thumbnail_url_template = ? + WHERE + provider = ? + AND provider_job_id = ? + "#, + SUCCESSFUL_STATUS, + maybe_batch_token_temp, + maybe_primary_media_token_temp, + maybe_primary_media_class, + args.maybe_primary_media_file_cdn_url, + args.maybe_primary_media_file_thumbnail_url_template, + provider_temp, + provider_job_id_temp, + ); + + // info!("query: {:?}", query.sql()); + + let res = query.execute(args.db.get_pool()).await?; + let rows_updated = res.rows_affected() > 0; + + Ok(rows_updated) +} 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 93b00b6c9b..99f08a361f 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 @@ -31,6 +31,8 @@ pub enum TaskModelType { Gemini25Flash, #[serde(rename = "gpt_image_1")] GptImage1, + #[serde(rename = "grok_image")] + GrokImage, #[serde(rename = "recraft_3")] Recraft3, @@ -84,6 +86,7 @@ impl TaskModelType { Self::FluxProKontextMax => "flux_pro_kontext_max", Self::Gemini25Flash => "gemini_25_flash", Self::GptImage1 => "gpt_image_1", + Self::GrokImage => "grok_image", Self::Recraft3 => "recraft_3", Self::Midjourney => "midjourney", // Video models @@ -114,6 +117,7 @@ impl TaskModelType { "flux_pro_kontext_max" => Ok(Self::FluxProKontextMax), "gemini_25_flash" => Ok(Self::Gemini25Flash), "gpt_image_1" => Ok(Self::GptImage1), + "grok_image" => Ok(Self::GrokImage), "recraft_3" => Ok(Self::Recraft3), "midjourney" => Ok(Self::Midjourney), // Video models @@ -147,6 +151,7 @@ impl TaskModelType { Self::FluxProKontextMax, Self::Gemini25Flash, Self::GptImage1, + Self::GrokImage, Self::Recraft3, Self::Midjourney, // Video models @@ -187,6 +192,7 @@ mod tests { assert_serialization(TaskModelType::FluxProKontextMax, "flux_pro_kontext_max"); assert_serialization(TaskModelType::Gemini25Flash, "gemini_25_flash"); assert_serialization(TaskModelType::GptImage1, "gpt_image_1"); + assert_serialization(TaskModelType::GrokImage, "grok_image"); assert_serialization(TaskModelType::Recraft3, "recraft_3"); assert_serialization(TaskModelType::Midjourney, "midjourney"); // Video models @@ -216,6 +222,7 @@ mod tests { assert_eq!(TaskModelType::FluxProKontextMax.to_str(), "flux_pro_kontext_max"); assert_eq!(TaskModelType::Gemini25Flash.to_str(), "gemini_25_flash"); assert_eq!(TaskModelType::GptImage1.to_str(), "gpt_image_1"); + assert_eq!(TaskModelType::GrokImage.to_str(), "grok_image"); assert_eq!(TaskModelType::Recraft3.to_str(), "recraft_3"); assert_eq!(TaskModelType::Midjourney.to_str(), "midjourney"); // Video models @@ -245,6 +252,7 @@ mod tests { assert_eq!(TaskModelType::from_str("flux_pro_kontext_max").unwrap(), TaskModelType::FluxProKontextMax); assert_eq!(TaskModelType::from_str("gemini_25_flash").unwrap(), TaskModelType::Gemini25Flash); assert_eq!(TaskModelType::from_str("gpt_image_1").unwrap(), TaskModelType::GptImage1); + assert_eq!(TaskModelType::from_str("grok_image").unwrap(), TaskModelType::GrokImage); assert_eq!(TaskModelType::from_str("recraft_3").unwrap(), TaskModelType::Recraft3); assert_eq!(TaskModelType::from_str("midjourney").unwrap(), TaskModelType::Midjourney); // Video models @@ -276,7 +284,7 @@ mod tests { #[test] fn all_variants() { let mut variants = TaskModelType::all_variants(); - assert_eq!(variants.len(), 22); + assert_eq!(variants.len(), 23); // Image models assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Dev)); assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Schnell)); @@ -287,6 +295,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(TaskModelType::FluxProKontextMax)); assert_eq!(variants.pop_first(), Some(TaskModelType::Gemini25Flash)); assert_eq!(variants.pop_first(), Some(TaskModelType::GptImage1)); + assert_eq!(variants.pop_first(), Some(TaskModelType::GrokImage)); assert_eq!(variants.pop_first(), Some(TaskModelType::Recraft3)); assert_eq!(variants.pop_first(), Some(TaskModelType::Midjourney)); // Video models diff --git a/frontend/libs/model-list/src/lib/lists/ImageModels.ts b/frontend/libs/model-list/src/lib/lists/ImageModels.ts index 8880321296..18e0f981ec 100644 --- a/frontend/libs/model-list/src/lib/lists/ImageModels.ts +++ b/frontend/libs/model-list/src/lib/lists/ImageModels.ts @@ -3,6 +3,21 @@ import { ModelCreator } from "../classes/metadata/ModelCreator.js"; import { ModelTag } from "../classes/metadata/ModelTag.js"; export const IMAGE_MODELS: ImageModel[] = [ + new ImageModel({ + id: "grok_image", + tauriId: "grok_image", + fullName: "Grok", + category: "image", + creator: ModelCreator.Grok, + selectorName: "Grok", + selectorDescription: "Fast af", + selectorBadges: ["10 sec."], + maxGenerationCount: 6, + defaultGenerationCount: 6, + canUseImagePrompt: false, + maxImagePromptCount: 6, + progressBarTime: 10000, + }), new ImageModel({ id: "midjourney", tauriId: "midjourney",