diff --git a/crates/api_clients/grok_client/src/error/grok_client_error.rs b/crates/api_clients/grok_client/src/error/grok_client_error.rs index f1c8f5be96..2efa87fdee 100644 --- a/crates/api_clients/grok_client/src/error/grok_client_error.rs +++ b/crates/api_clients/grok_client/src/error/grok_client_error.rs @@ -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."), diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/_metadata/websocket_messages/image_payloads.txt b/crates/api_clients/grok_client/src/requests/image_websocket/_metadata/websocket_messages/image_payloads.txt new file mode 100644 index 0000000000..1a77ccc7f1 --- /dev/null +++ b/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} 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 new file mode 100644 index 0000000000..56804491d0 --- /dev/null +++ b/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, +} + +#[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 { + 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::(&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(()) + } +} diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/messages/image_data_message.rs b/crates/api_clients/grok_client/src/requests/image_websocket/messages/image_data_message.rs index e69de29bb2..19db5a0a8b 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/messages/image_data_message.rs +++ b/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, + + /// UUID. + /// This turns out to be the file ID in the image URL (and `job_id`). + pub id: Option, + + /// UUID. + /// This turns out to be the file ID in the image URL (and `id`). + pub job_id: Option, + + /// Eg. null, "0", "50", "100" + pub percentage_complete: Option, + + /// The user's prompt + pub prompt: String, //pub prompt: Option, + + /// The X.ai enriched prompt + pub full_prompt: String, //pub full_prompt: Option, + + /// URL to the image. + pub url: String, //pub url: Option, + + // We don't need to decode - let's just download again + // /// Base64 encoded image blob. + // pub blob: Option, + + // /// NSFW flag + // pub r_rated: Option, + + // /// Name of the model used to generate the image + // /// eg. "imagine_h_1", "imagine_x_1" + // pub model_name: Option, +} + +/* +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 +} +*/ diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_server_message.rs b/crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_server_message.rs index 4ec735e27f..a3b17362ee 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_server_message.rs +++ b/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, pub prompt: Option, + + /// The enriched prompt pub full_prompt: Option, /// Base64 encoded image blob. 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 c171c7eb4a..f065930694 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,4 +1,5 @@ pub mod clonable_websocket; pub mod create_listen_websocket; pub mod messages; -pub mod prompt_websocket_image; \ No newline at end of file +pub mod prompt_websocket_image; +mod listen_for_websocket_images; \ 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 51e8df4851..113a061225 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 @@ -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; } diff --git a/crates/api_clients/grok_client/src/utils/mod.rs b/crates/api_clients/grok_client/src/utils/mod.rs index c3168d1512..22c38941b9 100644 --- a/crates/api_clients/grok_client/src/utils/mod.rs +++ b/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; \ No newline at end of file +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; diff --git a/crates/api_clients/grok_client/src/utils/scrub_blobs_for_debug_logging.rs b/crates/api_clients/grok_client/src/utils/scrub_blobs_for_debug_logging.rs new file mode 100644 index 0000000000..fa1d3404cc --- /dev/null +++ b/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 = 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() +}