wip

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

@@ -8,6 +8,7 @@ use std::sync::{Arc, RwLock};
 use std::time::Duration;
 use wreq::ws::message::Message;
 use wreq::ws::WebSocket;
+use crate::requests::listen_websocket::messages::websocket_server_message::WebsocketServerMessage;
 
 #[derive(Clone)]
 pub struct ClonableWebsocket {
@@ -45,6 +46,38 @@ impl ClonableWebsocket {
     self.send(message_json).await
   }
 
+  pub async fn get_response_with_timeout(&self, duration: Duration) -> Result<Option<WebsocketServerMessage>, GrokError> {
+    match self.websocket.write() {
+      Err(_) => Err(GrokClientError::WebsocketLockError.into()),
+      Ok(mut websocket) => {
+        let result = tokio::time::timeout(duration, websocket.try_next()).await;
+
+        match result {
+          Err(elapsed) => {
+            info!("Websocket try_next() elapsed without receiving a message: {:?}", elapsed);
+            Ok(None) // Timeout elapsed.
+          }
+          Ok(inner) => {
+            let maybe_message = inner.map_err(GrokClientError::WebsocketReadError)?;
+
+            match maybe_message {
+              None => Ok(None),
+              Some(Message::Text(text)) => {
+                let maybe_message = WebsocketServerMessage::from_json_str(&text)?;
+                Ok(Some(maybe_message))
+              },
+              Some(_) => {
+                warn!("Received non-text websocket message.");
+                Ok(None)
+              },
+            }
+          }
+        }
+      },
+    }
+  }
+
+
   pub async fn try_next_timeout(&self, duration: Duration) -> Result<Option<Message>, GrokClientError> {
     match self.websocket.write() {
       Err(_) => Err(GrokClientError::WebsocketLockError),

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

@@ -1,3 +1,3 @@
 pub mod websocket_client_message;
-mod websocket_server_message;
-mod image_data_message;
+pub mod websocket_server_message;
+pub mod image_data_message;

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

@@ -1,6 +1,7 @@
+use crate::error::grok_generic_api_error::GrokGenericApiError;
 use serde::Deserialize;
 
-#[derive(Deserialize, Clone)]
+#[derive(Deserialize, Clone, Debug)]
 #[serde(tag = "type")]
 pub enum WebsocketServerMessage {
   /// Json with root `type = "image"`
@@ -9,8 +10,7 @@ pub enum WebsocketServerMessage {
 
   /// Json with root `type = "json"`
   #[serde(rename = "json")]
-  JsonData {
-  },
+  JsonData(JsonDataMessage),
 
   /// Captures any other JSON messages.
   #[serde(untagged)]
@@ -19,7 +19,7 @@ pub enum WebsocketServerMessage {
 
 /// Images with a binary blob and URLs.
 /// We may receive several of these for a single prompt.
-#[derive(Deserialize, Clone)]
+#[derive(Deserialize, Clone, Debug)]
 pub struct ImageDataMessage {
   /// UUID.
   pub id: Option<String>,
@@ -50,11 +50,11 @@ pub struct ImageDataMessage {
   pub model_name: Option<String>,
 }
 
-#[derive(Deserialize, Clone)]
+#[derive(Deserialize, Clone, Debug)]
 pub struct JsonDataMessage {
   /// UUID.
   pub job_id: Option<String>,
-  
+
   /// UUID.
   pub request_id: Option<String>,
 
@@ -72,6 +72,12 @@ pub struct JsonDataMessage {
   pub image_id: Option<String>,
 }
 
+impl WebsocketServerMessage {
+  pub fn from_json_str(json_str: &str) -> Result<Self, GrokGenericApiError> {
+    Ok(serde_json::from_str(json_str)
+        .map_err(|err| GrokGenericApiError::SerdeResponseParseErrorWithBody(err, json_str.to_string()))?)
+  }
+}
 
 #[cfg(test)]
 mod tests {

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

@@ -25,6 +25,7 @@ mod tests {
   use errors::AnyhowResult;
   use std::io::Write;
   use std::time::Duration;
+  use crate::requests::listen_websocket::messages::websocket_server_message::WebsocketServerMessage;
 
   #[tokio::test]
   #[ignore] // manually test
@@ -34,7 +35,7 @@ mod tests {
 
     let cookies = get_test_cookies()?;
 
-    let mut websocket = create_listen_websocket(CreateListenWebsocketArgs {
+    let websocket = create_listen_websocket(CreateListenWebsocketArgs {
       cookies: &cookies,
     }).await?;
 
@@ -75,9 +76,28 @@ mod tests {
     //  }
     //}
 
+    //loop {
+    //  let maybe_message =
+    //      websocket.try_next_timeout(Duration::from_millis(1000)).await?;
+    //  match maybe_message {
+    //    None => {
+    //      println!("No message received within timeout.");
+    //      count = count + 1;
+    //    }
+    //    Some(message) => {
+    //      println!("[{count}] Received websocket message!");//: {:?}", message);
+    //      count = 0;
+    //    }
+    //  }
+    //  if count > 5 {
+    //    println!("No messages after 5 seconds");
+    //    break;
+    //  }
+    //}
+
     loop {
       let maybe_message =
-          websocket.try_next_timeout(Duration::from_millis(1000)).await?;
+          websocket.get_response_with_timeout(Duration::from_millis(1000)).await?;
 
       match maybe_message {
         None => {
@@ -85,12 +105,28 @@ mod tests {
           count = count + 1;
         }
         Some(message) => {
-          println!("[{count}] Received websocket message!");//: {:?}", message);
+          match message {
+            WebsocketServerMessage::ImageData(image) => {
+              println!("IMAGE: {:?}", image.percentage_complete);
+            }
+            WebsocketServerMessage::JsonData(json) => {
+              println!("JSON : {:?}", json.percentage_complete);
+            }
+            WebsocketServerMessage::Unknown(unknown) => {
+              println!("[UNKNOWN] websocket message: {:?}", unknown);
+            }
+          }
           count = 0;
         }
       }
 
-      if count > 5 {
+      // 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 > 30 {
         println!("No messages after 5 seconds");
         break;
       }