From 4f85d27a88997d40018d211e80bae95c48ec3bb3 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 15 Jul 2025 21:00:05 -0400 Subject: [PATCH] artcraft task polling (wip) --- .../src/jobs/list_session_jobs.rs | 43 ++++++ .../storyteller_client/src/jobs/mod.rs | 1 + .../api_clients/storyteller_client/src/lib.rs | 1 + .../src/utils/basic_json_get_request.rs | 63 ++++++++ .../storyteller_client/src/utils/mod.rs | 1 + .../artcraft/src/core/state/task_database.rs | 1 + .../artcraft/src/services/storyteller/mod.rs | 3 +- .../src/services/storyteller/threads/mod.rs | 1 + .../storyteller_task_polling_thread.rs | 144 ++++++++++++++++++ .../database/sqlite_tasks/src/connection.rs | 1 + .../schema/database/sqlite_tasks/src/error.rs | 14 +- .../list_tasks_by_provider_and_tokens.rs | 104 +++++++++++++ .../database/sqlite_tasks/src/queries/mod.rs | 3 +- .../src/common/responses/job_details.rs | 5 +- .../src/common/responses/media_links.rs | 7 +- .../src/jobs/list_session_jobs.rs | 10 +- 16 files changed, 384 insertions(+), 18 deletions(-) create mode 100644 crates/api_clients/storyteller_client/src/jobs/list_session_jobs.rs create mode 100644 crates/api_clients/storyteller_client/src/jobs/mod.rs create mode 100644 crates/api_clients/storyteller_client/src/utils/basic_json_get_request.rs create mode 100644 crates/desktop/artcraft/src/services/storyteller/threads/mod.rs create mode 100644 crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs create mode 100644 crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs diff --git a/crates/api_clients/storyteller_client/src/jobs/list_session_jobs.rs b/crates/api_clients/storyteller_client/src/jobs/list_session_jobs.rs new file mode 100644 index 0000000000..accf2a6043 --- /dev/null +++ b/crates/api_clients/storyteller_client/src/jobs/list_session_jobs.rs @@ -0,0 +1,43 @@ +use crate::credentials::storyteller_credential_set::StorytellerCredentialSet; +use crate::error::storyteller_error::StorytellerError; +use crate::utils::api_host::ApiHost; +use crate::utils::basic_json_get_request::basic_json_get_request; +use artcraft_api_defs::jobs::list_session_jobs::{ListSessionJobsSuccessResponse, LIST_SESSION_JOBS_URL_PATH}; +use enums::common::job_status_plus::JobStatusPlus; +use std::collections::HashSet; + +#[derive(Clone)] +pub enum States { + All, + Include(HashSet), + Exclude(HashSet), +} + +pub async fn list_session_jobs( + api_host: &ApiHost, + maybe_creds: Option<&StorytellerCredentialSet>, + job_states: States, +) -> Result { + + let url = match job_states { + States::All => + LIST_SESSION_JOBS_URL_PATH.to_string(), + States::Include(states) => + format!("{}?include_states={}", LIST_SESSION_JOBS_URL_PATH, states_string(&states)), + States::Exclude(states) => + format!("{}?exclude_states={}", LIST_SESSION_JOBS_URL_PATH, states_string(&states)), + }; + + Ok(basic_json_get_request( + api_host, + &url, + maybe_creds, + ).await?) +} + +fn states_string(states: &HashSet) -> String { + states.iter() + .map(|state| state.to_str()) + .collect::>() + .join(",") +} \ No newline at end of file diff --git a/crates/api_clients/storyteller_client/src/jobs/mod.rs b/crates/api_clients/storyteller_client/src/jobs/mod.rs new file mode 100644 index 0000000000..164271aad7 --- /dev/null +++ b/crates/api_clients/storyteller_client/src/jobs/mod.rs @@ -0,0 +1 @@ +pub mod list_session_jobs; \ No newline at end of file diff --git a/crates/api_clients/storyteller_client/src/lib.rs b/crates/api_clients/storyteller_client/src/lib.rs index 228b2137de..e0af3d9aac 100644 --- a/crates/api_clients/storyteller_client/src/lib.rs +++ b/crates/api_clients/storyteller_client/src/lib.rs @@ -3,6 +3,7 @@ pub mod credentials; pub mod error; pub mod generate; +pub mod jobs; pub mod media_files; pub mod recipes; pub mod shared_response_types; diff --git a/crates/api_clients/storyteller_client/src/utils/basic_json_get_request.rs b/crates/api_clients/storyteller_client/src/utils/basic_json_get_request.rs new file mode 100644 index 0000000000..7df8fc9673 --- /dev/null +++ b/crates/api_clients/storyteller_client/src/utils/basic_json_get_request.rs @@ -0,0 +1,63 @@ +use crate::credentials::storyteller_credential_set::StorytellerCredentialSet; +use crate::error::api_error::ApiError; +use crate::error::client_error::ClientError; +use crate::error::storyteller_error::StorytellerError; +use crate::utils::api_host::ApiHost; +use crate::utils::constants::{APPLICATION_JSON, USER_AGENT}; +use crate::utils::filter_bad_response::filter_bad_response; +use artcraft_api_defs::generate::image::generate_flux_1_schnell_text_to_image::{GenerateFlux1SchnellTextToImageRequest, GenerateFlux1SchnellTextToImageResponse, GENERATE_FLUX_1_SCHNELL_TEXT_TO_IMAGE_PATH}; +use log::{debug, info}; +use reqwest::Client; +use serde::de::DeserializeOwned; +use serde::{Deserialize, Serialize}; + +pub async fn basic_json_get_request( + api_host: &ApiHost, + route_path: &str, + maybe_creds: Option<&StorytellerCredentialSet>, +) -> Result { + + let url = get_route(api_host, route_path); + + info!("Requesting {:?}", &url); + + let client = Client::builder() + .gzip(true) + .build() + .map_err(|err| StorytellerError::Client(ClientError::from(err)))?; + + let mut request_builder = client.get(url) + .header("User-Agent", USER_AGENT) + .header("Accept", APPLICATION_JSON); + + if let Some(creds) = maybe_creds { + if let Some(header) = &creds.maybe_as_cookie_header() { + request_builder = request_builder.header("Cookie", header); + } + } + + let response = request_builder + .send() + .await + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + let response = filter_bad_response(response) + .await + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + let response_body = &response.text() + .await + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + debug!("Response body: {:?}", response_body); + + let response = serde_json::from_str(&response_body) + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + Ok(response) +} + +fn get_route(api_host: &ApiHost, route_path: &str) -> String { + let api_hostname = api_host.to_api_hostname_and_scheme(); + format!("{}{}", api_hostname, route_path) +} diff --git a/crates/api_clients/storyteller_client/src/utils/mod.rs b/crates/api_clients/storyteller_client/src/utils/mod.rs index defbbbc4c5..2211888f7f 100644 --- a/crates/api_clients/storyteller_client/src/utils/mod.rs +++ b/crates/api_clients/storyteller_client/src/utils/mod.rs @@ -1,4 +1,5 @@ pub mod api_host; +pub mod basic_json_get_request; pub mod basic_json_post_request; pub mod constants; pub mod filter_bad_response; diff --git a/crates/desktop/artcraft/src/core/state/task_database.rs b/crates/desktop/artcraft/src/core/state/task_database.rs index 9df07c4cd2..4467785f0d 100644 --- a/crates/desktop/artcraft/src/core/state/task_database.rs +++ b/crates/desktop/artcraft/src/core/state/task_database.rs @@ -2,6 +2,7 @@ use crate::core::state::data_dir::app_data_root::AppDataRoot; use errors::AnyhowResult; use sqlite_tasks::connection::TaskDbConnection; +#[derive(Clone)] pub struct TaskDatabase { connection: TaskDbConnection, } diff --git a/crates/desktop/artcraft/src/services/storyteller/mod.rs b/crates/desktop/artcraft/src/services/storyteller/mod.rs index 0d0b840411..51f42a9be7 100644 --- a/crates/desktop/artcraft/src/services/storyteller/mod.rs +++ b/crates/desktop/artcraft/src/services/storyteller/mod.rs @@ -1 +1,2 @@ -pub mod state; \ No newline at end of file +pub mod state; +pub mod threads; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/mod.rs b/crates/desktop/artcraft/src/services/storyteller/threads/mod.rs new file mode 100644 index 0000000000..09250df48d --- /dev/null +++ b/crates/desktop/artcraft/src/services/storyteller/threads/mod.rs @@ -0,0 +1 @@ +pub mod storyteller_task_polling_thread; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs new file mode 100644 index 0000000000..2b92838935 --- /dev/null +++ b/crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread.rs @@ -0,0 +1,144 @@ +use std::collections::HashMap; +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::services::sora::state::sora_credential_manager::SoraCredentialManager; +use crate::services::sora::state::sora_task_queue::SoraTaskQueue; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use errors::AnyhowResult; +use log::{error, info}; +use reqwest::Url; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use anyhow::anyhow; +use storyteller_client::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file; +use tauri::AppHandle; +use tempdir::TempDir; +use enums::common::generation_provider::GenerationProvider; +use enums::common::job_status_plus::JobStatusPlus; +use enums::tauri::tasks::task_status::TaskStatus; +use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task}; +use storyteller_client::error::api_error::ApiError; +use storyteller_client::error::storyteller_error::StorytellerError; +use storyteller_client::jobs::list_session_jobs::{list_session_jobs, States}; +use crate::core::state::task_database::TaskDatabase; + +pub async fn storyteller_task_polling_thread( + app_handle: AppHandle, + app_env_configs: AppEnvConfigs, + app_data_root: AppDataRoot, + task_database: TaskDatabase, + storyteller_creds_manager: StorytellerCredentialManager, +) -> ! { + loop { + let res = polling_loop( + &app_handle, + &app_env_configs, + &app_data_root, + &task_database, + &storyteller_creds_manager, + ).await; + if let Err(err) = res { + error!("An error occurred: {:?}", err); + } + // NB: Only sleep if an error occurs. + tokio::time::sleep(std::time::Duration::from_millis(30_000)).await; + } +} + +async fn polling_loop( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> AnyhowResult<()> { + loop { + // Wait before next request for jobs. + tokio::time::sleep(std::time::Duration::from_millis(5_000)).await; + + let creds = storyteller_creds_manager.get_credentials()?; + + let result = list_session_jobs( + &app_env_configs.storyteller_host, + creds.as_ref(), + States::All, + ).await; + + let jobs = match result { + Ok(result) => result.jobs, + Err(err) => { + match &err { + StorytellerError::Api(ApiError::TooManyRequests(message)) => { + error!("Too many requests (sleeping): {:?}", message); + tokio::time::sleep(std::time::Duration::from_millis(60_000)).await; + } + _ => {} + } + return Err(anyhow!(err)); + } + }; + + let job_ids = jobs.iter() + .map(|job| job.job_token.to_string()) + .collect::>(); + + let tasks = list_tasks_by_provider_and_tokens(ListTasksArgs { + db: task_database.get_connection(), + provider: GenerationProvider::Artcraft, + provider_job_ids: Some(job_ids), + }).await?; + + let tasks = tasks.tasks; + + let jobs_by_id = jobs.iter() + .map(|job| (job.job_token.to_string(), job)) + .collect::>(); + + let tasks_by_provider_job_id = tasks.iter() + .filter_map(|task| { + if let Some(provider_job_id) = &task.provider_job_id { + Some((provider_job_id.clone(), task.clone())) + } else { + None + } + }) + //.map(|task| (task.provider_job_id.clone(), task.clone())) + .collect::>(); + + for job in jobs.iter() { + match job.status.status { + JobStatusPlus::CompleteSuccess => {} // Fall-through. + _ => continue, + } + + let task = match tasks_by_provider_job_id.get(job.job_token.as_str()) { + Some(job) => job, + None => continue, + }; + + match task.status { + TaskStatus::CompleteSuccess => continue, + _ => {} // Fall-through. + } + + let event = GenerationCompleteEvent { + //media_file_token: result.media_file_token, + action: Some(GenerationAction::GenerateImage), // TODO + service: GenerationServiceProvider::Artcraft, // TODO + model: None, // TODO + }; + + if let Err(err) = event.send(&app_handle) { + error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open + } + } + + } +} diff --git a/crates/schema/database/sqlite_tasks/src/connection.rs b/crates/schema/database/sqlite_tasks/src/connection.rs index 50972d4670..2a46d832d0 100644 --- a/crates/schema/database/sqlite_tasks/src/connection.rs +++ b/crates/schema/database/sqlite_tasks/src/connection.rs @@ -3,6 +3,7 @@ use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode}; use sqlx::{Pool, Sqlite, SqlitePool}; use std::path::Path; +#[derive(Clone)] pub struct TaskDbConnection { pool: Pool, } diff --git a/crates/schema/database/sqlite_tasks/src/error.rs b/crates/schema/database/sqlite_tasks/src/error.rs index a04b9ec225..708c4fd196 100644 --- a/crates/schema/database/sqlite_tasks/src/error.rs +++ b/crates/schema/database/sqlite_tasks/src/error.rs @@ -3,12 +3,13 @@ use std::fmt::{Display, Formatter}; #[derive(Debug)] pub enum SqliteTasksError { - SqlxError(sqlx::Error), - //TaskNotFound, - //TaskAlreadyExists, - //InvalidTaskStatus, - //InvalidTaskType, - //InvalidGenerationProvider, + SqlxError(sqlx::Error), + TaskParseError(String), + //TaskNotFound, + //TaskAlreadyExists, + //InvalidTaskStatus, + //InvalidTaskType, + //InvalidGenerationProvider, } impl Error for SqliteTasksError {} @@ -17,6 +18,7 @@ impl Display for SqliteTasksError { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { SqliteTasksError::SqlxError(err) => write!(f, "SQLx error: {}", err), + SqliteTasksError::TaskParseError(err) => write!(f, "Task parse error: {}", err), } } } diff --git a/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs new file mode 100644 index 0000000000..36cce2dffb --- /dev/null +++ b/crates/schema/database/sqlite_tasks/src/queries/list_tasks_by_provider_and_tokens.rs @@ -0,0 +1,104 @@ +use crate::connection::TaskDbConnection; +use crate::error::SqliteTasksError; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_status::TaskStatus; +use enums::tauri::tasks::task_type::TaskType; +use sqlx::{QueryBuilder, Sqlite}; + +pub struct ListTasksArgs<'a> { + pub db: &'a TaskDbConnection, + pub provider: GenerationProvider, + pub provider_job_ids: Option>, +} + +pub struct TaskList { + pub tasks: Vec, +} + +#[derive(Debug, Clone)] +pub struct Task { + pub status: TaskStatus, + pub task_type: TaskType, + pub provider: GenerationProvider, + pub provider_job_id: Option, + pub frontend_subscriber_id: Option, + pub frontend_subscriber_payload: Option, +} + +#[derive(Debug)] +#[derive(sqlx::FromRow)] +struct RawTask { + status: String, + task_type: String, + provider: String, + provider_job_id: Option, + frontend_subscriber_id: Option, + frontend_subscriber_payload: Option, +} + +pub async fn list_tasks_by_provider_and_tokens( + args: ListTasksArgs<'_>, +) -> Result { + + let mut query_builder: QueryBuilder = QueryBuilder::new(r#" + SELECT + id, + task_status, + task_type, + provider, + provider_job_id, + frontend_subscriber_id, + frontend_subscriber_payload + FROM tasks + WHERE provider = ? + "#); + + // TODO(bt,2025-07-15): Fix this. The sqlx mysql queries never required temporaries + let provider = args.provider.to_string(); + + query_builder.push_bind(provider); + + if let Some(provider_job_ids) = args.provider_job_ids { + query_builder.push(" AND provider_job_id IN ("); + //for (i, job_id) in provider_job_ids.iter().enumerate() { + // if i > 0 { + // query_builder.push(", "); + // } + // query_builder.push_bind(job_id); + //} + //query_builder.push(")"); + + let mut separated = query_builder.separated(", "); + + for job_id in provider_job_ids.into_iter() { + separated.push_bind(job_id); + } + + separated.push_unseparated(") "); + } + + //let query = query_builder.build(); + + let query = query_builder.build_query_as::(); + let results = query.fetch_all(args.db.get_pool()).await?; + + let mut tasks: Vec = Vec::new(); + + for task in results { + tasks.push(Task { + status: TaskStatus::from_str(&task.status) + .map_err(|err| SqliteTasksError::TaskParseError(err))?, + task_type: TaskType::from_str(&task.task_type) + .map_err(|err| SqliteTasksError::TaskParseError(err))?, + provider: GenerationProvider::from_str(&task.provider) + .map_err(|err| SqliteTasksError::TaskParseError(err))?, + provider_job_id: task.provider_job_id, + frontend_subscriber_id: task.frontend_subscriber_id, + frontend_subscriber_payload: task.frontend_subscriber_payload, + }); + } + + Ok(TaskList { + tasks, + }) +} diff --git a/crates/schema/database/sqlite_tasks/src/queries/mod.rs b/crates/schema/database/sqlite_tasks/src/queries/mod.rs index cd1a4daa46..4f1fbc7424 100644 --- a/crates/schema/database/sqlite_tasks/src/queries/mod.rs +++ b/crates/schema/database/sqlite_tasks/src/queries/mod.rs @@ -1 +1,2 @@ -pub mod create_task; \ No newline at end of file +pub mod create_task; +pub mod list_tasks_by_provider_and_tokens; \ No newline at end of file diff --git a/crates/schema/public/api/artcraft_api_defs/src/common/responses/job_details.rs b/crates/schema/public/api/artcraft_api_defs/src/common/responses/job_details.rs index d6d572043b..1ff8f3229d 100644 --- a/crates/schema/public/api/artcraft_api_defs/src/common/responses/job_details.rs +++ b/crates/schema/public/api/artcraft_api_defs/src/common/responses/job_details.rs @@ -1,9 +1,10 @@ +use serde::Deserialize; use serde::Serialize; use tokens::tokens::media_files::MediaFileToken; use utoipa::ToSchema; /// Details about submitted live portrait jobs (request arguments only) -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct JobDetailsLivePortraitRequest { /// Source media token /// This can be an image or a video media file token @@ -20,7 +21,7 @@ pub struct JobDetailsLivePortraitRequest { } /// Details about submitted lipsync jobs (request arguments only) -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct JobDetailsLipsyncRequest { /// Media file token for the source audio. /// This is probably an audio file, but in the future we might pull audio from video. diff --git a/crates/schema/public/api/artcraft_api_defs/src/common/responses/media_links.rs b/crates/schema/public/api/artcraft_api_defs/src/common/responses/media_links.rs index a8868a65e1..367dae629d 100644 --- a/crates/schema/public/api/artcraft_api_defs/src/common/responses/media_links.rs +++ b/crates/schema/public/api/artcraft_api_defs/src/common/responses/media_links.rs @@ -1,9 +1,10 @@ +use serde::Deserialize; use serde::Serialize; -use utoipa::ToSchema; use url::Url; +use utoipa::ToSchema; /// Links to media file locations (bucket, CDN, etc.) -#[derive(Serialize, ToSchema, Debug, Clone, Eq, PartialEq)] +#[derive(Serialize, Deserialize, ToSchema, Debug, Clone, Eq, PartialEq)] pub struct MediaLinks { /// Primary link to the asset via the CDN. pub cdn_url: Url, @@ -20,7 +21,7 @@ pub struct MediaLinks { pub maybe_video_previews: Option, } -#[derive(Serialize, ToSchema, Debug, Clone, Eq, PartialEq)] +#[derive(Serialize, Deserialize, ToSchema, Debug, Clone, Eq, PartialEq)] pub struct VideoPreviews { /// A static single frame preview image of the video. pub still: Url, diff --git a/crates/schema/public/api/artcraft_api_defs/src/jobs/list_session_jobs.rs b/crates/schema/public/api/artcraft_api_defs/src/jobs/list_session_jobs.rs index f4b85fcc38..281fc20a2d 100644 --- a/crates/schema/public/api/artcraft_api_defs/src/jobs/list_session_jobs.rs +++ b/crates/schema/public/api/artcraft_api_defs/src/jobs/list_session_jobs.rs @@ -27,7 +27,7 @@ pub struct ListSessionJobsQueryParams { // TODO(bt,2024-04-23): Add the ability for users to dismiss completed/dead jobs from view. } -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct ListSessionJobsSuccessResponse { pub success: bool, @@ -35,7 +35,7 @@ pub struct ListSessionJobsSuccessResponse { pub jobs: Vec, } -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct ListSessionJobsItem { pub job_token: InferenceJobToken, @@ -48,7 +48,7 @@ pub struct ListSessionJobsItem { } /// Details about what the user requested for generation -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct ListSessionRequestDetailsResponse { pub inference_category: InferenceCategory, pub maybe_model_type: Option, @@ -73,7 +73,7 @@ pub struct ListSessionRequestDetailsResponse { } /// Details about the ongoing job status -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct ListSessionStatusDetailsResponse { /// Primary status from the database (a state machine). pub status: JobStatusPlus, @@ -107,7 +107,7 @@ pub struct ListSessionStatusDetailsResponse { } /// Details about the completed result (if any) -#[derive(Serialize, ToSchema)] +#[derive(Serialize, Deserialize, ToSchema)] pub struct ListSessionResultDetailsResponse { pub entity_type: String, pub entity_token: String,