Grok image retry

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;

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<GrokWebsocket>,
+}
+
+
+pub async fn prompt_websocket_image_with_retry(args: PromptWebsocketImageWithRetryArgs<'_>) -> Result<PromptWebsocketImageWithRetryResult, GrokError>
+{
+  let mut maybe_new_websocket : Option<GrokWebsocket> = 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<String> = 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(())
+  }
+}

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::<Vec<_>>();
-
       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,