prompt images wip

Cargo.lock

@@ -1998,7 +1998,7 @@ checksum = "d38f2da7a0a2c4ccf0065be06397cc26a81f4e528be095826eee9d4adbb8c60f"
 dependencies = [
  "byteorder",
  "fnv",
- "uuid 1.8.0",
+ "uuid 1.18.1",
 ]
 
 [[package]]
@@ -4690,6 +4690,7 @@ dependencies = [
  "testing",
  "tokio",
  "url",
+ "uuid 1.18.1",
  "workspace-build-acceleration",
  "wreq",
  "wreq-util",
@@ -5433,7 +5434,7 @@ dependencies = [
 name = "idempotency"
 version = "0.0.1"
 dependencies = [
- "uuid 1.8.0",
+ "uuid 1.18.1",
  "workspace-build-acceleration",
 ]
 
@@ -9193,7 +9194,7 @@ dependencies = [
  "rkyv_derive",
  "seahash",
  "tinyvec",
- "uuid 1.8.0",
+ "uuid 1.18.1",
 ]
 
 [[package]]
@@ -9587,7 +9588,7 @@ dependencies = [
  "serde",
  "serde_json",
  "url",
- "uuid 1.8.0",
+ "uuid 1.18.1",
 ]
 
 [[package]]
@@ -10825,7 +10826,7 @@ dependencies = [
  "tokio",
  "url",
  "url-builder",
- "uuid 1.8.0",
+ "uuid 1.18.1",
  "workspace-build-acceleration",
 ]
 
@@ -11390,7 +11391,7 @@ dependencies = [
  "thiserror 2.0.12",
  "time",
  "url",
- "uuid 1.8.0",
+ "uuid 1.18.1",
  "walkdir",
 ]
 
@@ -11615,7 +11616,7 @@ dependencies = [
  "toml 0.9.5",
  "url",
  "urlpattern",
- "uuid 1.8.0",
+ "uuid 1.18.1",
  "walkdir",
 ]
 
@@ -12814,12 +12815,14 @@ dependencies = [
 
 [[package]]
 name = "uuid"
-version = "1.8.0"
+version = "1.18.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a183cf7feeba97b4dd1c0d46788634f6221d87fa961b305bed08c851829efcc0"
+checksum = "2f87b8aa10b915a06587d0dec516c282ff295b475d94abf425d62b57710070a2"
 dependencies = [
- "getrandom 0.2.14",
+ "getrandom 0.3.1",
+ "js-sys",
  "serde",
+ "wasm-bindgen",
 ]
 
 [[package]]
@@ -14028,7 +14031,7 @@ dependencies = [
  "unicode-bidi",
  "unicode-normalization",
  "url",
- "uuid 1.8.0",
+ "uuid 1.18.1",
  "void",
  "winapi 0.3.9",
  "windows-sys 0.48.0",

Cargo.toml

@@ -262,6 +262,7 @@ thiserror = "2.0.12"
 tokio = { version = "1.45.1", features = ["macros"] }
 toml = "0.8.22"
 url = { version = "2.5.4", features = ["serde"] }
+uuid = "1.18.1"
 
 # NB: wreq gets past fingerprinting, but we have to deal with platform-specific build issues.
 #

crates/api_clients/grok_client/Cargo.toml

@@ -23,12 +23,14 @@ errors.workspace = true
 
 # External
 anyhow.workspace = true
+chrono.workspace = true
 futures.workspace = true
 log.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 tokio = { version = "1", features = ["full"] }
 url.workspace = true
+uuid.workspace = true
 wreq-util.workspace = true
 wreq.workspace = true
 

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

@@ -7,6 +7,18 @@ pub enum GrokClientError {
   /// An error was encountered in building the Wreq client
   WreqClientError(wreq::Error),
 
+  /// Error serializing a message to send the websocket
+  WebsocketRequestSerializationError(serde_json::Error),
+
+  /// Error locking the websocket for sending/receiving
+  WebsocketLockError,
+
+  /// Error reading from a websocket.
+  WebsocketReadError(wreq::Error),
+
+  /// Error sending to an open websocket.
+  WebsocketSendError(wreq::Error),
+
 //  /// An error reading the file for upload.
 //  FileForUploadReadError(std::io::Error),
 //
@@ -34,10 +46,6 @@ pub enum GrokClientError {
 //  /// Error parsing a request URL.
 //  UrlParseError(url::ParseError),
 //
-//
-//  /// Error serializing the sentinel token to JSON
-//  CouldNotSerializeSentinelToken(serde_json::Error),
-//
 //  /// Error serializing the sentinel store token to JSON (typically for persistent storage).
 //  CouldNotSerializeSentinelTokenStore(serde_json::Error),
 //
@@ -54,6 +62,10 @@ impl Display for GrokClientError {
   fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
     match self {
       Self::WreqClientError(err) => write!(f, "Wreq client error (during client creation): {}", err),
+      Self::WebsocketRequestSerializationError(err) => write!(f, "Websocket request serialization error: {}", err),
+      Self::WebsocketLockError => write!(f, "Websocket lock error"),
+      Self::WebsocketReadError(err) => write!(f, "Websocket read error: {}", err),
+      Self::WebsocketSendError(err) => write!(f, "Websocket send error: {}", err),
 
       //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/listen_websocket/clonable_websocket.rs

@@ -0,0 +1,63 @@
+use crate::error::grok_client_error::GrokClientError;
+use crate::error::grok_error::GrokError;
+use futures::stream::FusedStream;
+use futures::TryStreamExt;
+use log::warn;
+use serde::Serialize;
+use std::sync::{Arc, RwLock};
+use wreq::ws::message::Message;
+use wreq::ws::WebSocket;
+
+#[derive(Clone)]
+pub struct ClonableWebsocket {
+  pub(crate) websocket: Arc<RwLock<WebSocket>>,
+}
+
+impl ClonableWebsocket {
+  pub fn new(websocket: WebSocket) -> Self {
+    Self {
+      websocket: Arc::new(RwLock::new(websocket)),
+    }
+  }
+
+  pub async fn send(&self, message: String) -> Result<(), GrokError> {
+    match self.websocket.write() {
+      Err(err) => Err(GrokClientError::WebsocketLockError.into()),
+      Ok(mut websocket) => {
+        let message = Message::text(message);
+        websocket.send(message)
+            .await
+            .map_err(|err| {
+              GrokClientError::WreqClientError(err)
+            })?;
+        Ok(())
+      },
+    }
+  }
+
+  pub async fn send_serializable<T: Serialize>(&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 is_terminated(&self) -> Result<bool, GrokClientError> {
+    match self.websocket.read() {
+      Err(_) => Err(GrokClientError::WebsocketLockError.into()),
+      Ok(websocket) => Ok(websocket.is_terminated()),
+    }
+  }
+
+  pub async fn try_next(&self) -> Result<Option<Message>, GrokClientError> {
+    match self.websocket.write() {
+      Err(_) => Err(GrokClientError::WebsocketLockError),
+      Ok(mut websocket) => websocket
+          .try_next()
+          .await
+          .map_err(|err| GrokClientError::WebsocketReadError(err)),
+    }
+  }
+}

crates/api_clients/grok_client/src/requests/listen_websocket/messages/metadata/create_image_request.json

@@ -0,0 +1,21 @@
+{
+  "type": "conversation.item.create",
+  "timestamp": 1761010203630,
+  "item": {
+    "type": "message",
+    "content": [
+      {
+        "requestId": "fc39cefd-fcf9-4775-b490-7e17ce670060",
+        "text": "a dog typing on a keyboard",
+        "type": "input_text",
+        "properties": {
+          "section_count": 0,
+          "is_kids_mode": false,
+          "enable_nsfw": true,
+          "skip_upsampler": false,
+          "is_initial": false
+        }
+      }
+    ]
+  }
+}

crates/api_clients/grok_client/src/requests/listen_websocket/messages/mod.rs

@@ -0,0 +1 @@
+pub mod websocket_request;

crates/api_clients/grok_client/src/requests/listen_websocket/messages/websocket_request.rs

@@ -0,0 +1,75 @@
+use chrono::Utc;
+use serde::Serialize;
+use uuid::Uuid;
+
+#[derive(Serialize)]
+pub struct WebsocketRequest {
+  /// Type of message
+  /// eg. "conversation.item.create"
+  pub r#type: String,
+
+  /// Unix timestamp
+  pub timestamp: u64,
+
+  pub item: WebsocketRequestItem,
+}
+
+#[derive(Serialize)]
+pub struct WebsocketRequestItem {
+  /// Type of item
+  /// eg. "message"
+  pub r#type: String,
+
+  pub content: Vec<WebsocketRequestItemContent>,
+}
+
+#[derive(Serialize)]
+pub struct WebsocketRequestItemContent {
+  /// UUID
+  #[serde(rename="requestId")]
+  pub request_id: String,
+
+  /// Type of item content
+  /// eg. "input_text"
+  pub r#type: String,
+
+  /// The prompt for the request
+  pub text: String,
+
+  pub properties: WebsocketRequestItemContentProperties,
+}
+
+#[derive(Serialize)]
+pub struct WebsocketRequestItemContentProperties {
+  pub section_count: usize,
+  pub is_kids_mode: bool,
+  pub enable_nsfw: bool,
+  pub skip_upsampler: bool,
+  pub is_initial: bool,
+}
+
+impl WebsocketRequest {
+  pub fn new_image_prompt(prompt: &str) -> Self {
+    Self {
+      r#type: "conversation.item.create".to_string(),
+      timestamp: Utc::now().timestamp_millis() as u64,
+      item: WebsocketRequestItem {
+        r#type: "message".to_string(),
+        content: vec![
+          WebsocketRequestItemContent {
+            request_id: Uuid::new_v4().to_string(),
+            r#type: "input_text".to_string(),
+            text: prompt.to_string(),
+            properties: WebsocketRequestItemContentProperties {
+              section_count: 0,
+              is_kids_mode: false,
+              enable_nsfw: true,
+              skip_upsampler: false,
+              is_initial: false,
+            },
+          }
+        ],
+      },
+    }
+  }
+}

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

@@ -1,2 +1,6 @@
 pub (crate) mod cookies;
+
+pub mod prompt_websocket_image;
+pub mod clonable_websocket;
 pub mod create_listen_websocket;
+pub mod messages;

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

@@ -0,0 +1,71 @@
+use log::warn;
+use wreq::ws::WebSocket;
+use crate::error::grok_client_error::GrokClientError;
+use crate::error::grok_error::GrokError;
+use crate::requests::listen_websocket::clonable_websocket::ClonableWebsocket;
+use crate::requests::listen_websocket::messages::websocket_request::{WebsocketRequest, WebsocketRequestItem};
+
+pub struct PromptWebsocketImageArgs<'a> {
+  pub websocket_wrapped: ClonableWebsocket,
+  pub prompt: &'a str,
+}
+
+
+pub async fn prompt_websocket_image(args: PromptWebsocketImageArgs<'_>) -> Result<(), GrokError> {
+  let message = WebsocketRequest::new_image_prompt(args.prompt);
+
+  args.websocket_wrapped.send_serializable(&message).await?;
+
+  Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+  use std::io::Write;
+  use anyhow::anyhow;
+  use futures::TryStreamExt;
+  use log::{warn, LevelFilter};
+  use errors::AnyhowResult;
+  use crate::requests::listen_websocket::clonable_websocket::ClonableWebsocket;
+  use crate::requests::listen_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs};
+  use crate::requests::listen_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs};
+  use crate::test_utils::get_test_cookies::get_test_cookies;
+  use crate::test_utils::setup_test_logging::setup_test_logging;
+
+  #[tokio::test]
+  #[ignore] // manually test
+  async fn prompt() -> AnyhowResult<()> {
+    setup_test_logging(LevelFilter::Trace);
+    //let cookies = SESSION_COOKIES_WITHOUT_CF_CLEARANCE;
+
+    let cookies = get_test_cookies()?;
+
+    let mut 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 skateboard",
+    }).await?;
+
+    println!("Reading...");
+    std::io::stdout().flush()?;
+
+    if let Some(message) = websocket.try_next().await? {
+      warn!("Received websocket message: {:?}", message);
+    }
+
+    log::logger().flush();
+
+    assert_eq!(1,2);
+
+    Ok(())
+  }
+}