From 421ef353a2d1b3edd23f453b2a3b5ea8d21db585 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 23 Apr 2025 23:24:25 -0400 Subject: [PATCH] continued refactor --- .../src/recipes/wait_for_image_gen_status.rs | 2 +- .../src/requests/image_gen/mod.rs | 1 - .../src/requests/job/sora_job_status.rs | 43 +++++++++++++------ .../v2/process_single_sora_job.rs | 3 +- 4 files changed, 32 insertions(+), 17 deletions(-) diff --git a/crates/lib/api_clients/openai_sora_client/src/recipes/wait_for_image_gen_status.rs b/crates/lib/api_clients/openai_sora_client/src/recipes/wait_for_image_gen_status.rs index dbf88733dc..79a9a71d89 100644 --- a/crates/lib/api_clients/openai_sora_client/src/recipes/wait_for_image_gen_status.rs +++ b/crates/lib/api_clients/openai_sora_client/src/recipes/wait_for_image_gen_status.rs @@ -1,6 +1,6 @@ use crate::credentials::SoraCredentials; use crate::requests::image_gen::image_gen_status::TaskStatus; -use crate::requests::image_gen::{get_image_gen_status, StatusRequest, TaskResponse, VideoGenStatusResponse}; +use crate::requests::image_gen::image_gen_status::{get_image_gen_status, StatusRequest, TaskResponse, VideoGenStatusResponse}; use errors::AnyhowResult; pub async fn wait_for_image_gen_status(task_id: &String, credentials: &SoraCredentials, retry_limit: Option) -> AnyhowResult { diff --git a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/mod.rs b/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/mod.rs index fb9c54de2f..f0247756fb 100644 --- a/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/mod.rs +++ b/crates/lib/api_clients/openai_sora_client/src/requests/image_gen/mod.rs @@ -4,4 +4,3 @@ pub mod common; pub mod sora_image_gen_remix; pub mod sora_image_gen_simple; pub mod image_gen_status; -pub use image_gen_status::{get_image_gen_status, save_generations_to_dir, Generation, StatusRequest, TaskResponse, VideoGenStatusResponse}; diff --git a/crates/lib/api_clients/openai_sora_client/src/requests/job/sora_job_status.rs b/crates/lib/api_clients/openai_sora_client/src/requests/job/sora_job_status.rs index c33f9541f9..15e95ca335 100644 --- a/crates/lib/api_clients/openai_sora_client/src/requests/job/sora_job_status.rs +++ b/crates/lib/api_clients/openai_sora_client/src/requests/job/sora_job_status.rs @@ -1,29 +1,45 @@ +use crate::creds::credential_migration::CredentialMigrationRef; +use crate::sora_error::SoraError; + 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"; // https://sora.com/backend/notif?limit=100&before=task_01jqwwrkvgeqp8jsf5mqk1jceh -use crate::credentials::SoraCredentials; - pub async fn sora_job_status( task_id: &str, - credentials: &SoraCredentials, -) -> anyhow::Result<()> { + credentials: CredentialMigrationRef<'_>, +) -> Result<(), SoraError> { let url = format!("https://sora.com/backend/notif?limit=100&before={}", task_id); + let mut cookie; + let mut authorization_header; + + match credentials { + CredentialMigrationRef::Legacy(creds) => { + cookie = creds.cookie.clone(); + authorization_header = creds.authorization_header_value(); + } + CredentialMigrationRef::New(creds) => { + cookie = creds.cookies.to_string(); + // TODO(bt,2025-04-23): We're using a Sora payload error in place of application state error. Surface this differently. + authorization_header = creds.jwt_bearer_token.as_ref() + .ok_or(SoraError::NoBearerTokenAvailable)? + .to_authorization_header_value(); + } + } + let client = reqwest::Client::new(); let request = client .get(url) - .header("Authorization", &credentials.bearer_token) + .header("Authorization", &authorization_header) .header("User-Agent", USER_AGENT) - .header("Cookie", &credentials.cookie); + .header("Cookie", &cookie); let response = request.send().await?; let json_response = &response.text().await?; - println!(" >>> response = {:?}", json_response); - //let response = serde_json::from_str(json_response)?; Ok(()) @@ -31,11 +47,12 @@ pub async fn sora_job_status( #[cfg(test)] mod tests { - use std::fs::read_to_string; - use errors::AnyhowResult; - use testing::test_file_path::test_file_path; use crate::credentials::SoraCredentials; - use crate::job::sora_job_status::sora_job_status; + use crate::creds::credential_migration::CredentialMigrationRef; + use crate::requests::job::sora_job_status::sora_job_status; + use errors::AnyhowResult; + use std::fs::read_to_string; + use testing::test_file_path::test_file_path; #[ignore] #[tokio::test] @@ -57,7 +74,7 @@ mod tests { let task_id = "task_01jqwwrkvgeqp8jsf5mqk1jceh"; - let result = sora_job_status(task_id, &creds).await?; + let result = sora_job_status(task_id, CredentialMigrationRef::Legacy(&creds)).await?; Ok(()) } diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/v2/process_single_sora_job.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/v2/process_single_sora_job.rs index 9e348e4c13..79667ff46f 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/v2/process_single_sora_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/v2/process_single_sora_job.rs @@ -29,8 +29,7 @@ use mysql_queries::queries::batch_generations::insert_batch_generation_records:: use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; use openai_sora_client::credentials::SoraCredentials; use openai_sora_client::recipes::wait_for_image_gen_status::wait_for_image_gen_status; -use openai_sora_client::requests::image_gen::image_gen_status::TaskStatus; -use openai_sora_client::requests::image_gen::Generation; +use openai_sora_client::requests::image_gen::image_gen_status::{Generation, TaskStatus}; use r2d2_redis::r2d2::PooledConnection; use r2d2_redis::redis::Commands; use r2d2_redis::RedisConnectionManager;