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 360107db13..4a192cb224 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 @@ -117,6 +117,7 @@ mod tests { use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; use crate::requests::image_websocket::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::listen_for_websocket_images::{listen_for_websocket_images, ListenForWebsocketImagesArgs}; + use crate::requests::image_websocket::messages::websocket_client_message::ClientMessageAspectRatio; use crate::requests::image_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; @@ -147,6 +148,7 @@ mod tests { let _result = prompt_websocket_image(PromptWebsocketImageArgs { websocket: &mut websocket, prompt, + aspect_ratio: ClientMessageAspectRatio::WideThreeByTwo, }).await?; println!("Reading..."); 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 a0d0a38999..3a35791dd3 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 @@ -73,6 +73,7 @@ mod tests { use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; use crate::requests::image_websocket::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::listen_for_websocket_request_id::{listen_for_websocket_request_id, ListenForWebsocketRequestIdArgs}; + use crate::requests::image_websocket::messages::websocket_client_message::ClientMessageAspectRatio; use crate::requests::image_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; @@ -103,6 +104,7 @@ mod tests { let _result = prompt_websocket_image(PromptWebsocketImageArgs { websocket: &mut websocket, prompt, + aspect_ratio: ClientMessageAspectRatio::Square, }).await?; println!("Reading..."); diff --git a/crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_client_message.rs b/crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_client_message.rs index 5b3400a0c3..88941f69e9 100644 --- a/crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_client_message.rs +++ b/crates/api_clients/grok_client/src/requests/image_websocket/messages/websocket_client_message.rs @@ -46,11 +46,55 @@ pub struct ClientMessageItemContentProperties { pub enable_nsfw: bool, pub skip_upsampler: bool, pub is_initial: bool, + pub aspect_ratio: ClientMessageAspectRatio, +} + +#[derive(Serialize, Clone, Copy)] +pub enum ClientMessageAspectRatio { + #[serde(rename = "2:3")] + TallTwoByThree, + + #[serde(rename = "3:2")] + WideThreeByTwo, + + #[serde(rename = "1:1")] + Square, } impl WebsocketClientMessage { /// Create a new image prompt websocket client message - pub fn new_image_prompt(prompt: &str) -> Self { + pub fn new_image_prompt(prompt: &str, aspect_ratio: ClientMessageAspectRatio) -> Self { + /* + { + "type":"conversation.item.create", + "timestamp":1764222761003, + "item": { + "type":"message", + "content": [ + { + "requestId":"efdde1fc-2427-4f49-bd90-3e54df0e6294", + "text":"fighter jet in the taco bell drive thru", + "type":"input_text", + "properties": { + "section_count":0, + "is_kids_mode":false, + "enable_nsfw":true, + "skip_upsampler":false, + "is_initial":false, + "aspect_ratio":"2:3" + } + } + ] + } + } + + "2:3" + "3:2" + "1:1" + + {"type":"conversation.item.create","timestamp":1764223129572,"item":{"type":"message","content":[{"requestId":"3825ffdd-0613-48ba-9a50-fb75bbfbdba0","text":"tank at mcdonalds","type":"input_text","properties":{"section_count":0,"is_kids_mode":false,"enable_nsfw":true,"skip_upsampler":false,"is_initial":false,"aspect_ratio":"1:1"}}]}} + {"type":"conversation.item.create","timestamp":1764223055299,"item":{"type":"message","content":[{"requestId":"86361ecf-bdcf-4159-bfcb-f1097dc7499d","text":"an attack helicopter at Burger King","type":"input_scroll","properties":{"section_count":0,"is_kids_mode":false,"enable_nsfw":true,"skip_upsampler":false,"is_initial":false,"aspect_ratio":"3:2"}}]}} + */ Self { r#type: "conversation.item.create".to_string(), timestamp: Utc::now().timestamp_millis() as u64, @@ -67,6 +111,7 @@ impl WebsocketClientMessage { enable_nsfw: true, skip_upsampler: false, is_initial: false, + aspect_ratio, }, } ], 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 7b223f5e5e..10b9c859d9 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,15 +1,16 @@ use crate::error::grok_error::GrokError; use crate::requests::image_websocket::grok_websocket::GrokWebsocket; -use crate::requests::image_websocket::messages::websocket_client_message::WebsocketClientMessage; +use crate::requests::image_websocket::messages::websocket_client_message::{ClientMessageAspectRatio, WebsocketClientMessage}; pub struct PromptWebsocketImageArgs<'a> { pub websocket: &'a mut GrokWebsocket, pub prompt: &'a str, + pub aspect_ratio: ClientMessageAspectRatio, } pub async fn prompt_websocket_image(args: PromptWebsocketImageArgs<'_>) -> Result<(), GrokError> { - let message = WebsocketClientMessage::new_image_prompt(args.prompt); + let message = WebsocketClientMessage::new_image_prompt(args.prompt, args.aspect_ratio); args.websocket.send_serializable(&message).await?; @@ -21,6 +22,7 @@ mod tests { use crate::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; use crate::requests::image_websocket::grok_websocket::GrokWebsocket; use crate::requests::image_websocket::grok_wrapped_websocket::GrokWrappedWebsocket; + use crate::requests::image_websocket::messages::websocket_client_message::ClientMessageAspectRatio; use crate::requests::image_websocket::messages::websocket_server_message::WebsocketServerMessage; use crate::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; use crate::test_utils::get_test_cookies::get_test_cookies; @@ -50,6 +52,7 @@ mod tests { let result = prompt_websocket_image(PromptWebsocketImageArgs { websocket: &mut websocket, prompt: "a dog riding a motorcycle", + aspect_ratio: ClientMessageAspectRatio::WideThreeByTwo, }).await?; println!("Reading...");