polling wip

crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs

@@ -78,8 +78,10 @@ pub async fn handle_tauri_startup(
   tauri::async_runtime::spawn(midjourney_long_polling_thread(
     app.clone(),
     app_env_configs.clone(),
+    root.clone(),
     task_database.clone(),
     mj_creds_manager.clone(),
+    storyteller_creds_manager.clone(),
   ));
 
   spawn_discord_presence_thread()?;

crates/desktop/artcraft/src/services/midjourney/mod.rs

@@ -1,4 +1,5 @@
-pub mod threads;
-pub mod state;
 pub mod commands;
+pub mod state;
+pub mod threads;
+pub mod utils;
 pub mod windows;

crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs

@@ -1,31 +1,47 @@
 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::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
 use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
-use errors::AnyhowResult;
-use log::error;
-use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
-use midjourney_client::endpoints::imagine::{imagine, ImagineRequest};
-use tauri::AppHandle;
 use enums::common::generation_provider::GenerationProvider;
 use enums::tauri::tasks::task_status::TaskStatus;
+use errors::AnyhowResult;
+use log::{error, info};
+use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
+use midjourney_client::endpoints::imagine::{imagine, ImagineRequest};
 use midjourney_client::utils::get_image_url::get_image_url;
-use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs};
+use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task};
+use std::collections::HashMap;
+use std::fs::File;
+use std::io::Write;
+use std::path::PathBuf;
+use tauri::AppHandle;
+use url::Url;
+use storyteller_client::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs};
+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::services::midjourney::utils::download_midjourney_image::download_midjourney_image;
 
 /// This thread is responsible for picking up tasks that fell through the cracks of
 /// the faster websocket thread.
 pub async fn midjourney_long_polling_thread(
   app_handle: AppHandle,
   app_env_configs: AppEnvConfigs,
+  app_data_root: AppDataRoot,
   task_database: TaskDatabase,
   creds: MidjourneyCredentialManager,
+  storyteller_creds_manager: StorytellerCredentialManager,
 ) -> ! {
   loop {
     let res = polling_loop(
       &app_handle,
       &app_env_configs,
+      &app_data_root,
       &task_database,
       &creds,
+      &storyteller_creds_manager,
     ).await;
     if let Err(err) = res {
       error!("An error occurred: {:?}", err);
@@ -38,8 +54,10 @@ pub async fn midjourney_long_polling_thread(
 async fn polling_loop(
   app_handle: &AppHandle,
   app_env_configs: &AppEnvConfigs,
+  app_data_root: &AppDataRoot,
   task_database: &TaskDatabase,
   creds: &MidjourneyCredentialManager,
+  storyteller_creds_manager: &StorytellerCredentialManager,
 ) -> AnyhowResult<()> {
   loop {
     if !creds.session_appears_active()? {
@@ -47,6 +65,9 @@ async fn polling_loop(
       continue;
     }
 
+    // TODO: Graceful wait, fix this long function body
+    let storyteller_creds = storyteller_creds_manager.get_credentials_required()?;
+
     let cookies = creds.maybe_copy_cookie_store()?;
 
     let cookies = match cookies {
@@ -80,40 +101,108 @@ async fn polling_loop(
       page_size: None,
     }).await?;
 
-    let job_ids = result.items.iter()
-        .filter_map(|item| item.id.as_ref())
-        .map(|id| id.to_string())
-        .collect::<Vec<_>>();
+    let midjourney_items = result.items;
 
-    let tasks = list_tasks_by_provider_and_tokens(ListTasksArgs {
+    let midjourney_items_by_id = {
+      let mut hash = HashMap::new();
+      for item in midjourney_items.iter() {
+        if let Some(id) = &item.id {
+          hash.insert(id.to_string(), item.clone());
+        }
+      }
+      hash
+    };
+
+    //let midjourney_item_ids = midjourney_items.iter()
+    //    .filter_map(|item| item.id.as_ref())
+    //    .map(|id| id.to_string())
+    //    .collect::<Vec<_>>();
+
+    let midjourney_item_ids = midjourney_items_by_id
+        .keys()
+        .map(|id| id.to_string())
+        .collect();
+
+    let local_tasks = list_tasks_by_provider_and_tokens(ListTasksArgs {
       db: task_database.get_connection(),
       provider: GenerationProvider::Midjourney,
-      provider_job_ids: Some(job_ids),
+      provider_job_ids: Some(midjourney_item_ids),
     }).await?;
 
-    let tasks = tasks.tasks;
+    let local_tasks = local_tasks.tasks;
 
-    for task in tasks {
-      let job_id = match task.provider_job_id {
-        Some(job_id) => job_id,
+    // Map of Midjourney Job ID to Local Task.
+    let local_tasks_by_midjourney_job_id = local_tasks.iter()
+        .filter(|task| match task.status {
+          TaskStatus::Pending => true,
+          TaskStatus::Started => true,
+          TaskStatus::AttemptFailed => true,
+          _ => false,
+        })
+        .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: If we introduce another job polling mechanism, we may need to handle concurrency.
+
+    for (midjourney_job_id, local_task) in local_tasks_by_midjourney_job_id.iter() {
+      let midjourney_item = match midjourney_items_by_id.get(midjourney_job_id) {
+        Some(item) => item,
         None => continue,
       };
 
-      match task.status {
-        TaskStatus::Pending => {}, // Fall-through
-        TaskStatus::Started => {}, // Fall-through
-        TaskStatus::AttemptFailed => {}, // Fall-through
-        _ => continue,
-      }
-
       for index in 0..4 {
-        let url = get_image_url(&job_id, index)?;
+        info!("Downloading generated Midjourney file...");
+        let download_path = download_midjourney_image(midjourney_job_id, index, app_data_root).await?;
 
-        println!("Image URL for job {} index {}: {}", job_id, index, url);
+        info!("Uploading to backend...");
+
+        let result = upload_image_media_file_from_file(UploadImageFromFileArgs {
+          api_host: &app_env_configs.storyteller_host,
+          maybe_creds: Some(&storyteller_creds),
+          path: download_path,
+          is_intermediate_system_file: false,
+        }).await?;
+
+        info!("Uploaded to API backend: {:?}", result.media_file_token);
+
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
       }
+
+      let event = GenerationCompleteEvent {
+        //media_file_token: result.media_file_token,
+        action: Some(GenerationAction::GenerateImage),
+        service: GenerationServiceProvider::Midjourney,
+        model: None,
+      };
+
+      if let Err(err) = event.send(&app_handle) {
+        error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open
+      }
+
+      tokio::time::sleep(std::time::Duration::from_millis(2_000)).await;
     }
 
-    //println!("Response: {:?}\n\n", tasks);
+    //for task in tasks {
+    //  let job_id = match task.provider_job_id {
+    //    Some(job_id) => job_id,
+    //    None => continue,
+    //  };
+    //  match task.status {
+    //    TaskStatus::Pending => {}, // Fall-through
+    //    TaskStatus::Started => {}, // Fall-through
+    //    TaskStatus::AttemptFailed => {}, // Fall-through
+    //    _ => {
+    //      tasks_by_provider_job_id.remove(&job_id);
+    //      continue
+    //    },
+    //  }
+    //}
 
     tokio::time::sleep(std::time::Duration::from_millis(60_000)).await;
   }

crates/desktop/artcraft/src/services/midjourney/utils/download_midjourney_image.rs

@@ -0,0 +1,27 @@
+use crate::core::state::data_dir::app_data_root::AppDataRoot;
+use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
+use errors::AnyhowResult;
+use midjourney_client::utils::get_image_url::get_image_url;
+use std::fs::File;
+use std::io::Write;
+use std::path::PathBuf;
+use url::Url;
+
+pub async fn download_midjourney_image(job_id: &str, index: u8, app_data_root: &AppDataRoot) -> AnyhowResult<PathBuf> {
+  let url = get_image_url(job_id, index)?;
+  let parsed_url = Url::parse(&url)?;
+
+  let response = reqwest::get(&url).await?;
+  let image_bytes = response.bytes().await?;
+
+  let ext = parsed_url.path().split(".").last().unwrap_or("png");
+
+  let tempdir = app_data_root.temp_dir().path();
+  let download_filename = format!("{}_{}.{}", job_id, index, ext);
+  let download_path = tempdir.join(download_filename);
+
+  let mut file = File::create(&download_path)?;
+  file.write_all(&image_bytes)?;
+
+  Ok(download_path)
+}

crates/desktop/artcraft/src/services/midjourney/utils/mod.rs

@@ -0,0 +1 @@
+pub mod download_midjourney_image;