diff --git a/Cargo.lock b/Cargo.lock index ba2126b324..129667152f 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index 59f27455df..065c2a88c7 100644 --- a/Cargo.toml +++ b/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. # diff --git a/crates/api_clients/grok_client/Cargo.toml b/crates/api_clients/grok_client/Cargo.toml index 5a4eb481f9..320ffe9533 100644 --- a/crates/api_clients/grok_client/Cargo.toml +++ b/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 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 2b05d6aaa3..6e9015e969 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 @@ -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."), 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 new file mode 100644 index 0000000000..635e41f6da --- /dev/null +++ b/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>, +} + +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(&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 { + match self.websocket.read() { + Err(_) => Err(GrokClientError::WebsocketLockError.into()), + Ok(websocket) => Ok(websocket.is_terminated()), + } + } + + pub async fn try_next(&self) -> Result, GrokClientError> { + match self.websocket.write() { + Err(_) => Err(GrokClientError::WebsocketLockError), + Ok(mut websocket) => websocket + .try_next() + .await + .map_err(|err| GrokClientError::WebsocketReadError(err)), + } + } +} diff --git a/crates/api_clients/grok_client/src/requests/listen_websocket/messages/metadata/create_image_request.json b/crates/api_clients/grok_client/src/requests/listen_websocket/messages/metadata/create_image_request.json new file mode 100644 index 0000000000..5ea498f971 --- /dev/null +++ b/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 + } + } + ] + } +} 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 new file mode 100644 index 0000000000..432675fd47 --- /dev/null +++ b/crates/api_clients/grok_client/src/requests/listen_websocket/messages/mod.rs @@ -0,0 +1 @@ +pub mod websocket_request; \ No newline at end of file diff --git a/crates/api_clients/grok_client/src/requests/listen_websocket/messages/websocket_request.rs b/crates/api_clients/grok_client/src/requests/listen_websocket/messages/websocket_request.rs new file mode 100644 index 0000000000..2bfe52954f --- /dev/null +++ b/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, +} + +#[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, + }, + } + ], + }, + } + } +} diff --git a/crates/api_clients/grok_client/src/requests/listen_websocket/mod.rs b/crates/api_clients/grok_client/src/requests/listen_websocket/mod.rs index 604c51aa88..b9d8942f4c 100644 --- a/crates/api_clients/grok_client/src/requests/listen_websocket/mod.rs +++ b/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; \ No newline at end of file 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 new file mode 100644 index 0000000000..ed7d03ab06 --- /dev/null +++ b/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(()) + } +} \ No newline at end of file