From e956c32652d46a7ab5d71d6465c941fb866fddbb Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 3 Apr 2025 17:06:00 -0400 Subject: [PATCH] read tokens from headers, cleanup --- .../openai_sora_client/src/credentials.rs | 12 +------ .../src/upload/upload_media_from_file.rs | 2 +- .../src/upload/upload_media_http_request.rs | 10 ++++-- ...enqueue_studio_image_generation_handler.rs | 31 +++++++++++++++---- 4 files changed, 34 insertions(+), 21 deletions(-) diff --git a/crates/lib/api_clients/openai_sora_client/src/credentials.rs b/crates/lib/api_clients/openai_sora_client/src/credentials.rs index 15a9da377b..00e4ff6862 100644 --- a/crates/lib/api_clients/openai_sora_client/src/credentials.rs +++ b/crates/lib/api_clients/openai_sora_client/src/credentials.rs @@ -1,7 +1,7 @@ use reqwest::RequestBuilder; use errors::AnyhowResult; // NB: It appears that the sentinel may require a matching user agent. -const USER_AGENT : &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36"; +pub (crate) const USER_AGENT : &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36"; #[derive(Clone)] pub struct SoraCredentials { @@ -29,10 +29,6 @@ impl SoraCredentials { pub fn add_credential_headers_to_request(&self, request: RequestBuilder) -> RequestBuilder { let bearer_header = self.authorization_header_value(); - println!(">>> BEARER HEADER = {}", bearer_header); - //println!(">>> COOKIE HEADER = {}", self.cookie); - println!(">>> SENTINEL HEADER = {}", self.sentinel.clone().unwrap_or("None".to_string())); - let mut request = request .header("User-Agent", USER_AGENT) .header("Cookie", &self.cookie) @@ -42,12 +38,6 @@ impl SoraCredentials { request = request.header("OpenAI-Sentinel-Token", sentinel); } - /* - Without sentinel token: - status: "{\n \"error\": {\n \"message\": \"Hmmm something didn't look right with your request. Please try again later or visit https://help.openai.com if this issue persists.\",\n \"type\": \"invalid_request_error\",\n \"param\": null,\n \"code\": \"sentinel_block\"\n }\n}" - Error: missing field `id` at line 8 column 1 - - */ request } } diff --git a/crates/lib/api_clients/openai_sora_client/src/upload/upload_media_from_file.rs b/crates/lib/api_clients/openai_sora_client/src/upload/upload_media_from_file.rs index cbedc2e4da..d7f00d1556 100644 --- a/crates/lib/api_clients/openai_sora_client/src/upload/upload_media_from_file.rs +++ b/crates/lib/api_clients/openai_sora_client/src/upload/upload_media_from_file.rs @@ -56,7 +56,7 @@ mod tests { let bearer = read_to_string(test_file_path("test_data/temp/bearer.txt")?)?; let bearer = bearer.trim().to_string(); - let image_path = test_file_path("test_data/image/mochi.jpg")?; + let image_path = test_file_path("test_data/image/juno.jpg")?; // media_01jqyqgqpwf40tkcapq5bmaz5d let creds = SoraCredentials { bearer_token: bearer, diff --git a/crates/lib/api_clients/openai_sora_client/src/upload/upload_media_http_request.rs b/crates/lib/api_clients/openai_sora_client/src/upload/upload_media_http_request.rs index b44ee45b0f..c4774322db 100644 --- a/crates/lib/api_clients/openai_sora_client/src/upload/upload_media_http_request.rs +++ b/crates/lib/api_clients/openai_sora_client/src/upload/upload_media_http_request.rs @@ -1,8 +1,8 @@ +use crate::credentials::{SoraCredentials, USER_AGENT}; use errors::AnyhowResult; use reqwest::multipart::{Form, Part}; use reqwest::Client; use serde::Deserialize; -use crate::credentials::SoraCredentials; const SORA_UPLOAD_MEDIA_URL: &str = "https://sora.com/backend/uploads"; @@ -60,8 +60,12 @@ pub (crate) async fn upload_media_http_request(file_bytes: Vec, filename: St // Make API request let client = Client::new(); - let request_builder = client.post(SORA_UPLOAD_MEDIA_URL).multipart(form); - let request_builder = credentials.add_credential_headers_to_request(request_builder); + let request_builder = client.post(SORA_UPLOAD_MEDIA_URL) + .multipart(form) + .header("User-Agent", USER_AGENT) + .header("Cookie", &credentials.cookie) + .header("Authorization", credentials.authorization_header_value()); + let response = request_builder.send().await?; // Check response status diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_studio/prompt/enqueue_studio_image_generation_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_studio/prompt/enqueue_studio_image_generation_handler.rs index 68ba303e5e..c720101481 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/image_studio/prompt/enqueue_studio_image_generation_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_studio/prompt/enqueue_studio_image_generation_handler.rs @@ -68,6 +68,9 @@ const DEBUG_HEADER_NAME: &str = "enable-debug-mode"; /// This is useful for catching the live logs or intercepting the job. const ROUTING_TAG_HEADER_NAME: &str = "routing-tag"; +const SORA_BEARER_HEADER_NAME : &str = "sora-bearer"; +const SORA_COOKIE_HEADER_NAME : &str = "sora-cookie"; +const SORA_SENTINEL_HEADER_NAME : &str = "sora-sentinel"; #[derive(Deserialize, ToSchema)] pub struct EnqueueStudioImageGenRequest { @@ -281,14 +284,30 @@ pub async fn enqueue_studio_image_generation_handler( // ==================== HANDLE SORA CREDENTIALS ==================== // + let sora_bearer = get_request_header_optional(&http_request, SORA_BEARER_HEADER_NAME); + let sora_cookie = get_request_header_optional(&http_request, SORA_COOKIE_HEADER_NAME); + let sora_sentinel = get_request_header_optional(&http_request, SORA_SENTINEL_HEADER_NAME); - let sora_credentials = match SoraCredentials::from_env() { - Ok(creds) => creds, - Err(err) => { - error!("Failed to load Sora credentials from environment: {:?}", err); - return Err(EnqueueImageGenRequestError::ServerError); + let sora_credentials; + + match (sora_bearer, sora_cookie, sora_sentinel) { + (Some(bearer), Some(cookie), Some(sentinel)) => { + sora_credentials = SoraCredentials { + bearer_token: bearer, + cookie, + sentinel: Some(sentinel), + }; } - }; + _ => { + sora_credentials = match SoraCredentials::from_env() { + Ok(creds) => creds, + Err(err) => { + error!("Failed to load Sora credentials from environment: {:?}", err); + return Err(EnqueueImageGenRequestError::ServerError); + } + }; + } + } // ==================== HANDLE SORA UPLOAD ==================== //