diff --git a/crates/api_clients/grok_client/src/requests/listen_websocket/clonable_websocket.rs b/crates/api_clients/grok_client/src/requests/listen_websocket/clonable_websocket.rs index 56c1798384..d20ed4fa5d 100644 --- a/crates/api_clients/grok_client/src/requests/listen_websocket/clonable_websocket.rs +++ b/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, 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, GrokClientError> { match self.websocket.write() { Err(_) => Err(GrokClientError::WebsocketLockError), diff --git a/crates/api_clients/grok_client/src/requests/listen_websocket/messages/mod.rs b/crates/api_clients/grok_client/src/requests/listen_websocket/messages/mod.rs index 88fb5fee0e..055b1b8f2a 100644 --- a/crates/api_clients/grok_client/src/requests/listen_websocket/messages/mod.rs +++ b/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; \ No newline at end of file +pub mod websocket_server_message; +pub mod image_data_message; \ No newline at end of file diff --git a/crates/api_clients/grok_client/src/requests/listen_websocket/messages/websocket_server_message.rs b/crates/api_clients/grok_client/src/requests/listen_websocket/messages/websocket_server_message.rs index 02f72a95c8..bb5289431d 100644 --- a/crates/api_clients/grok_client/src/requests/listen_websocket/messages/websocket_server_message.rs +++ b/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, @@ -50,11 +50,11 @@ pub struct ImageDataMessage { pub model_name: Option, } -#[derive(Deserialize, Clone)] +#[derive(Deserialize, Clone, Debug)] pub struct JsonDataMessage { /// UUID. pub job_id: Option, - + /// UUID. pub request_id: Option, @@ -72,6 +72,12 @@ pub struct JsonDataMessage { pub image_id: Option, } +impl WebsocketServerMessage { + pub fn from_json_str(json_str: &str) -> Result { + Ok(serde_json::from_str(json_str) + .map_err(|err| GrokGenericApiError::SerdeResponseParseErrorWithBody(err, json_str.to_string()))?) + } +} #[cfg(test)] mod tests { diff --git a/crates/api_clients/grok_client/src/requests/listen_websocket/prompt_websocket_image.rs b/crates/api_clients/grok_client/src/requests/listen_websocket/prompt_websocket_image.rs index b16a092d0e..4dd655bfec 100644 --- a/crates/api_clients/grok_client/src/requests/listen_websocket/prompt_websocket_image.rs +++ b/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; }