mirror of
https://github.com/storytold/artcraft.git
synced 2026-10-20 17:02:53 -04:00
tauri task failure reasons (#1130)
This commit is contained in:
Brandon Thomas
+17
-5
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"db_name": "SQLite",
|
||||
"query": "\n SELECT\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_caller,\n frontend_subscriber_id,\n frontend_subscriber_payload,\n on_complete_primary_media_file_token,\n on_complete_primary_media_file_class,\n on_complete_batch_token,\n on_complete_primary_media_file_cdn_url,\n on_complete_primary_media_file_thumbnail_url_template,\n created_at as \"created_at: DateTime<Utc>\",\n updated_at as \"updated_at: DateTime<Utc>\",\n completed_at as \"completed_at: DateTime<Utc>\"\n FROM tasks\n WHERE is_dismissed_by_user == 0\n ",
|
||||
"query": "\n SELECT\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_caller,\n frontend_subscriber_id,\n frontend_subscriber_payload,\n on_complete_primary_media_file_token,\n on_complete_primary_media_file_class,\n on_complete_batch_token,\n on_complete_primary_media_file_cdn_url,\n on_complete_primary_media_file_thumbnail_url_template,\n on_failure_type,\n on_failure_message,\n created_at as \"created_at: DateTime<Utc>\",\n updated_at as \"updated_at: DateTime<Utc>\",\n completed_at as \"completed_at: DateTime<Utc>\"\n FROM tasks\n WHERE is_dismissed_by_user == 0\n ",
|
||||
"describe": {
|
||||
"columns": [
|
||||
{
|
||||
@@ -74,18 +74,28 @@
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"name": "created_at: DateTime<Utc>",
|
||||
"name": "on_failure_type",
|
||||
"ordinal": 14,
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"name": "on_failure_message",
|
||||
"ordinal": 15,
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"name": "created_at: DateTime<Utc>",
|
||||
"ordinal": 16,
|
||||
"type_info": "Int64"
|
||||
},
|
||||
{
|
||||
"name": "updated_at: DateTime<Utc>",
|
||||
"ordinal": 15,
|
||||
"ordinal": 17,
|
||||
"type_info": "Int64"
|
||||
},
|
||||
{
|
||||
"name": "completed_at: DateTime<Utc>",
|
||||
"ordinal": 16,
|
||||
"ordinal": 18,
|
||||
"type_info": "Int64"
|
||||
}
|
||||
],
|
||||
@@ -107,10 +117,12 @@
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
]
|
||||
},
|
||||
"hash": "47f94c127abc20683b25c257aed3c2190b9bd29eb829d61392f08e10fd986767"
|
||||
"hash": "2436eb7bf4b0e23a899311a67a77f71d89eb35461e2113e43cd9d221bfa91c9a"
|
||||
}
|
||||
@@ -16,7 +16,9 @@
|
||||
-- on_complete_primary_media_file_class,
|
||||
-- on_complete_primary_media_file_cdn_url,
|
||||
-- on_complete_primary_media_file_thumbnail_url_template,
|
||||
--
|
||||
-- tasks_v6.sqlite - added two fields:
|
||||
-- on_failure_type
|
||||
-- on_failure_message
|
||||
|
||||
CREATE TABLE tasks (
|
||||
-- Task auto-incrementing primary key.
|
||||
@@ -84,6 +86,15 @@ CREATE TABLE tasks (
|
||||
-- The URL template for the thumbnail of the primary media file.
|
||||
on_complete_primary_media_file_thumbnail_url_template TEXT,
|
||||
|
||||
-- OPTIONAL.
|
||||
-- An enum-like string describing the failure reason (even for non-terminal failures).
|
||||
-- The enum is `TaskFailureType`, which is a close analogue of the backend API's `FrontendFailureCategory`.
|
||||
on_failure_type TEXT,
|
||||
|
||||
-- OPTIONAL.
|
||||
-- A human-readable message describing the failure reason (even for non-terminal failures).
|
||||
on_failure_message TEXT,
|
||||
|
||||
created_at INTEGER NOT NULL DEFAULT (unixepoch('now')),
|
||||
updated_at INTEGER NOT NULL DEFAULT (unixepoch('now')),
|
||||
completed_at INTEGER DEFAULT NULL
|
||||
|
||||
@@ -18,6 +18,7 @@ use serde_derive::{Deserialize, Serialize};
|
||||
use sqlite_tasks::queries::list_tasks_for_frontend::list_tasks_for_frontend;
|
||||
use artcraft_client::endpoints::media_files::delete_media_file::delete_media_file;
|
||||
use tauri::{AppHandle, State};
|
||||
use enums::tauri::tasks::task_failure_type::TaskFailureType;
|
||||
use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass;
|
||||
use tokens::tokens::batch_generations::BatchGenerationToken;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
@@ -31,12 +32,20 @@ pub struct GetTaskQueueCommandResponse {
|
||||
#[derive(Serialize)]
|
||||
pub struct TaskQueueItem {
|
||||
pub id: TaskId,
|
||||
|
||||
pub task_status: TaskStatus,
|
||||
pub task_type: TaskType,
|
||||
pub model_type: Option<TaskModelType>,
|
||||
|
||||
pub provider: Option<GenerationProvider>,
|
||||
pub provider_job_id: Option<String>,
|
||||
|
||||
/// If the item is done, these will be filled out.
|
||||
pub completed_item: Option<CompletedItemData>,
|
||||
|
||||
/// If the item failed, this will be filled out.
|
||||
pub failure_reason: Option<FailedItemData>,
|
||||
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
@@ -63,6 +72,12 @@ pub struct MediaFileData {
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct FailedItemData {
|
||||
pub failure_type: TaskFailureType,
|
||||
pub failure_message: Option<String>,
|
||||
}
|
||||
|
||||
impl SerializeMarker for GetTaskQueueCommandResponse {}
|
||||
|
||||
#[tauri::command]
|
||||
@@ -103,6 +118,7 @@ pub async fn handle_request(
|
||||
|
||||
for task in tasks.tasks.into_iter() {
|
||||
let mut completed_item = None;
|
||||
let mut failure_reason = None;
|
||||
|
||||
if task.status == TaskStatus::CompleteSuccess {
|
||||
let token_and_url = task.on_complete_primary_media_file_token
|
||||
@@ -123,6 +139,19 @@ pub async fn handle_request(
|
||||
} else {
|
||||
warn!("Task {} is marked complete but has no primary media file token or URL.", task.id);
|
||||
}
|
||||
} else {
|
||||
// If either failure field is present, fill out the failure report.
|
||||
if let Some(failure_message) = task.on_failure_message.as_deref() {
|
||||
failure_reason = Some(FailedItemData {
|
||||
failure_type: task.on_failure_type.unwrap_or(TaskFailureType::Unknown),
|
||||
failure_message: Some(failure_message.to_string()),
|
||||
});
|
||||
} else if let Some(failure_type) = task.on_failure_type {
|
||||
failure_reason = Some(FailedItemData {
|
||||
failure_type,
|
||||
failure_message: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
transformed_tasks.push(TaskQueueItem {
|
||||
@@ -136,6 +165,7 @@ pub async fn handle_request(
|
||||
updated_at: task.updated_at,
|
||||
completed_at: task.completed_at,
|
||||
completed_item,
|
||||
failure_reason,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+1
-1
@@ -4,7 +4,7 @@ use crate::core::state::artcraft_usage_tracker::artcraft_usage_tracker::Artcraft
|
||||
use crate::core::state::task_database::TaskDatabase;
|
||||
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
|
||||
use crate::services::storyteller::threads::storyteller_activity_thread::storyteller_activity_thread;
|
||||
use crate::services::storyteller::threads::storyteller_task_polling_thread::storyteller_task_polling_thread;
|
||||
use crate::services::storyteller::threads::storyteller_task_polling_thread::storyteller_task_polling_thread::storyteller_task_polling_thread;
|
||||
use errors::AnyhowResult;
|
||||
use tauri::{AppHandle, Manager};
|
||||
|
||||
|
||||
@@ -7,4 +7,4 @@
|
||||
/// We increment this whenever we make a change to the database schema.
|
||||
/// This prevents deadlock on startup if the schema has changed from the
|
||||
/// previous version.
|
||||
pub const TASK_DATABASE_VERSION: u32 = 5;
|
||||
pub const TASK_DATABASE_VERSION: u32 = 6;
|
||||
|
||||
@@ -1,3 +1,2 @@
|
||||
pub mod events;
|
||||
pub mod storyteller_activity_thread;
|
||||
pub mod storyteller_task_polling_thread;
|
||||
+4
-22
@@ -1,34 +1,16 @@
|
||||
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
|
||||
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
|
||||
use crate::core::state::artcraft_platform_info::ArtcraftPlatformInfo;
|
||||
use crate::core::state::artcraft_usage_tracker::artcraft_usage_tracker::ArtcraftUsageTracker;
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
|
||||
use crate::core::state::os_platform::OsPlatform;
|
||||
use crate::core::state::task_database::TaskDatabase;
|
||||
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
|
||||
use anyhow::anyhow;
|
||||
use artcraft_api_defs::analytics::log_active_user::LogAppActiveUserRequest;
|
||||
use artcraft_api_defs::jobs::list_session_jobs::ListSessionJobsItem;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::common::job_status_plus::JobStatusPlus;
|
||||
use enums::tauri::tasks::task_status::TaskStatus;
|
||||
use enums::tauri::tasks::task_type::TaskType;
|
||||
use errors::AnyhowResult;
|
||||
use log::{debug, error, info};
|
||||
use os_info::Info;
|
||||
use reqwest::Url;
|
||||
use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs};
|
||||
use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
|
||||
use std::time::Instant;
|
||||
use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
|
||||
use artcraft_client::endpoints::analytics::log_active_user::log_active_user;
|
||||
use artcraft_client::endpoints::analytics::log_active_user_v2::log_active_user_v2;
|
||||
use artcraft_client::endpoints::jobs::list_session_jobs::{list_session_jobs, States};
|
||||
use artcraft_client::endpoints::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file;
|
||||
use artcraft_client::error::api_error::ApiError;
|
||||
use artcraft_client::error::storyteller_error::StorytellerError;
|
||||
use tauri::AppHandle;
|
||||
use errors::AnyhowResult;
|
||||
use log::{debug, error, info};
|
||||
use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs};
|
||||
use std::time::Instant;
|
||||
use tokens::tokens::app_session::AppSessionToken;
|
||||
|
||||
const CLIENT_NAME : &str = "artcraft";
|
||||
|
||||
-339
@@ -1,339 +0,0 @@
|
||||
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::enum_conversion::generation_provider::to_generation_service_provider;
|
||||
use crate::core::utils::enum_conversion::task_type::to_generation_action;
|
||||
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 crate::services::storyteller::threads::events::maybe_handle_frontend_caller_notification::maybe_handle_frontend_caller_notification;
|
||||
use crate::services::storyteller::threads::events::maybe_handle_inpainting_complete_event::maybe_handle_inpainting_complete_event;
|
||||
use crate::services::storyteller::threads::events::maybe_handle_text_to_image_complete_event::maybe_handle_text_to_image_complete_event;
|
||||
use crate::services::storyteller::threads::events::maybe_send_background_removal_complete_event::maybe_send_background_removal_complete_event;
|
||||
use anyhow::anyhow;
|
||||
use artcraft_api_defs::jobs::list_session_jobs::ListSessionJobsItem;
|
||||
use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template;
|
||||
use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::common::job_status_plus::JobStatusPlus;
|
||||
use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass;
|
||||
use enums::tauri::tasks::task_status::TaskStatus;
|
||||
use enums::tauri::tasks::task_type::TaskType;
|
||||
use errors::AnyhowResult;
|
||||
use log::{error, info};
|
||||
use reqwest::Url;
|
||||
use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs};
|
||||
use sqlite_tasks::queries::task::Task;
|
||||
use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs};
|
||||
use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
|
||||
use std::collections::HashMap;
|
||||
use std::fs::File;
|
||||
use std::io::Write;
|
||||
use std::path::PathBuf;
|
||||
use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
|
||||
use artcraft_client::endpoints::jobs::list_session_jobs::{list_session_jobs, States};
|
||||
use artcraft_client::endpoints::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file;
|
||||
use artcraft_client::error::api_error::ApiError;
|
||||
use artcraft_client::error::storyteller_error::StorytellerError;
|
||||
use tauri::AppHandle;
|
||||
use tempdir::TempDir;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
|
||||
pub async fn storyteller_task_polling_thread(
|
||||
app_handle: AppHandle,
|
||||
app_env_configs: AppEnvConfigs,
|
||||
task_database: TaskDatabase,
|
||||
storyteller_creds_manager: StorytellerCredentialManager,
|
||||
) -> ! {
|
||||
loop {
|
||||
let res = polling_loop(
|
||||
&app_handle,
|
||||
&app_env_configs,
|
||||
&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,
|
||||
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::<Vec<_>>();
|
||||
|
||||
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::<HashMap<String, _>>();
|
||||
|
||||
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::<HashMap<String, Task>>();
|
||||
|
||||
for job in jobs.iter() {
|
||||
let task = match tasks_by_provider_job_id.get(job.job_token.as_str()) {
|
||||
Some(task) => task,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
match job.status.status {
|
||||
JobStatusPlus::CompleteSuccess => {
|
||||
match task.status {
|
||||
TaskStatus::CompleteSuccess => continue,
|
||||
_ => {}
|
||||
}
|
||||
handle_successful_job(app_handle, app_env_configs, creds.as_ref(), job, task, task_database).await?;
|
||||
}
|
||||
JobStatusPlus::CompleteFailure => {
|
||||
match task.status {
|
||||
TaskStatus::CompleteFailure => continue,
|
||||
_ => {}
|
||||
}
|
||||
handle_failed_job(app_handle, job, task, task_database).await?;
|
||||
}
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_successful_job(
|
||||
app_handle: &AppHandle,
|
||||
app_env_configs: &AppEnvConfigs,
|
||||
creds: Option<&StorytellerCredentialSet>,
|
||||
job: &ListSessionJobsItem,
|
||||
task: &Task,
|
||||
task_database: &TaskDatabase,
|
||||
) -> AnyhowResult<()> {
|
||||
let maybe_primary_media_file_token = job.maybe_result
|
||||
.as_ref()
|
||||
.map(|result| MediaFileToken::new_from_str(&result.entity_token));
|
||||
|
||||
let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs {
|
||||
db: task_database.get_connection(),
|
||||
task_id: &task.id,
|
||||
maybe_batch_token: job.maybe_result
|
||||
.as_ref()
|
||||
.map(|result| result.maybe_batch_token.as_ref())
|
||||
.flatten(),
|
||||
maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(),
|
||||
maybe_primary_media_file_class: get_media_file_class(job),
|
||||
maybe_primary_media_file_thumbnail_url_template: get_thumbnail_template(job),
|
||||
maybe_primary_media_file_cdn_url: job.maybe_result
|
||||
.as_ref()
|
||||
.map(|result| result.media_links.cdn_url.as_str()),
|
||||
}).await?;
|
||||
|
||||
if !updated {
|
||||
return Ok(()); // If anything breaks with queries, don't spam events.
|
||||
}
|
||||
|
||||
send_additional_success_events(app_handle, app_env_configs, creds, job, task).await;
|
||||
|
||||
let service = to_generation_service_provider(task.provider);
|
||||
let action = to_generation_action(task.task_type);
|
||||
|
||||
let event = GenerationCompleteEvent {
|
||||
action: Some(action),
|
||||
service,
|
||||
model: None, // TODO
|
||||
};
|
||||
|
||||
if let Err(err) = event.send(app_handle) {
|
||||
error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_failed_job(
|
||||
app_handle: &AppHandle,
|
||||
job: &ListSessionJobsItem,
|
||||
task: &Task,
|
||||
task_database: &TaskDatabase,
|
||||
) -> AnyhowResult<()> {
|
||||
info!("Marking storyteller job as failed: {:?}", task.id);
|
||||
|
||||
update_task_status(UpdateTaskArgs {
|
||||
db: task_database.get_connection(),
|
||||
task_id: &task.id,
|
||||
status: TaskStatus::CompleteFailure,
|
||||
}).await?;
|
||||
|
||||
let service = to_generation_service_provider(task.provider);
|
||||
let action = to_generation_action(task.task_type);
|
||||
|
||||
let event = GenerationFailedEvent {
|
||||
action,
|
||||
service,
|
||||
model: None,
|
||||
reason: None,
|
||||
};
|
||||
|
||||
event.send_infallible(app_handle);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_additional_success_events(
|
||||
app_handle: &AppHandle,
|
||||
app_env_configs: &AppEnvConfigs,
|
||||
creds: Option<&StorytellerCredentialSet>,
|
||||
job: &ListSessionJobsItem,
|
||||
task: &Task
|
||||
) {
|
||||
info!("Attempting to dispatch events for completed Storyteller job: {:?}", task);
|
||||
|
||||
let result = maybe_handle_frontend_caller_notification(
|
||||
app_handle,
|
||||
app_env_configs,
|
||||
creds,
|
||||
task,
|
||||
job,
|
||||
).await;
|
||||
|
||||
if let Err(err) = result {
|
||||
error!("Failed to send generation complete event: {:?}", err);
|
||||
}
|
||||
|
||||
//let result = maybe_handle_text_to_image_complete_event(
|
||||
// app_handle,
|
||||
// app_env_configs,
|
||||
// creds,
|
||||
// task,
|
||||
// job,
|
||||
//).await;
|
||||
|
||||
//if let Err(err) = result {
|
||||
// error!("Failed to send text-to-image complete event: {:?}", err);
|
||||
//}
|
||||
|
||||
//let result = maybe_handle_inpainting_complete_event(
|
||||
// app_handle,
|
||||
// app_env_configs,
|
||||
// creds,
|
||||
// task,
|
||||
// job,
|
||||
//).await;
|
||||
|
||||
//if let Err(err) = result {
|
||||
// error!("Failed to send image inpainting complete event: {:?}", err);
|
||||
//}
|
||||
|
||||
//let result = maybe_send_background_removal_complete_event(
|
||||
// app_handle,
|
||||
// task,
|
||||
// job,
|
||||
//).await;
|
||||
|
||||
//if let Err(err) = result {
|
||||
// error!("Failed to send background removal complete event: {:?}", err);
|
||||
//}
|
||||
}
|
||||
|
||||
fn get_thumbnail_template<'a>(job: &'a ListSessionJobsItem) -> Option<&'a str> {
|
||||
let links = match job.maybe_result.as_ref() {
|
||||
None => return None,
|
||||
Some(result) => &result.media_links,
|
||||
};
|
||||
|
||||
media_links_to_thumbnail_template(links)
|
||||
}
|
||||
|
||||
fn get_media_file_class<'a>(job: &'a ListSessionJobsItem) -> Option<TaskMediaFileClass> {
|
||||
match job.request.inference_category {
|
||||
InferenceCategory::BackgroundRemoval => return Some(TaskMediaFileClass::Image),
|
||||
InferenceCategory::ImageGeneration => return Some(TaskMediaFileClass::Image),
|
||||
InferenceCategory::VideoGeneration => return Some(TaskMediaFileClass::Video),
|
||||
InferenceCategory::ObjectGeneration => return Some(TaskMediaFileClass::Dimensional),
|
||||
_ => {}, // Fall-through
|
||||
}
|
||||
|
||||
let result = match job.maybe_result.as_ref() {
|
||||
None => return None,
|
||||
Some(result) => result,
|
||||
};
|
||||
|
||||
let url = result.media_links.cdn_url.as_str();
|
||||
|
||||
if url.ends_with("jpg")
|
||||
|| url.ends_with("jpeg")
|
||||
|| url.ends_with("png")
|
||||
{
|
||||
return Some(TaskMediaFileClass::Image);
|
||||
}
|
||||
|
||||
if url.ends_with("mp4")
|
||||
|| url.ends_with("webm")
|
||||
{
|
||||
return Some(TaskMediaFileClass::Video);
|
||||
}
|
||||
|
||||
if url.ends_with("glb") {
|
||||
return Some(TaskMediaFileClass::Dimensional);
|
||||
}
|
||||
|
||||
if url.ends_with("wav")
|
||||
|| url.ends_with("mp3")
|
||||
{
|
||||
return Some(TaskMediaFileClass::Audio);
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
|
||||
use crate::core::events::generation_events::generation_failed_event::GenerationFailedEvent;
|
||||
use crate::core::state::task_database::TaskDatabase;
|
||||
use crate::core::utils::enum_conversion::generation_provider::to_generation_service_provider;
|
||||
use crate::core::utils::enum_conversion::task_type::to_generation_action;
|
||||
use artcraft_api_defs::jobs::list_session_jobs::ListSessionJobsItem;
|
||||
use enums::tauri::tasks::task_status::TaskStatus;
|
||||
use errors::AnyhowResult;
|
||||
use log::info;
|
||||
use sqlite_tasks::queries::task::Task;
|
||||
use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
|
||||
use tauri::AppHandle;
|
||||
|
||||
pub async fn handle_failed_job(
|
||||
app_handle: &AppHandle,
|
||||
job: &ListSessionJobsItem,
|
||||
task: &Task,
|
||||
task_database: &TaskDatabase,
|
||||
) -> AnyhowResult<()> {
|
||||
info!("Marking storyteller job as failed: {:?}", task.id);
|
||||
|
||||
update_task_status(UpdateTaskArgs {
|
||||
db: task_database.get_connection(),
|
||||
task_id: &task.id,
|
||||
status: TaskStatus::CompleteFailure,
|
||||
}).await?;
|
||||
|
||||
let service = to_generation_service_provider(task.provider);
|
||||
let action = to_generation_action(task.task_type);
|
||||
|
||||
let event = GenerationFailedEvent {
|
||||
action,
|
||||
service,
|
||||
model: None,
|
||||
reason: None,
|
||||
};
|
||||
|
||||
event.send_infallible(app_handle);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
+177
@@ -0,0 +1,177 @@
|
||||
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
|
||||
use crate::core::events::generation_events::common::GenerationAction;
|
||||
use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
|
||||
use crate::core::events::sendable_event_trait::SendableEvent;
|
||||
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
|
||||
use crate::core::state::task_database::TaskDatabase;
|
||||
use crate::core::utils::enum_conversion::generation_provider::to_generation_service_provider;
|
||||
use crate::core::utils::enum_conversion::task_type::to_generation_action;
|
||||
use super::events::maybe_handle_frontend_caller_notification::maybe_handle_frontend_caller_notification;
|
||||
use artcraft_api_defs::jobs::list_session_jobs::ListSessionJobsItem;
|
||||
use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template;
|
||||
use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
|
||||
use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
|
||||
use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass;
|
||||
use errors::AnyhowResult;
|
||||
use log::error;
|
||||
use log::info;
|
||||
use sqlite_tasks::queries::task::Task;
|
||||
use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs};
|
||||
use tauri::AppHandle;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
|
||||
pub async fn handle_successful_job(
|
||||
app_handle: &AppHandle,
|
||||
app_env_configs: &AppEnvConfigs,
|
||||
creds: Option<&StorytellerCredentialSet>,
|
||||
job: &ListSessionJobsItem,
|
||||
task: &Task,
|
||||
task_database: &TaskDatabase,
|
||||
) -> AnyhowResult<()> {
|
||||
let maybe_primary_media_file_token = job.maybe_result
|
||||
.as_ref()
|
||||
.map(|result| MediaFileToken::new_from_str(&result.entity_token));
|
||||
|
||||
let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs {
|
||||
db: task_database.get_connection(),
|
||||
task_id: &task.id,
|
||||
maybe_batch_token: job.maybe_result
|
||||
.as_ref()
|
||||
.map(|result| result.maybe_batch_token.as_ref())
|
||||
.flatten(),
|
||||
maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(),
|
||||
maybe_primary_media_file_class: get_media_file_class(job),
|
||||
maybe_primary_media_file_thumbnail_url_template: get_thumbnail_template(job),
|
||||
maybe_primary_media_file_cdn_url: job.maybe_result
|
||||
.as_ref()
|
||||
.map(|result| result.media_links.cdn_url.as_str()),
|
||||
}).await?;
|
||||
|
||||
if !updated {
|
||||
return Ok(()); // If anything breaks with queries, don't spam events.
|
||||
}
|
||||
|
||||
send_additional_success_events(app_handle, app_env_configs, creds, job, task).await;
|
||||
|
||||
let service = to_generation_service_provider(task.provider);
|
||||
let action = to_generation_action(task.task_type);
|
||||
|
||||
let event = GenerationCompleteEvent {
|
||||
action: Some(action),
|
||||
service,
|
||||
model: None, // TODO
|
||||
};
|
||||
|
||||
if let Err(err) = event.send(app_handle) {
|
||||
error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_additional_success_events(
|
||||
app_handle: &AppHandle,
|
||||
app_env_configs: &AppEnvConfigs,
|
||||
creds: Option<&StorytellerCredentialSet>,
|
||||
job: &ListSessionJobsItem,
|
||||
task: &Task
|
||||
) {
|
||||
info!("Attempting to dispatch events for completed Storyteller job: {:?}", task);
|
||||
|
||||
let result = maybe_handle_frontend_caller_notification(
|
||||
app_handle,
|
||||
app_env_configs,
|
||||
creds,
|
||||
task,
|
||||
job,
|
||||
).await;
|
||||
|
||||
if let Err(err) = result {
|
||||
error!("Failed to send generation complete event: {:?}", err);
|
||||
}
|
||||
|
||||
//let result = maybe_handle_text_to_image_complete_event(
|
||||
// app_handle,
|
||||
// app_env_configs,
|
||||
// creds,
|
||||
// task,
|
||||
// job,
|
||||
//).await;
|
||||
|
||||
//if let Err(err) = result {
|
||||
// error!("Failed to send text-to-image complete event: {:?}", err);
|
||||
//}
|
||||
|
||||
//let result = maybe_handle_inpainting_complete_event(
|
||||
// app_handle,
|
||||
// app_env_configs,
|
||||
// creds,
|
||||
// task,
|
||||
// job,
|
||||
//).await;
|
||||
|
||||
//if let Err(err) = result {
|
||||
// error!("Failed to send image inpainting complete event: {:?}", err);
|
||||
//}
|
||||
|
||||
//let result = maybe_send_background_removal_complete_event(
|
||||
// app_handle,
|
||||
// task,
|
||||
// job,
|
||||
//).await;
|
||||
|
||||
//if let Err(err) = result {
|
||||
// error!("Failed to send background removal complete event: {:?}", err);
|
||||
//}
|
||||
}
|
||||
|
||||
fn get_thumbnail_template<'a>(job: &'a ListSessionJobsItem) -> Option<&'a str> {
|
||||
let links = match job.maybe_result.as_ref() {
|
||||
None => return None,
|
||||
Some(result) => &result.media_links,
|
||||
};
|
||||
|
||||
media_links_to_thumbnail_template(links)
|
||||
}
|
||||
|
||||
fn get_media_file_class(job: &ListSessionJobsItem) -> Option<TaskMediaFileClass> {
|
||||
match job.request.inference_category {
|
||||
InferenceCategory::BackgroundRemoval => return Some(TaskMediaFileClass::Image),
|
||||
InferenceCategory::ImageGeneration => return Some(TaskMediaFileClass::Image),
|
||||
InferenceCategory::VideoGeneration => return Some(TaskMediaFileClass::Video),
|
||||
InferenceCategory::ObjectGeneration => return Some(TaskMediaFileClass::Dimensional),
|
||||
_ => {}, // Fall-through
|
||||
}
|
||||
|
||||
let result = match job.maybe_result.as_ref() {
|
||||
None => return None,
|
||||
Some(result) => result,
|
||||
};
|
||||
|
||||
let url = result.media_links.cdn_url.as_str();
|
||||
|
||||
if url.ends_with("jpg")
|
||||
|| url.ends_with("jpeg")
|
||||
|| url.ends_with("png")
|
||||
{
|
||||
return Some(TaskMediaFileClass::Image);
|
||||
}
|
||||
|
||||
if url.ends_with("mp4")
|
||||
|| url.ends_with("webm")
|
||||
{
|
||||
return Some(TaskMediaFileClass::Video);
|
||||
}
|
||||
|
||||
if url.ends_with("glb") {
|
||||
return Some(TaskMediaFileClass::Dimensional);
|
||||
}
|
||||
|
||||
if url.ends_with("wav")
|
||||
|| url.ends_with("mp3")
|
||||
{
|
||||
return Some(TaskMediaFileClass::Audio);
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
+4
@@ -0,0 +1,4 @@
|
||||
mod events;
|
||||
mod handle_storyteller_failed_job;
|
||||
mod handle_storyteller_successful_job;
|
||||
pub mod storyteller_task_polling_thread;
|
||||
+126
@@ -0,0 +1,126 @@
|
||||
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
|
||||
use crate::core::state::task_database::TaskDatabase;
|
||||
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
|
||||
use crate::services::storyteller::threads::storyteller_task_polling_thread::handle_storyteller_failed_job::handle_failed_job;
|
||||
use crate::services::storyteller::threads::storyteller_task_polling_thread::handle_storyteller_successful_job::handle_successful_job;
|
||||
use anyhow::anyhow;
|
||||
use artcraft_client::endpoints::jobs::list_session_jobs::{list_session_jobs, States};
|
||||
use artcraft_client::error::api_error::ApiError;
|
||||
use artcraft_client::error::storyteller_error::StorytellerError;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::common::job_status_plus::JobStatusPlus;
|
||||
use enums::tauri::tasks::task_status::TaskStatus;
|
||||
use errors::AnyhowResult;
|
||||
use log::error;
|
||||
use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs};
|
||||
use sqlite_tasks::queries::task::Task;
|
||||
use std::collections::HashMap;
|
||||
use tauri::AppHandle;
|
||||
|
||||
pub async fn storyteller_task_polling_thread(
|
||||
app_handle: AppHandle,
|
||||
app_env_configs: AppEnvConfigs,
|
||||
task_database: TaskDatabase,
|
||||
storyteller_creds_manager: StorytellerCredentialManager,
|
||||
) -> ! {
|
||||
loop {
|
||||
let res = polling_loop(
|
||||
&app_handle,
|
||||
&app_env_configs,
|
||||
&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,
|
||||
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::<Vec<_>>();
|
||||
|
||||
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::<HashMap<String, _>>();
|
||||
|
||||
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::<HashMap<String, Task>>();
|
||||
|
||||
for job in jobs.iter() {
|
||||
let task = match tasks_by_provider_job_id.get(job.job_token.as_str()) {
|
||||
Some(task) => task,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
match job.status.status {
|
||||
JobStatusPlus::CompleteSuccess => {
|
||||
match task.status {
|
||||
TaskStatus::CompleteSuccess => continue, // NB: We're done with this task.
|
||||
_ => {}
|
||||
}
|
||||
handle_successful_job(app_handle, app_env_configs, creds.as_ref(), job, task, task_database).await?;
|
||||
}
|
||||
JobStatusPlus::CompleteFailure => {
|
||||
match task.status {
|
||||
TaskStatus::CompleteFailure => continue, // NB: We're done with this task.
|
||||
_ => {}
|
||||
}
|
||||
handle_failed_job(app_handle, job, task, task_database).await?;
|
||||
}
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@ use crate::connection::TaskDbConnection;
|
||||
use crate::error::SqliteTasksError;
|
||||
use chrono::{DateTime, Utc};
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::tauri::tasks::task_failure_type::TaskFailureType;
|
||||
use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass;
|
||||
use enums::tauri::tasks::task_model_type::TaskModelType;
|
||||
use enums::tauri::tasks::task_status::TaskStatus;
|
||||
@@ -30,6 +31,8 @@ pub struct TaskItem {
|
||||
pub on_complete_batch_token: Option<BatchGenerationToken>,
|
||||
pub on_complete_primary_media_file_cdn_url: Option<String>,
|
||||
pub on_complete_primary_media_file_thumbnail_url_template: Option<String>,
|
||||
pub on_failure_type: Option<TaskFailureType>,
|
||||
pub on_failure_message: Option<String>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
@@ -56,6 +59,8 @@ pub async fn list_tasks_for_frontend(
|
||||
on_complete_batch_token,
|
||||
on_complete_primary_media_file_cdn_url,
|
||||
on_complete_primary_media_file_thumbnail_url_template,
|
||||
on_failure_type,
|
||||
on_failure_message,
|
||||
created_at as "created_at: DateTime<Utc>",
|
||||
updated_at as "updated_at: DateTime<Utc>",
|
||||
completed_at as "completed_at: DateTime<Utc>"
|
||||
@@ -92,6 +97,10 @@ pub async fn list_tasks_for_frontend(
|
||||
on_complete_batch_token: raw.on_complete_batch_token.map(|t| BatchGenerationToken::new_from_str(&t)),
|
||||
on_complete_primary_media_file_cdn_url: raw.on_complete_primary_media_file_cdn_url,
|
||||
on_complete_primary_media_file_thumbnail_url_template: raw.on_complete_primary_media_file_thumbnail_url_template,
|
||||
on_failure_type: raw.on_failure_type
|
||||
.map(|t| TaskFailureType::from_str(&t))
|
||||
.transpose()?,
|
||||
on_failure_message: raw.on_failure_message,
|
||||
created_at: raw.created_at,
|
||||
updated_at: raw.updated_at,
|
||||
completed_at: raw.completed_at,
|
||||
@@ -118,6 +127,8 @@ struct TaskItemRaw {
|
||||
on_complete_batch_token: Option<String>,
|
||||
on_complete_primary_media_file_cdn_url: Option<String>,
|
||||
on_complete_primary_media_file_thumbnail_url_template: Option<String>,
|
||||
on_failure_type: Option<String>,
|
||||
on_failure_message: Option<String>,
|
||||
created_at: DateTime<Utc>,
|
||||
updated_at: DateTime<Utc>,
|
||||
completed_at: Option<DateTime<Utc>>,
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod task_failure_type;
|
||||
pub mod task_media_file_class;
|
||||
pub mod task_model_type;
|
||||
pub mod task_status;
|
||||
|
||||
@@ -0,0 +1,195 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use crate::by_table::generic_inference_jobs::frontend_failure_category::FrontendFailureCategory;
|
||||
use crate::error::enum_error::EnumError;
|
||||
#[cfg(test)]
|
||||
use strum::EnumCount;
|
||||
#[cfg(test)]
|
||||
use strum::EnumIter;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
/// Failure type for tasks in the Tauri desktop app.
|
||||
///
|
||||
/// Mirrors the relevant variants from `FrontendFailureCategory` so the desktop
|
||||
/// client can display localized failure information without depending / breaking on the
|
||||
/// server-side enum directly.
|
||||
#[cfg_attr(test, derive(EnumIter, EnumCount))]
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd, Serialize, Deserialize, ToSchema)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TaskFailureType {
|
||||
/// Catch-all for unknown failures.
|
||||
Unknown,
|
||||
|
||||
RuleBansUserImage,
|
||||
RuleBansUserImageWithFaces,
|
||||
RuleBansUserTextPrompt,
|
||||
RuleBansUserContent,
|
||||
|
||||
RuleBansGeneratedVideo,
|
||||
RuleBansGeneratedAudio,
|
||||
RuleBansGeneratedContent,
|
||||
|
||||
/// No reason given for generation failure, but this matches what we were told.
|
||||
GenerationFailed,
|
||||
}
|
||||
|
||||
impl_enum_display_and_debug_using_to_str!(TaskFailureType);
|
||||
|
||||
impl TaskFailureType {
|
||||
|
||||
/// Convert the web API's `FrontendFailureCategory` to a Tauri-facing type, if there is a matching variant.
|
||||
/// If there isn't a matching variant, return `Unknown`.
|
||||
pub fn from_frontend_failure_category(category: FrontendFailureCategory) -> Self {
|
||||
match category {
|
||||
FrontendFailureCategory::ModelRulesViolation => Self::RuleBansUserContent, // NB: This is a legacy enum value.
|
||||
FrontendFailureCategory::RuleBansUserImage => Self::RuleBansUserImage,
|
||||
FrontendFailureCategory::RuleBansUserImageWithFaces => Self::RuleBansUserImageWithFaces,
|
||||
FrontendFailureCategory::RuleBansUserTextPrompt => Self::RuleBansUserTextPrompt,
|
||||
FrontendFailureCategory::RuleBansUserContent => Self::RuleBansUserContent,
|
||||
FrontendFailureCategory::RuleBansGeneratedVideo => Self::RuleBansGeneratedVideo,
|
||||
FrontendFailureCategory::RuleBansGeneratedAudio => Self::RuleBansGeneratedAudio,
|
||||
FrontendFailureCategory::RuleBansGeneratedContent => Self::RuleBansGeneratedContent,
|
||||
FrontendFailureCategory::GenerationFailed => Self::GenerationFailed,
|
||||
_ => Self::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn to_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Unknown => "unknown",
|
||||
Self::RuleBansUserImage => "rule_bans_user_image",
|
||||
Self::RuleBansUserImageWithFaces => "rule_bans_user_image_with_faces",
|
||||
Self::RuleBansUserTextPrompt => "rule_bans_user_text_prompt",
|
||||
Self::RuleBansUserContent => "rule_bans_user_content",
|
||||
Self::RuleBansGeneratedVideo => "rule_bans_generated_video",
|
||||
Self::RuleBansGeneratedAudio => "rule_bans_generated_audio",
|
||||
Self::RuleBansGeneratedContent => "rule_bans_generated_content",
|
||||
Self::GenerationFailed => "generation_failed",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_str(value: &str) -> Result<Self, EnumError> {
|
||||
match value {
|
||||
"unknown" => Ok(Self::Unknown),
|
||||
"rule_bans_user_image" => Ok(Self::RuleBansUserImage),
|
||||
"rule_bans_user_image_with_faces" => Ok(Self::RuleBansUserImageWithFaces),
|
||||
"rule_bans_user_text_prompt" => Ok(Self::RuleBansUserTextPrompt),
|
||||
"rule_bans_user_content" => Ok(Self::RuleBansUserContent),
|
||||
"rule_bans_generated_video" => Ok(Self::RuleBansGeneratedVideo),
|
||||
"rule_bans_generated_audio" => Ok(Self::RuleBansGeneratedAudio),
|
||||
"rule_bans_generated_content" => Ok(Self::RuleBansGeneratedContent),
|
||||
"generation_failed" => Ok(Self::GenerationFailed),
|
||||
_ => Err(EnumError::CouldNotConvertFromString(value.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn all_variants() -> BTreeSet<Self> {
|
||||
BTreeSet::from([
|
||||
Self::Unknown,
|
||||
Self::RuleBansUserImage,
|
||||
Self::RuleBansUserImageWithFaces,
|
||||
Self::RuleBansUserTextPrompt,
|
||||
Self::RuleBansUserContent,
|
||||
Self::RuleBansGeneratedVideo,
|
||||
Self::RuleBansGeneratedAudio,
|
||||
Self::RuleBansGeneratedContent,
|
||||
Self::GenerationFailed,
|
||||
])
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::tauri::tasks::task_failure_type::TaskFailureType;
|
||||
use crate::test_helpers::assert_serialization;
|
||||
|
||||
mod explicit_checks {
|
||||
use super::*;
|
||||
use crate::error::enum_error::EnumError;
|
||||
|
||||
#[test]
|
||||
fn test_serialization() {
|
||||
assert_serialization(TaskFailureType::Unknown, "unknown");
|
||||
assert_serialization(TaskFailureType::RuleBansUserImage, "rule_bans_user_image");
|
||||
assert_serialization(TaskFailureType::RuleBansUserImageWithFaces, "rule_bans_user_image_with_faces");
|
||||
assert_serialization(TaskFailureType::RuleBansUserTextPrompt, "rule_bans_user_text_prompt");
|
||||
assert_serialization(TaskFailureType::RuleBansUserContent, "rule_bans_user_content");
|
||||
assert_serialization(TaskFailureType::RuleBansGeneratedVideo, "rule_bans_generated_video");
|
||||
assert_serialization(TaskFailureType::RuleBansGeneratedAudio, "rule_bans_generated_audio");
|
||||
assert_serialization(TaskFailureType::RuleBansGeneratedContent, "rule_bans_generated_content");
|
||||
assert_serialization(TaskFailureType::GenerationFailed, "generation_failed");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn to_str() {
|
||||
assert_eq!(TaskFailureType::Unknown.to_str(), "unknown");
|
||||
assert_eq!(TaskFailureType::RuleBansUserImage.to_str(), "rule_bans_user_image");
|
||||
assert_eq!(TaskFailureType::RuleBansUserImageWithFaces.to_str(), "rule_bans_user_image_with_faces");
|
||||
assert_eq!(TaskFailureType::RuleBansUserTextPrompt.to_str(), "rule_bans_user_text_prompt");
|
||||
assert_eq!(TaskFailureType::RuleBansUserContent.to_str(), "rule_bans_user_content");
|
||||
assert_eq!(TaskFailureType::RuleBansGeneratedVideo.to_str(), "rule_bans_generated_video");
|
||||
assert_eq!(TaskFailureType::RuleBansGeneratedAudio.to_str(), "rule_bans_generated_audio");
|
||||
assert_eq!(TaskFailureType::RuleBansGeneratedContent.to_str(), "rule_bans_generated_content");
|
||||
assert_eq!(TaskFailureType::GenerationFailed.to_str(), "generation_failed");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_str() {
|
||||
assert_eq!(TaskFailureType::from_str("unknown").unwrap(), TaskFailureType::Unknown);
|
||||
assert_eq!(TaskFailureType::from_str("rule_bans_user_image").unwrap(), TaskFailureType::RuleBansUserImage);
|
||||
assert_eq!(TaskFailureType::from_str("rule_bans_user_image_with_faces").unwrap(), TaskFailureType::RuleBansUserImageWithFaces);
|
||||
assert_eq!(TaskFailureType::from_str("rule_bans_user_text_prompt").unwrap(), TaskFailureType::RuleBansUserTextPrompt);
|
||||
assert_eq!(TaskFailureType::from_str("rule_bans_user_content").unwrap(), TaskFailureType::RuleBansUserContent);
|
||||
assert_eq!(TaskFailureType::from_str("rule_bans_generated_video").unwrap(), TaskFailureType::RuleBansGeneratedVideo);
|
||||
assert_eq!(TaskFailureType::from_str("rule_bans_generated_audio").unwrap(), TaskFailureType::RuleBansGeneratedAudio);
|
||||
assert_eq!(TaskFailureType::from_str("rule_bans_generated_content").unwrap(), TaskFailureType::RuleBansGeneratedContent);
|
||||
assert_eq!(TaskFailureType::from_str("generation_failed").unwrap(), TaskFailureType::GenerationFailed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_str_err() {
|
||||
let result = TaskFailureType::from_str("asdf");
|
||||
assert!(result.is_err());
|
||||
if let Err(EnumError::CouldNotConvertFromString(value)) = result {
|
||||
assert_eq!(value, "asdf");
|
||||
} else {
|
||||
panic!("Expected EnumError::CouldNotConvertFromString");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_variants() {
|
||||
let mut variants = TaskFailureType::all_variants();
|
||||
assert_eq!(variants.len(), 9);
|
||||
assert_eq!(variants.pop_first(), Some(TaskFailureType::Unknown));
|
||||
assert_eq!(variants.pop_first(), Some(TaskFailureType::RuleBansUserImage));
|
||||
assert_eq!(variants.pop_first(), Some(TaskFailureType::RuleBansUserImageWithFaces));
|
||||
assert_eq!(variants.pop_first(), Some(TaskFailureType::RuleBansUserTextPrompt));
|
||||
assert_eq!(variants.pop_first(), Some(TaskFailureType::RuleBansUserContent));
|
||||
assert_eq!(variants.pop_first(), Some(TaskFailureType::RuleBansGeneratedVideo));
|
||||
assert_eq!(variants.pop_first(), Some(TaskFailureType::RuleBansGeneratedAudio));
|
||||
assert_eq!(variants.pop_first(), Some(TaskFailureType::RuleBansGeneratedContent));
|
||||
assert_eq!(variants.pop_first(), Some(TaskFailureType::GenerationFailed));
|
||||
assert_eq!(variants.pop_first(), None);
|
||||
}
|
||||
}
|
||||
|
||||
mod mechanical_checks {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn variant_length() {
|
||||
use strum::IntoEnumIterator;
|
||||
assert_eq!(TaskFailureType::all_variants().len(), TaskFailureType::iter().len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round_trip() {
|
||||
for variant in TaskFailureType::all_variants() {
|
||||
assert_eq!(variant, TaskFailureType::from_str(variant.to_str()).unwrap());
|
||||
assert_eq!(variant, TaskFailureType::from_str(&format!("{}", variant)).unwrap());
|
||||
assert_eq!(variant, TaskFailureType::from_str(&format!("{:?}", variant)).unwrap());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user