tool to poll images

crates/api_clients/grok_client/src/error/grok_client_error.rs

@@ -1,6 +1,6 @@
+use base64::DecodeError;
 use std::error::Error;
 use std::fmt::{Display, Formatter};
-use base64::DecodeError;
 
 #[derive(Debug)]
 pub enum GrokClientError {
@@ -43,6 +43,9 @@ pub enum GrokClientError {
 
   /// Signature algorithm isn't working with inputs
   BadSignatureInputs,
+  
+  /// Something is broken with timeout math
+  TimeoutMathBroken,
 
 //  /// An error reading the file for upload.
 //  FileForUploadReadError(std::io::Error),
@@ -99,6 +102,7 @@ impl Display for GrokClientError {
       Self::InvalidVerificationTokenBytes => write!(f, "Invalid verification token bytes"),
       Self::ScriptLogicOutOfDate => write!(f, "Script logic out of date"),
       Self::BadSignatureInputs => write!(f, "Bad signature inputs"),
+      Self::TimeoutMathBroken => write!(f, "Timeout math is broken"),
 
       //Self::FileForUploadReadError(err) => write!(f, "Error reading file for upload: {}", err),
       //Self::FileForUploadHasInvalidPath => write!(f, "The file path provided for upload is invalid."),

crates/api_clients/grok_client/src/requests/image_websocket/_metadata/websocket_messages/image_payloads.txt

@@ -0,0 +1,8 @@
+
+Received message: {"type":"image","blob":"...","job_id":"260270db-2515-470b-a730-1fe87e11d21c","grid_index":0,"is_alteration":false,"order":0,"section_id":0,"request_id":"57b94c87-ae13-4ade-8fc0-ebd39fb511ec","url":"https://imagine-public.x.ai/imagine-public/images/260270db-2515-470b-a730-1fe87e11d21c.png","id":"260270db-2515-470b-a730-1fe87e11d21c","prompt":"a dog riding a motorcycle","full_prompt":"a dog riding a motorcycle","percentage_complete":null,"model_name":"imagine_x_1","suggested_moderation_rule":null,"r_rated":false}
+
+Received message: {"type":"image","blob":"...","job_id":"260270db-2515-470b-a730-1fe87e11d21c","grid_index":0,"is_alteration":false,"order":0,"section_id":0,"request_id":"57b94c87-ae13-4ade-8fc0-ebd39fb511ec","url":"https://imagine-public.x.ai/imagine-public/images/260270db-2515-470b-a730-1fe87e11d21c.png","id":"260270db-2515-470b-a730-1fe87e11d21c","prompt":"a dog riding a motorcycle","full_prompt":"a dog riding a motorcycle","percentage_complete":null,"model_name":"imagine_x_1","suggested_moderation_rule":null,"r_rated":false}
+
+Received message: {"type":"image","blob":"...","job_id":"260270db-2515-470b-a730-1fe87e11d21c","grid_index":0,"is_alteration":false,"order":0,"section_id":0,"request_id":"57b94c87-ae13-4ade-8fc0-ebd39fb511ec","url":"https://imagine-public.x.ai/imagine-public/images/260270db-2515-470b-a730-1fe87e11d21c.png","id":"260270db-2515-470b-a730-1fe87e11d21c","prompt":"a dog riding a motorcycle","full_prompt":"A dog sits in a motorcycle sidecar on a paved road with a natural background.","percentage_complete":50.0,"model_name":"imagine_x_1","suggested_moderation_rule":null,"r_rated":false}
+
+Received message: {"type":"image","blob":"...","job_id":"260270db-2515-470b-a730-1fe87e11d21c","grid_index":0,"is_alteration":false,"order":0,"section_id":0,"request_id":"57b94c87-ae13-4ade-8fc0-ebd39fb511ec","url":"https://imagine-public.x.ai/imagine-public/images/260270db-2515-470b-a730-1fe87e11d21c.png","id":"260270db-2515-470b-a730-1fe87e11d21c","prompt":"a dog riding a motorcycle","full_prompt":"A dog sits in a motorcycle sidecar on a paved road with a natural background.","percentage_complete":100.0,"model_name":"imagine_x_1","suggested_moderation_rule":"MODERATION_RULE_DEFAULT","r_rated":false}

crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_images.rs

@@ -0,0 +1,154 @@
+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::messages::image_data_message::MessageImageData;
+use log::info;
+use std::time::{Duration, Instant};
+use wreq::ws::message::Message;
+
+// Number of images we get per prompt by default
+const DEFAULT_IMAGE_COUNT: usize = 6;
+
+pub struct ListenForWebsocketImagesArgs<'a> {
+  pub websocket: &'a ClonableWebsocket,
+  pub timeout: Duration,
+}
+
+#[derive(Clone, Debug)]
+pub struct ImageResults {
+  pub images: Vec<ImageData>,
+}
+
+#[derive(Clone, Debug)]
+pub struct ImageData {
+  pub url: String,
+  pub user_prompt: String,
+  pub enriched_prompt: String,
+}
+
+pub async fn listen_for_website_images(args: ListenForWebsocketImagesArgs<'_>) -> Result<ImageResults, GrokError> {
+  let start = Instant::now();
+  let end_at = start.checked_add(args.timeout);
+  let end_at = match end_at {
+    Some(end_at) => end_at,
+    None => return Err(GrokClientError::TimeoutMathBroken.into()),
+  };
+
+  let mut images = Vec::new();
+
+  loop {
+    if end_at < Instant::now() {
+      break;
+    }
+
+    let maybe_message = args.websocket.try_next_timeout(Duration::from_millis(500)).await?;
+
+    let maybe_message = match maybe_message {
+      None => None,
+      Some(Message::Text(text)) => Some(text.to_string()),
+      Some(_) => {
+        // Got a ping, pong, binary payload, etc.
+        None
+      },
+    };
+
+    let message = match maybe_message {
+      Some(message) => message,
+      None => continue,
+    };
+
+    let maybe_image_data = serde_json::from_str::<MessageImageData>(&message);
+
+    let image_data = match maybe_image_data {
+      Ok(image_data) => image_data,
+      Err(_err) => {
+        // NB: Might not have been a parse error, since we're parsing irrelevant message payloads too.
+        // This is slightly dangerous to do this way as we might mask true parsing errors with image payloads.
+        continue;
+      }
+    };
+
+    match image_data.percentage_complete {
+      None => continue,
+      Some(percent) => {
+        if percent < 100.0 {
+          continue;
+        }
+      }
+    }
+
+    images.push(ImageData {
+      url : image_data.url,
+      enriched_prompt: image_data.full_prompt,
+      user_prompt: image_data.prompt,
+    });
+
+    if images.len() >= DEFAULT_IMAGE_COUNT {
+      info!("{} images generated; we're done polling.", images.len());
+      break;
+    }
+  }
+
+  Ok(ImageResults {
+    images,
+  })
+}
+
+
+#[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::listen_for_websocket_images::{listen_for_website_images, ListenForWebsocketImagesArgs};
+  use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs};
+  use crate::test_utils::get_test_cookies::get_test_cookies;
+  use errors::AnyhowResult;
+  use log::LevelFilter;
+  use std::io::Write;
+  use std::time::Duration;
+  use crate::test_utils::setup_test_logging::setup_test_logging;
+
+  #[tokio::test]
+  #[ignore] // manually test
+  async fn prompt() -> AnyhowResult<()> {
+    setup_test_logging(LevelFilter::Info);
+
+    let cookies = get_test_cookies()?;
+
+    let websocket = create_listen_websocket(CreateListenWebsocketArgs {
+      cookies: &cookies,
+    }).await?;
+
+    let websocket = ClonableWebsocket::new(websocket);
+
+    println!("Sending...");
+    std::io::stdout().flush()?;
+
+    let _result = prompt_websocket_image(PromptWebsocketImageArgs {
+      websocket_wrapped: websocket.clone(),
+      prompt: "a dog riding a motorcycle",
+    }).await?;
+
+    println!("Reading...");
+    std::io::stdout().flush()?;
+
+    println!("Polling.");
+
+    let images = listen_for_website_images(ListenForWebsocketImagesArgs {
+      websocket: &websocket,
+      timeout: Duration::from_millis(10_000),
+    }).await?;
+
+    println!("Done polling. Images: {}", images.images.len());
+
+    for image in images.images {
+      println!("Image: {:?}", image);
+    }
+
+    log::logger().flush();
+
+    assert_eq!(1,2);
+
+    Ok(())
+  }
+}

