diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/clonable_websocket.rs b/crates/api_clients/grok_client/src/requests/image_websocket/grok_websocket.rs similarity index 98% rename from crates/api_clients/grok_client/src/requests/image_websocket/clonable_websocket.rs rename to crates/api_clients/grok_client/src/requests/image_websocket/grok_websocket.rs index d9e5cf2b11..e607b1678c 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/clonable_websocket.rs +++ b/crates/api_clients/grok_client/src/requests/image_websocket/grok_websocket.rs @@ -11,11 +11,11 @@ use wreq::ws::message::Message; use wreq::ws::WebSocket; #[derive(Clone)] -pub struct ClonableWebsocket { +pub struct GrokWebsocket { pub(crate) websocket: Arc>, } -impl ClonableWebsocket { +impl GrokWebsocket { pub fn new(websocket: WebSocket) -> Self { Self { websocket: Arc::new(RwLock::new(websocket)), @@ -77,7 +77,6 @@ impl ClonableWebsocket { } } - 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/image_websocket/listen_for_websocket_images.rs b/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_images.rs index 5354b322e9..6f4211a621 100644 --- 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 @@ -1,7 +1,7 @@ use crate::datatypes::api::request_id::RequestId; 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::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::messages::message_image_data::MessageImageData; use crate::utils::scrub_blobs_for_debug_logging::scrub_blobs_for_debug_logging; use log::info; @@ -12,7 +12,7 @@ use wreq::ws::message::Message; const DEFAULT_IMAGE_COUNT: usize = 6; pub struct ListenForWebsocketImagesArgs<'a> { - pub websocket: &'a ClonableWebsocket, + pub websocket: &'a GrokWebsocket, pub timeout: Duration, } @@ -113,7 +113,7 @@ pub async fn listen_for_websocket_images(args: ListenForWebsocketImagesArgs<'_>) #[cfg(test)] mod tests { - use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket; + use crate::requests::image_websocket::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; use crate::requests::image_websocket::listen_for_websocket_images::{listen_for_websocket_images, ListenForWebsocketImagesArgs}; use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; @@ -137,7 +137,7 @@ mod tests { cookies: &cookies, }).await?; - let websocket = ClonableWebsocket::new(websocket); + let websocket = GrokWebsocket::new(websocket); println!("Sending..."); std::io::stdout().flush()?; diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs b/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs index d5c15b083a..84c313783d 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs +++ b/crates/api_clients/grok_client/src/requests/image_websocket/listen_for_websocket_request_id.rs @@ -1,13 +1,13 @@ use crate::datatypes::api::request_id::RequestId; 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::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::messages::message_request_id::MessageRequestId; use std::time::{Duration, Instant}; use wreq::ws::message::Message; pub struct ListenForWebsocketRequestIdArgs<'a> { - pub websocket: &'a ClonableWebsocket, + pub websocket: &'a GrokWebsocket, pub timeout: Duration, } @@ -70,7 +70,7 @@ pub async fn listen_for_websocket_request_id(args: ListenForWebsocketRequestIdAr #[cfg(test)] mod tests { - use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket; + use crate::requests::image_websocket::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; use crate::requests::image_websocket::listen_for_websocket_request_id::{listen_for_websocket_request_id, ListenForWebsocketRequestIdArgs}; use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; @@ -94,7 +94,7 @@ mod tests { cookies: &cookies, }).await?; - let websocket = ClonableWebsocket::new(websocket); + let websocket = GrokWebsocket::new(websocket); println!("Sending..."); std::io::stdout().flush()?; 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 23b5af662c..b15666c033 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,4 @@ -pub mod clonable_websocket; +pub mod grok_websocket; pub mod create_listen_websocket; pub mod listen_for_websocket_images; pub mod listen_for_websocket_request_id; 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 113a061225..1f2d4859d8 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 @@ -1,9 +1,9 @@ use crate::error::grok_error::GrokError; -use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket; +use crate::requests::image_websocket::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::messages::websocket_client_message::WebsocketClientMessage; pub struct PromptWebsocketImageArgs<'a> { - pub websocket_wrapped: ClonableWebsocket, + pub websocket_wrapped: GrokWebsocket, pub prompt: &'a str, } @@ -18,7 +18,7 @@ pub async fn prompt_websocket_image(args: PromptWebsocketImageArgs<'_>) -> Resul #[cfg(test)] mod tests { - use crate::requests::image_websocket::clonable_websocket::ClonableWebsocket; + use crate::requests::image_websocket::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; use crate::requests::image_websocket::messages::websocket_server_message::WebsocketServerMessage; use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; @@ -40,7 +40,7 @@ mod tests { cookies: &cookies, }).await?; - let websocket = ClonableWebsocket::new(websocket); + let websocket = GrokWebsocket::new(websocket); println!("Sending..."); std::io::stdout().flush()?; diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_websocket_manager.rs b/crates/desktop/artcraft/src/services/grok/state/grok_websocket_manager.rs new file mode 100644 index 0000000000..d912d9495c --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/state/grok_websocket_manager.rs @@ -0,0 +1,7 @@ +use grok_client::requests::image_websocket::grok_websocket::GrokWebsocket; +use std::sync::{Arc, RwLock}; + +#[derive(Clone)] +pub struct GrokWebsocketHolder { + websocket: Arc>, +} \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/grok/state/mod.rs b/crates/desktop/artcraft/src/services/grok/state/mod.rs index 9ea2b969bd..d2c195dc4d 100644 --- a/crates/desktop/artcraft/src/services/grok/state/mod.rs +++ b/crates/desktop/artcraft/src/services/grok/state/mod.rs @@ -1,3 +1,4 @@ pub mod grok_credential_manager; pub mod grok_serializable_state; -mod grok_credential_holder; \ No newline at end of file +mod grok_credential_holder; +mod grok_websocket_manager; \ No newline at end of file