updating to handle sora 2.0 task polling

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_failed_generations.rs

@@ -5,6 +5,7 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::core::state::task_database::TaskDatabase;
 use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
+use crate::services::sora::threads::sora_task_polling::helpers::handle_successful_generations::GenerationItem;
 use enums::tauri::tasks::task_status;
 use errors::AnyhowResult;
 use log::info;
@@ -16,11 +17,15 @@ use std::collections::HashMap;
 use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
 use tauri::AppHandle;
 
-pub async fn handle_failed_generations(
+pub struct FailedGeneration {
+  pub reason: Option<String>,
+}
+
+pub async fn handle_classic_failed_generations(
   app_handle: &AppHandle,
   task_database: &TaskDatabase,
   local_sqlite_tasks_by_sora_task_id: &HashMap<String, Task>,
-  sora_failed_tasks_by_id: &HashMap<TaskId, PartialTaskResponse>,
+  sora_failed_tasks_by_id: &HashMap<TaskId, FailedGeneration>,
   sora_task_queue: &SoraTaskQueue,
 ) -> AnyhowResult<()> {
 
@@ -31,7 +36,7 @@ pub async fn handle_failed_generations(
         action: GenerationAction::GenerateImage,
         service: GenerationServiceProvider::Sora,
         model: None,
-        reason: None,
+        reason: task.reason.clone(),
       };
 
       event.send_infallible(&app_handle);

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs

@@ -2,7 +2,6 @@ use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
 use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
 use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
 use crate::core::events::generation_events::generation_failed_event::GenerationFailedEvent;
-use crate::core::events::sendable_event_trait::SendableEvent;
 use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
@@ -36,13 +35,23 @@ use storyteller_client::endpoints::prompts::create_prompt::create_prompt;
 use tauri::AppHandle;
 use tempdir::TempDir;
 
-pub async fn handle_successful_generations(
+pub struct SuccessfulGeneration {
+  pub prompt: Option<String>,
+  pub items: Vec<GenerationItem>,
+}
+
+pub struct GenerationItem {
+  pub item_id: String,
+  pub url: String,
+}
+
+pub async fn handle_classic_successful_generations(
   app_handle: &AppHandle,
   app_data_root: &AppDataRoot,
   app_env_configs: &AppEnvConfigs,
   task_database: &TaskDatabase,
   storyteller_creds: &StorytellerCredentialSet,
-  succeeded_tasks_by_id: &HashMap<TaskId, PartialTaskResponse>,
+  succeeded_tasks_by_id: &HashMap<TaskId, SuccessfulGeneration>,
   sqlite_tasks_by_sora_task_id: &HashMap<String, Task>,
 ) -> AnyhowResult<()> {
 
@@ -51,7 +60,7 @@ pub async fn handle_successful_generations(
       continue; // Task is irrelevant - previously completed, generated elsewhere, etc.
     }
 
-    info!("Task succeeded: {:?}", task.id);
+    info!("Task succeeded: {:?}", task_id);
 
     let request = CreatePromptRequest {
       uuid_idempotency_token: generate_random_uuid(),
@@ -69,7 +78,7 @@ pub async fn handle_successful_generations(
 
     info!("Created prompt: {:?}", &prompt_response.prompt_token);
 
-    for (i, generation) in task.generations.iter().enumerate() {
+    for (_i, generation) in task.items.iter().enumerate() {
       info!("Downloading generated file...");
       let download_path = download_generation(generation, &app_data_root).await?;
 
@@ -115,7 +124,7 @@ pub async fn handle_successful_generations(
 }
 
 
-async fn download_generation(generation: &PartialGeneration, app_data_root: &AppDataRoot) -> AnyhowResult<PathBuf> {
+async fn download_generation(generation: &GenerationItem, app_data_root: &AppDataRoot) -> AnyhowResult<PathBuf> {
   let url = Url::parse(&generation.url)?;
 
   let response = reqwest::get(&generation.url).await?;
@@ -124,7 +133,7 @@ async fn download_generation(generation: &PartialGeneration, app_data_root: &App
   let ext = url.path().split(".").last().unwrap_or("png");
 
   let tempdir = app_data_root.temp_dir().path();
-  let download_filename = format!("{}.{}", generation.id, ext);
+  let download_filename = format!("{}.{}", generation.item_id, ext);
   let download_path = tempdir.join(download_filename);
 
   let mut file = File::create(&download_path)?;

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/mod.rs

@@ -0,0 +1,4 @@
+pub (super) mod handle_successful_generations;
+pub (super) mod handle_failed_generations;
+pub (super) mod poll_classic_sora_tasks;
+pub (super) mod poll_sora_2_tasks;

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_classic_sora_tasks.rs

@@ -0,0 +1,122 @@
+use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
+use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
+use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
+use crate::core::events::generation_events::generation_failed_event::GenerationFailedEvent;
+use crate::core::events::sendable_event_trait::SendableEvent;
+use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
+use crate::core::state::data_dir::app_data_root::AppDataRoot;
+use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
+use crate::core::state::task_database::TaskDatabase;
+use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES;
+use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
+use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
+use crate::services::sora::threads::sora_task_polling::helpers::handle_failed_generations::{handle_classic_failed_generations, FailedGeneration};
+use crate::services::sora::threads::sora_task_polling::helpers::handle_successful_generations::{handle_classic_successful_generations, GenerationItem, SuccessfulGeneration};
+use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
+use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
+use enums::common::generation_provider::GenerationProvider;
+use enums::common::model_type::ModelType;
+use enums::tauri::tasks::task_status;
+use errors::AnyhowResult;
+use idempotency::uuid::generate_random_uuid;
+use log::{error, info, warn};
+use once_cell::sync::Lazy;
+use openai_sora_client::creds::sora_credential_set::SoraCredentialSet;
+use openai_sora_client::recipes::list_classic_sora_tasks_with_session_auto_renew::list_classic_sora_tasks_with_session_auto_renew;
+use openai_sora_client::requests::common::task_id::TaskId;
+use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::TaskStatus;
+use reqwest::Url;
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList};
+use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
+use std::collections::{HashMap, HashSet};
+use std::fs::File;
+use std::io::Write;
+use std::path::PathBuf;
+use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs};
+use storyteller_client::endpoints::prompts::create_prompt::create_prompt;
+use tauri::AppHandle;
+use tempdir::TempDir;
+
+pub async fn poll_classic_sora_tasks(
+  app_handle: &AppHandle,
+  app_env_configs: &AppEnvConfigs,
+  task_database: &TaskDatabase,
+  sora_creds_manager: &SoraCredentialManager,
+  sora_creds: &SoraCredentialSet,
+  storyteller_creds_manager: &StorytellerCredentialManager,
+  sora_task_queue: &SoraTaskQueue,
+  app_data_root: &AppDataRoot,
+  local_sqlite_tasks_by_sora_task_id: &HashMap<String, Task>,
+) -> AnyhowResult<()> {
+
+  let (sora_response, maybe_new_creds) =
+      list_classic_sora_tasks_with_session_auto_renew(&sora_creds).await?;
+
+  if let Some(new_creds) = maybe_new_creds {
+    info!("Saving new credentials.");
+    sora_creds_manager.set_credentials(&new_creds)?;
+  }
+
+  let sora_items = sora_response.task_responses;
+
+  let mut sora_succeeded_tasks_by_id = HashMap::new();
+  let mut sora_failed_tasks_by_id = HashMap::new();
+
+  let storyteller_creds = storyteller_creds_manager.get_credentials_required()?;
+
+  for task in sora_items.iter() {
+
+    match &task.status {
+      TaskStatus::Succeeded => {
+        sora_succeeded_tasks_by_id.insert(
+          task.id.clone(), 
+          SuccessfulGeneration {
+            prompt: task.prompt.clone(),
+            items: task.generations.iter()
+                .map(|gen| {
+                  GenerationItem {
+                    item_id: gen.id.clone(),
+                    url: gen.url.clone(),
+                  }
+                })
+                .collect()
+          });
+      }
+      TaskStatus::Failed => {
+        sora_failed_tasks_by_id.insert(
+          task.id.clone(), 
+          FailedGeneration {
+            reason: None, // TODO: Add reason if available.
+          }
+        );
+      }
+      TaskStatus::Queued => {}
+      TaskStatus::Running => {}
+      TaskStatus::Unknown(unknown_status) => {
+        warn!("Unknown task status: {:?}", unknown_status);
+      }
+    }
+  }
+
+  // Clear dead tasks.
+  handle_classic_failed_generations(
+    &app_handle,
+    &task_database,
+    &local_sqlite_tasks_by_sora_task_id,
+    &sora_failed_tasks_by_id,
+    &sora_task_queue,
+  ).await?;
+
+  // Process succeeded tasks.
+  handle_classic_successful_generations(
+    &app_handle,
+    &app_data_root,
+    &app_env_configs,
+    &task_database,
+    &storyteller_creds,
+    &sora_succeeded_tasks_by_id,
+    &local_sqlite_tasks_by_sora_task_id,
+  ).await?;
+
+  Ok(())
+}

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/poll_sora_2_tasks.rs

@@ -0,0 +1,116 @@
+use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
+use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
+use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
+use crate::core::events::generation_events::generation_failed_event::GenerationFailedEvent;
+use crate::core::events::sendable_event_trait::SendableEvent;
+use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
+use crate::core::state::data_dir::app_data_root::AppDataRoot;
+use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
+use crate::core::state::task_database::TaskDatabase;
+use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES;
+use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
+use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
+use crate::services::sora::threads::sora_task_polling::helpers::handle_failed_generations::{handle_classic_failed_generations, FailedGeneration};
+use crate::services::sora::threads::sora_task_polling::helpers::handle_successful_generations::{handle_classic_successful_generations, GenerationItem, SuccessfulGeneration};
+use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
+use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
+use enums::common::generation_provider::GenerationProvider;
+use enums::common::model_type::ModelType;
+use enums::tauri::tasks::task_status;
+use errors::AnyhowResult;
+use idempotency::uuid::generate_random_uuid;
+use log::{error, info, warn};
+use once_cell::sync::Lazy;
+use openai_sora_client::creds::sora_credential_set::SoraCredentialSet;
+use openai_sora_client::recipes::list_classic_sora_tasks_with_session_auto_renew::list_classic_sora_tasks_with_session_auto_renew;
+use openai_sora_client::recipes::list_sora2_drafts::list_sora2_drafts_with_session_auto_renew::list_sora2_drafts_with_session_auto_renew;
+use openai_sora_client::requests::common::task_id::TaskId;
+use openai_sora_client::requests::list_classic_tasks::list_classic_tasks::TaskStatus;
+use openai_sora_client::requests::list_sora2_drafts::list_sora2_drafts::Draft;
+use reqwest::Url;
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList};
+use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
+use std::collections::{HashMap, HashSet};
+use std::fs::File;
+use std::io::Write;
+use std::path::PathBuf;
+use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs};
+use storyteller_client::endpoints::prompts::create_prompt::create_prompt;
+use tauri::AppHandle;
+use tempdir::TempDir;
+
+pub async fn poll_sora_2_tasks(
+  app_handle: &AppHandle,
+  app_env_configs: &AppEnvConfigs,
+  task_database: &TaskDatabase,
+  sora_creds_manager: &SoraCredentialManager,
+  sora_creds: &SoraCredentialSet,
+  storyteller_creds_manager: &StorytellerCredentialManager,
+  sora_task_queue: &SoraTaskQueue,
+  app_data_root: &AppDataRoot,
+  local_sqlite_tasks_by_sora_task_id: &HashMap<String, Task>,
+) -> AnyhowResult<()> {
+
+  let (sora_response, maybe_new_creds) =
+      list_sora2_drafts_with_session_auto_renew(&sora_creds).await?;
+
+  if let Some(new_creds) = maybe_new_creds {
+    info!("Saving new credentials.");
+    sora_creds_manager.set_credentials(&new_creds)?;
+  }
+
+  let drafts = sora_response.drafts;
+
+  let mut sora_succeeded_drafts_by_id = HashMap::new();
+  let mut sora_failed_drafts_by_id = HashMap::new();
+
+  let storyteller_creds = storyteller_creds_manager.get_credentials_required()?;
+
+  for draft in drafts.into_iter() {
+    match draft {
+      Draft::Success(success_draft) => {
+        sora_succeeded_drafts_by_id.insert(
+          success_draft.task_id,
+          SuccessfulGeneration {
+            prompt: Some(success_draft.prompt),
+            items: vec![GenerationItem {
+              item_id: success_draft.id,
+              url: success_draft.url,
+            }]
+          });
+      }
+      Draft::Failure(failed_draft) => {
+        sora_failed_drafts_by_id.insert(
+          failed_draft.task_id,
+          FailedGeneration {
+            reason: Some(failed_draft.reason_message)
+          }
+        );
+      }
+    }
+  }
+
+  // Clear dead tasks.
+  handle_classic_failed_generations(
+    &app_handle,
+    &task_database,
+    &local_sqlite_tasks_by_sora_task_id,
+    &sora_failed_drafts_by_id,
+    &sora_task_queue,
+  ).await?;
+
+
+  // Process succeeded tasks.
+  handle_classic_successful_generations(
+    &app_handle,
+    &app_data_root,
+    &app_env_configs,
+    &task_database,
+    &storyteller_creds,
+    &sora_succeeded_drafts_by_id,
+    &local_sqlite_tasks_by_sora_task_id,
+  ).await?;
+
+  Ok(())
+}
+

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/mod.rs

@@ -1,3 +1,2 @@
 pub mod sora_task_polling_thread;
-mod handle_successful_generations;
-mod handle_failed_generations;
+pub mod helpers;

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/sora_task_polling_thread.rs

@@ -10,8 +10,8 @@ use crate::core::state::task_database::TaskDatabase;
 use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES;
 use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
 use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
-use crate::services::sora::threads::sora_task_polling::handle_failed_generations::handle_failed_generations;
-use crate::services::sora::threads::sora_task_polling::handle_successful_generations::handle_successful_generations;
+use crate::services::sora::threads::sora_task_polling::helpers::handle_failed_generations::handle_classic_failed_generations;
+use crate::services::sora::threads::sora_task_polling::helpers::handle_successful_generations::handle_classic_successful_generations;
 use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
 use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
 use enums::common::generation_provider::GenerationProvider;
@@ -36,6 +36,8 @@ use storyteller_client::endpoints::media_files::upload_image_media_file_from_fil
 use storyteller_client::endpoints::prompts::create_prompt::create_prompt;
 use tauri::AppHandle;
 use tempdir::TempDir;
+use crate::services::sora::threads::sora_task_polling::helpers::poll_classic_sora_tasks::poll_classic_sora_tasks;
+use crate::services::sora::threads::sora_task_polling::helpers::poll_sora_2_tasks::poll_sora_2_tasks;
 
 pub async fn sora_task_polling_thread(
   app_handle: AppHandle,
@@ -87,7 +89,20 @@ async fn local_task_polling_loop(
 
     let creds = sora_creds_manager.get_credentials_required()?;
 
-    poll_sora_tasks(
+    // Map of Sora Task ID to Local Task.
+    let local_sqlite_tasks_by_sora_task_id = local_sqlite_tasks.tasks.iter()
+        .filter_map(|task| {
+          if let Some(provider_job_id) = &task.provider_job_id {
+            Some((provider_job_id.clone(), task.clone()))
+          } else {
+            None
+          }
+        })
+        .collect::<HashMap<String, Task>>();
+
+    // TODO: Only poll if we have classic items
+    
+    poll_classic_sora_tasks(
       &app_handle,
       &app_env_configs,
       &task_database,
@@ -96,91 +111,24 @@ async fn local_task_polling_loop(
       &storyteller_creds_manager,
       &sora_task_queue,
       &app_data_root,
-      local_sqlite_tasks,
+      &local_sqlite_tasks_by_sora_task_id,
+    ).await?;
+
+    // TODO: Only poll if we have new items
+    
+    poll_sora_2_tasks(
+      &app_handle,
+      &app_env_configs,
+      &task_database,
+      &sora_creds_manager,
+      &creds,
+      &storyteller_creds_manager,
+      &sora_task_queue,
+      &app_data_root,
+      &local_sqlite_tasks_by_sora_task_id,
     ).await?;
 
     tokio::time::sleep(std::time::Duration::from_millis(2_000)).await;
   }
 }
 
-async fn poll_sora_tasks(
-  app_handle: &AppHandle,
-  app_env_configs: &AppEnvConfigs,
-  task_database: &TaskDatabase,
-  sora_creds_manager: &SoraCredentialManager,
-  sora_creds: &SoraCredentialSet,
-  storyteller_creds_manager: &StorytellerCredentialManager,
-  sora_task_queue: &SoraTaskQueue,
-  app_data_root: &AppDataRoot,
-  local_sqlite_tasks: TaskList,
-) -> AnyhowResult<()> {
-
-  let local_tasks = local_sqlite_tasks.tasks;
-
-  // Map of Sora Task ID to Local Task.
-  let local_sqlite_tasks_by_sora_task_id = local_tasks.iter()
-      .filter_map(|task| {
-        if let Some(provider_job_id) = &task.provider_job_id {
-          Some((provider_job_id.clone(), task.clone()))
-        } else {
-          None
-        }
-      })
-      .collect::<HashMap<String, Task>>();
-
-  let (sora_response, maybe_new_creds) =
-      list_classic_sora_tasks_with_session_auto_renew(&sora_creds).await?;
-
-  if let Some(new_creds) = maybe_new_creds {
-    info!("Saving new credentials.");
-    sora_creds_manager.set_credentials(&new_creds)?;
-  }
-
-  let sora_items = sora_response.task_responses;
-
-  let mut sora_succeeded_tasks_by_id = HashMap::new();
-  let mut sora_failed_tasks_by_id = HashMap::new();
-
-  let storyteller_creds = storyteller_creds_manager.get_credentials_required()?;
-
-  for task in sora_items.iter() {
-
-    match &task.status {
-      TaskStatus::Succeeded => {
-        sora_succeeded_tasks_by_id.insert(task.id.clone(), task.clone());
-      }
-      TaskStatus::Failed => {
-        sora_failed_tasks_by_id.insert(task.id.clone(), task.clone());
-      }
-      TaskStatus::Queued => {}
-      TaskStatus::Running => {}
-      TaskStatus::Unknown(unknown_status) => {
-        warn!("Unknown task status: {:?}", unknown_status);
-      }
-    }
-  }
-
-  // Clear dead tasks.
-  handle_failed_generations(
-    &app_handle,
-    &task_database,
-    &local_sqlite_tasks_by_sora_task_id,
-    &sora_failed_tasks_by_id,
-    &sora_task_queue,
-  ).await?;
-
-
-  // Process succeeded tasks.
-  handle_successful_generations(
-    &app_handle,
-    &app_data_root,
-    &app_env_configs,
-    &task_database,
-    &storyteller_creds,
-    &sora_succeeded_tasks_by_id,
-    &local_sqlite_tasks_by_sora_task_id,
-  ).await?;
-
-  Ok(())
-}
-