diff --git a/crates/api_clients/grok_client/src/recipes/mod.rs b/crates/api_clients/grok_client/src/recipes/mod.rs index e774ecd05b..68bb011571 100644 --- a/crates/api_clients/grok_client/src/recipes/mod.rs +++ b/crates/api_clients/grok_client/src/recipes/mod.rs @@ -1,3 +1,4 @@ +pub mod prompt_websocket_image_with_retry; pub mod request_client_secrets; pub mod upload_image_and_generate_video; pub mod upload_image_and_generate_video_with_retry; \ No newline at end of file diff --git a/crates/api_clients/grok_client/src/recipes/prompt_websocket_image_with_retry.rs b/crates/api_clients/grok_client/src/recipes/prompt_websocket_image_with_retry.rs new file mode 100644 index 0000000000..8d2d90ee01 --- /dev/null +++ b/crates/api_clients/grok_client/src/recipes/prompt_websocket_image_with_retry.rs @@ -0,0 +1,189 @@ +use crate::credentials::grok_cookies::GrokCookies; +use crate::error::grok_client_error::GrokClientError; +use crate::error::grok_error::GrokError; +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::messages::websocket_client_message::{ClientMessageAspectRatio, WebsocketClientMessage}; +use log::{error, info}; + +pub struct PromptWebsocketImageWithRetryArgs<'a> { + pub websocket: &'a mut GrokWebsocket, + pub prompt: &'a str, + pub aspect_ratio: ClientMessageAspectRatio, + pub cookies: &'a GrokCookies, +} + +pub struct PromptWebsocketImageWithRetryResult { + pub maybe_new_websocket: Option, +} + + +pub async fn prompt_websocket_image_with_retry(args: PromptWebsocketImageWithRetryArgs<'_>) -> Result +{ + let mut maybe_new_websocket : Option = None; + + for i in 0..3 { + info!("Starting Grok Websocket image generation (attempt {}) ...", (i+1)); + + let message = WebsocketClientMessage::new_image_prompt(args.prompt, args.aspect_ratio); + + let result = + if let Some(websocket) = maybe_new_websocket.as_mut() { + websocket.send_serializable(&message).await + } else { + args.websocket.send_serializable(&message).await + }; + + match result { + Ok(()) => { + break; + } + Err(GrokError::Client(GrokClientError::WreqClientError(err))) => { + error!("Error sending message to websocket (WreqClientError): {}", err); + } + Err(GrokError::Client(err)) => { + error!("Error sending message to websocket (GrokClientError): {}", err); + } + Err(GrokError::ApiSpecific(err)) => { + error!("Error sending message to websocket (GrokSpecificApiError): {}", err); + } + Err(GrokError::ApiGeneric(err)) => { + error!("Error sending message to websocket (GrokGenericApiError): {}", err); + } + } + + let result = create_listen_websocket(CreateListenWebsocketArgs { + cookies: args.cookies.as_str(), + }).await?; + + maybe_new_websocket = Some(GrokWebsocket::new(result)); + } + + Ok(PromptWebsocketImageWithRetryResult { + maybe_new_websocket, + }) +} + +#[cfg(test)] +mod tests { + use crate::recipes::prompt_websocket_image_with_retry::{prompt_websocket_image_with_retry, PromptWebsocketImageWithRetryArgs}; + 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_client_message::ClientMessageAspectRatio; + 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, get_typed_test_cookies}; + use errors::AnyhowResult; + use futures::SinkExt; + use log::warn; + use std::io::Write; + use std::time::Duration; + use wreq::ws::message::{CloseCode, Message, Utf8Bytes}; + + #[tokio::test] + #[ignore] // manually test + async fn prompt() -> AnyhowResult<()> { + //setup_test_logging(LevelFilter::Trace); + + let cookies = get_typed_test_cookies()?; + + let websocket = create_listen_websocket(CreateListenWebsocketArgs { + cookies: cookies.as_str(), + }).await?; + + //let websocket = GrokWrappedWebsocket::new(websocket); + let mut websocket = GrokWebsocket::new(websocket); + + //println!("Sending..."); + //std::io::stdout().flush()?; + + //println!("For our test, we'll close the websocket and retry..."); + //websocket.websocket + // .close(CloseCode::NORMAL, Utf8Bytes::try_from(Vec::new())?).await?; + + //let mut websocket = GrokWebsocket::new(websocket.websocket); + + let result = prompt_websocket_image_with_retry(PromptWebsocketImageWithRetryArgs { + websocket: &mut websocket, + prompt: "a dog riding a motorcycle", + aspect_ratio: ClientMessageAspectRatio::WideThreeByTwo, + cookies: &cookies, + }).await?; + + println!("Reading..."); + std::io::stdout().flush()?; + + let mut count = 0; + + println!("Polling."); + + loop { + //let maybe_message = + // websocket.get_response_with_timeout(Duration::from_millis(1000)).await?; + + let maybe_message = + websocket.try_next_timeout(Duration::from_millis(1000)).await?; + + let mut maybe_raw_text : Option = None; + + let maybe_message = match maybe_message { + None => None, + Some(Message::Text(text)) => { + let maybe_message = WebsocketServerMessage::from_json_str(text.as_str())?; + let mut text = text.as_str().to_string(); + text.truncate(100); + //println!("Received message: {}", text); + maybe_raw_text = Some(text); + Some(maybe_message) + }, + Some(_) => { + warn!("Received non-text websocket message."); + None + }, + }; + + match maybe_message { + None => { + println!("No message received within timeout."); + count = count + 1; + } + Some(message) => { + match message { + WebsocketServerMessage::Image(image) => { + println!("IMAGE: {:?}", image.percentage_complete); + } + WebsocketServerMessage::Json(json) => { + println!("JSON : {:?}", json.percentage_complete); + } + WebsocketServerMessage::Unknown(unknown) => { + //let typ = unknown.get("type"); + let unknown_string = unknown.as_str(); + println!("[UNKNOWN] websocket message: {:?}", maybe_raw_text); + } + } + count = 0; + } + } + + // TODO: Make sure we capture all the events for images + // TODO: Explore the other APIs + // TODO understand the IDs + // TODO: Video websocket APIs + // TODO: Upload APIs + + if count > 5 { + println!("No messages after 5 seconds"); + break; + } + } + + println!("Done polling."); + + log::logger().flush(); + + assert_eq!(1,2); + + Ok(()) + } +} 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 index 4c065d4b2b..69f22880ca 100644 --- 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 @@ -18,6 +18,7 @@ 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::recipes::prompt_websocket_image_with_retry::{prompt_websocket_image_with_retry, PromptWebsocketImageWithRetryArgs}; 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; @@ -86,47 +87,33 @@ async fn inner_loop( continue; } - let maybe_cookie_header = grok_creds.maybe_copy_cookie_header_string()?; + let maybe_cookies = grok_creds.maybe_copy_typed_cookies()?; - let cookie_header = match maybe_cookie_header { - Some(cookie_header) => cookie_header, + //let maybe_cookie_header = grok_creds.maybe_copy_cookie_header_string()?; + + let cookies = match maybe_cookies { + Some(cookies) => cookies, None => { tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; continue; - } + }, }; + //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, + cookies: cookies.as_str(), }).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 { @@ -137,14 +124,20 @@ async fn inner_loop( } }; - info!("Prompting Grok websocket: {}", prompt_item.prompt); + info!("Prompt received over in-memory queue. Prompting Grok websocket: {}", prompt_item.prompt); - let _result = prompt_websocket_image(PromptWebsocketImageArgs { + let result = prompt_websocket_image_with_retry(PromptWebsocketImageWithRetryArgs { websocket: &mut websocket, prompt: &prompt_item.prompt, aspect_ratio: prompt_item.aspect_ratio, + cookies: &cookies, }).await?; + if let Some(new_websocket) = result.maybe_new_websocket { + info!("Replacing websocket with new one..."); + websocket = new_websocket; + } + let images = listen_for_websocket_images(ListenForWebsocketImagesArgs { websocket: &mut websocket, timeout: Duration::from_millis(30_000), @@ -164,19 +157,6 @@ async fn inner_loop( } } -//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,