crates/api_clients/grok_client/src/requests/image_websocket/messages/image_data_message.rs

@@ -0,0 +1,63 @@
+use serde::Deserialize;
+
+
+/// Images with a binary blob and URLs.
+/// We may receive several of these for a single prompt.
+#[derive(Deserialize, Clone, Debug)]
+pub struct MessageImageData {
+  /// UUID.
+  /// Probably an identifier for the original request/prompt.
+  pub request_id: Option<String>,
+
+  /// UUID.
+  /// This turns out to be the file ID in the image URL (and `job_id`).
+  pub id: Option<String>,
+
+  /// UUID.
+  /// This turns out to be the file ID in the image URL (and `id`).
+  pub job_id: Option<String>,
+
+  /// Eg. null, "0", "50", "100"
+  pub percentage_complete: Option<f32>,
+
+  /// The user's prompt
+  pub prompt: String,  //pub prompt: Option<String>,
+
+  /// The X.ai enriched prompt
+  pub full_prompt: String, //pub full_prompt: Option<String>,
+
+  /// URL to the image.
+  pub url: String, //pub url: Option<String>,
+  
+  // We don't need to decode - let's just download again
+  // /// Base64 encoded image blob.
+  // pub blob: Option<String>,
+
+  // /// NSFW flag
+  // pub r_rated: Option<bool>,
+
+  // /// Name of the model used to generate the image
+  // /// eg. "imagine_h_1", "imagine_x_1"
+  // pub model_name: Option<String>,
+}
+
+/*
+Received websocket message: {
+  "type": "image",
+  "blob": "...",
+  "job_id": "260270db-2515-470b-a730-1fe87e11d21c",
+  "grid_index": 0,
+  "is_alteration": false,
+  "order": 0,
+  "section_id": 0,
+  "request_id": "57b94c87-ae13-4ade-8fc0-ebd39fb511ec",
+  "url": "https://imagine-public.x.ai/imagine-public/images/260270db-2515-470b-a730-1fe87e11d21c.png",
+  "id": "260270db-2515-470b-a730-1fe87e11d21c",
+  "prompt": "a dog riding a motorcycle",
+  "full_prompt": "A dog sits in a motorcycle sidecar on a paved road with a natural background.",
+  "percentage_complete": 100.0,
+  "model_name": "imagine_x_1",
+  "suggested_moderation_rule": "MODERATION_RULE_DEFAULT",
+  "r_rated":false
+}
+*/

crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_server_message.rs

@@ -35,6 +35,8 @@ pub struct ImageDataMessage {
   pub percentage_complete: Option<f32>,
 
   pub prompt: Option<String>,
+  
+  /// The enriched prompt
   pub full_prompt: Option<String>,
 
   /// Base64 encoded image blob.

crates/api_clients/grok_client/src/requests/image_websocket/mod.rs

@@ -1,4 +1,5 @@
 pub mod clonable_websocket;
 pub mod create_listen_websocket;
 pub mod messages;
-pub mod prompt_websocket_image;
+pub mod prompt_websocket_image;
+mod listen_for_websocket_images;

crates/api_clients/grok_client/src/requests/image_websocket/prompt_websocket_image.rs

@@ -33,7 +33,6 @@ mod tests {
   #[ignore] // manually test
   async fn prompt() -> AnyhowResult<()> {
     //setup_test_logging(LevelFilter::Trace);
-    //let cookies = SESSION_COOKIES_WITHOUT_CF_CLEARANCE;
 
     let cookies = get_test_cookies()?;
 
@@ -46,7 +45,6 @@ mod tests {
     println!("Sending...");
     std::io::stdout().flush()?;
 
-
     let result = prompt_websocket_image(PromptWebsocketImageArgs {
       websocket_wrapped: websocket.clone(),
       prompt: "a dog riding a motorcycle",
@@ -151,7 +149,7 @@ mod tests {
       // TODO: Video websocket APIs
       // TODO: Upload APIs
 
-      if count > 30 {
+      if count > 5 {
         println!("No messages after 5 seconds");
         break;
       }

crates/api_clients/grok_client/src/utils/mod.rs

@@ -1,3 +1,6 @@
-pub mod create_firefox_client;
 pub mod user_and_file_id_to_image_url;
-pub mod user_and_file_id_to_video_url;
+pub mod user_and_file_id_to_video_url;
+
+// Just for the crate
+pub (crate) mod create_firefox_client;
+pub (crate) mod scrub_blobs_for_debug_logging;

crates/api_clients/grok_client/src/utils/scrub_blobs_for_debug_logging.rs

@@ -0,0 +1,12 @@
+use once_cell::sync::Lazy;
+use regex::Regex;
+
+static LARGE_BLOB_REGEX: Lazy<Regex> = Lazy::new(|| {
+  Regex::new(r#"[^"]{100,}"#)
+      .expect("Regex should parse")
+});
+
+/// Helpful for debugging large websocket payloads with in-stream base64-encoded image binaries.
+pub fn scrub_blobs_for_debug_logging(blob: &str) -> String {
+  LARGE_BLOB_REGEX.replace(blob, "...").to_string()
+}