mirror of
https://github.com/storytold/artcraft.git
synced 2026-10-21 09:22:41 -04:00
Compare commits
8
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
de0776009b | ||
|
|
11f8885bf6 | ||
|
|
0e2926b179 | ||
|
|
671e02edcb | ||
|
|
10f8649089 | ||
|
|
3aeea6d2b2 | ||
|
|
abccceb5bd | ||
|
|
9add7ee1cd |
+35
-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 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 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": [
|
||||
{
|
||||
@@ -49,18 +49,43 @@
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"name": "created_at: DateTime<Utc>",
|
||||
"name": "on_complete_primary_media_file_token",
|
||||
"ordinal": 9,
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"name": "on_complete_primary_media_file_class",
|
||||
"ordinal": 10,
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"name": "on_complete_batch_token",
|
||||
"ordinal": 11,
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"name": "on_complete_primary_media_file_cdn_url",
|
||||
"ordinal": 12,
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"name": "on_complete_primary_media_file_thumbnail_url_template",
|
||||
"ordinal": 13,
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"name": "created_at: DateTime<Utc>",
|
||||
"ordinal": 14,
|
||||
"type_info": "Int64"
|
||||
},
|
||||
{
|
||||
"name": "updated_at: DateTime<Utc>",
|
||||
"ordinal": 10,
|
||||
"ordinal": 15,
|
||||
"type_info": "Int64"
|
||||
},
|
||||
{
|
||||
"name": "completed_at: DateTime<Utc>",
|
||||
"ordinal": 11,
|
||||
"ordinal": 16,
|
||||
"type_info": "Int64"
|
||||
}
|
||||
],
|
||||
@@ -77,10 +102,15 @@
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
]
|
||||
},
|
||||
"hash": "a142dc3fdf5b6e741c268d8f2e9996a28c214f5835681297d58a9431da8de939"
|
||||
"hash": "47f94c127abc20683b25c257aed3c2190b9bd29eb829d61392f08e10fd986767"
|
||||
}
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
{
|
||||
"db_name": "SQLite",
|
||||
"query": "\n UPDATE tasks\n SET\n task_status = ?,\n on_complete_batch_token = ?,\n on_complete_primary_media_file_token = ?,\n on_complete_primary_media_file_class = ?,\n on_complete_primary_media_file_cdn_url = ?,\n on_complete_primary_media_file_thumbnail_url_template = ?\n WHERE id = ?\n ",
|
||||
"describe": {
|
||||
"columns": [],
|
||||
"parameters": {
|
||||
"Right": 7
|
||||
},
|
||||
"nullable": []
|
||||
},
|
||||
"hash": "50a2449b87abed7df3302a74ecd0e8ab87940077a247636680cb09b2fe488321"
|
||||
}
|
||||
@@ -10,6 +10,13 @@
|
||||
-- tasks_v2.sqlite - added model_type (nullable)
|
||||
-- tasks_v3.sqlite - added frontend_caller, comments
|
||||
-- tasks_v4.sqlite - added is_dismissed_by_user
|
||||
-- tasks_v5.sqlite - added five fields:
|
||||
-- on_complete_batch_token,
|
||||
-- on_complete_primary_media_file_token,
|
||||
-- on_complete_primary_media_file_class,
|
||||
-- on_complete_primary_media_file_cdn_url,
|
||||
-- on_complete_primary_media_file_thumbnail_url_template,
|
||||
--
|
||||
|
||||
CREATE TABLE tasks (
|
||||
-- Task auto-incrementing primary key.
|
||||
@@ -54,6 +61,29 @@ CREATE TABLE tasks (
|
||||
-- Whether the user has dismissed the task from view.
|
||||
is_dismissed_by_user INTEGER NOT NULL DEFAULT 0,
|
||||
|
||||
-- OPTIONAL.
|
||||
-- When the generation completes, the batch token (if any).
|
||||
-- This will point to a collection of multiple files.
|
||||
on_complete_batch_token TEXT,
|
||||
|
||||
-- OPTIONAL.
|
||||
-- When the generation completes, the media token of the "primary" media file.
|
||||
-- We'll probably always generate homogenous types, but just in case this is the first or primary item in a batch.
|
||||
on_complete_primary_media_file_token TEXT,
|
||||
|
||||
-- OPTIONAL.
|
||||
-- When the generation completes, the file type of the primary media file.
|
||||
-- We'll probably always generate homogenous types, but just in case this is the first or primary item in a batch.
|
||||
on_complete_primary_media_file_class TEXT,
|
||||
|
||||
-- OPTIONAL.
|
||||
-- The URL to the media file.
|
||||
on_complete_primary_media_file_cdn_url TEXT,
|
||||
|
||||
-- OPTIONAL.
|
||||
-- The URL template for the thumbnail of the primary media file.
|
||||
on_complete_primary_media_file_thumbnail_url_template 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,8 @@ use serde_derive::{Deserialize, Serialize};
|
||||
use sqlite_tasks::queries::list_tasks_for_frontend::list_tasks_for_frontend;
|
||||
use storyteller_client::endpoints::media_files::delete_media_file::delete_media_file;
|
||||
use tauri::{AppHandle, State};
|
||||
use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass;
|
||||
use tokens::tokens::batch_generations::BatchGenerationToken;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
use tokens::tokens::sqlite::tasks::TaskId;
|
||||
|
||||
@@ -34,11 +36,33 @@ pub struct TaskQueueItem {
|
||||
pub model_type: Option<TaskModelType>,
|
||||
pub provider: Option<GenerationProvider>,
|
||||
pub provider_job_id: Option<String>,
|
||||
pub completed_item: Option<CompletedItemData>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct CompletedItemData {
|
||||
pub primary_media_file: MediaFileData,
|
||||
|
||||
/// The type of file(s) generated.
|
||||
pub media_file_class: Option<TaskMediaFileClass>,
|
||||
|
||||
/// If generated in a batch, we probably have a batch token we can query.
|
||||
pub maybe_batch_token: Option<BatchGenerationToken>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct MediaFileData {
|
||||
pub token: MediaFileToken,
|
||||
pub cdn_url: String,
|
||||
pub maybe_thumbnail_url_template: Option<String>,
|
||||
|
||||
// NB: The frontend wants this.
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl SerializeMarker for GetTaskQueueCommandResponse {}
|
||||
|
||||
#[tauri::command]
|
||||
@@ -74,17 +98,45 @@ pub async fn handle_request(
|
||||
let tasks = list_tasks_for_frontend(task_database.get_connection())
|
||||
.await?;
|
||||
|
||||
Ok(tasks.tasks.into_iter()
|
||||
.map(|task| TaskQueueItem {
|
||||
id: task.id,
|
||||
task_status: task.status,
|
||||
task_type: task.task_type,
|
||||
model_type: task.model_type,
|
||||
provider: task.provider,
|
||||
provider_job_id: task.provider_job_id,
|
||||
created_at: task.created_at,
|
||||
updated_at: task.updated_at,
|
||||
completed_at: task.completed_at,
|
||||
})
|
||||
.collect())
|
||||
let mut transformed_tasks = Vec::with_capacity(tasks.tasks.len());
|
||||
|
||||
for task in tasks.tasks.into_iter() {
|
||||
let mut completed_item = None;
|
||||
|
||||
if task.status == TaskStatus::CompleteSuccess {
|
||||
let token_and_url = task.on_complete_primary_media_file_token
|
||||
.zip(task.on_complete_primary_media_file_cdn_url);
|
||||
|
||||
if let Some((primary_media_file_token, media_file_url)) = token_and_url{
|
||||
completed_item = Some(CompletedItemData {
|
||||
primary_media_file: MediaFileData {
|
||||
token: primary_media_file_token,
|
||||
cdn_url: media_file_url.clone(),
|
||||
maybe_thumbnail_url_template: task.on_complete_primary_media_file_thumbnail_url_template.clone(),
|
||||
// NB: This isn't the exact completion date. Also, fallback to now if missing.
|
||||
created_at: task.completed_at.unwrap_or_else(Utc::now),
|
||||
},
|
||||
media_file_class: task.on_complete_primary_media_file_class,
|
||||
maybe_batch_token: task.on_complete_batch_token,
|
||||
});
|
||||
} else {
|
||||
warn!("Task {} is marked complete but has no primary media file token or URL.", task.id);
|
||||
}
|
||||
}
|
||||
|
||||
transformed_tasks.push(TaskQueueItem {
|
||||
id: task.id,
|
||||
task_status: task.status,
|
||||
task_type: task.task_type,
|
||||
model_type: task.model_type,
|
||||
provider: task.provider,
|
||||
provider_job_id: task.provider_job_id,
|
||||
created_at: task.created_at,
|
||||
updated_at: task.updated_at,
|
||||
completed_at: task.completed_at,
|
||||
completed_item,
|
||||
})
|
||||
}
|
||||
|
||||
Ok(transformed_tasks)
|
||||
}
|
||||
|
||||
@@ -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 = 4;
|
||||
pub const TASK_DATABASE_VERSION: u32 = 5;
|
||||
|
||||
+78
-4
@@ -17,26 +17,30 @@ use crate::services::storyteller::threads::events::maybe_handle_text_to_image_co
|
||||
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 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, 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 storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
|
||||
use storyteller_client::error::api_error::ApiError;
|
||||
use storyteller_client::error::storyteller_error::StorytellerError;
|
||||
use storyteller_client::endpoints::jobs::list_session_jobs::{list_session_jobs, States};
|
||||
use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file;
|
||||
use storyteller_client::error::api_error::ApiError;
|
||||
use storyteller_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,
|
||||
@@ -134,10 +138,23 @@ async fn polling_loop(
|
||||
_ => {} // Fall-through.
|
||||
}
|
||||
|
||||
let updated = update_task_status(UpdateTaskArgs {
|
||||
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,
|
||||
status: TaskStatus::CompleteSuccess,
|
||||
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 {
|
||||
@@ -212,3 +229,60 @@ async fn send_additional_events(
|
||||
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,
|
||||
};
|
||||
|
||||
// NB: We only populate video previews for video tasks.
|
||||
if let Some(video) = links.maybe_video_previews.as_ref() {
|
||||
return Some(&video.animated_thumbnail_template);
|
||||
}
|
||||
|
||||
// Last option - use the image thumbnail template if it exists.
|
||||
links.maybe_thumbnail_template.as_deref()
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
@@ -2,10 +2,13 @@ use crate::connection::TaskDbConnection;
|
||||
use crate::error::SqliteTasksError;
|
||||
use chrono::{DateTime, Utc};
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass;
|
||||
use enums::tauri::tasks::task_model_type::TaskModelType;
|
||||
use enums::tauri::tasks::task_status::TaskStatus;
|
||||
use enums::tauri::tasks::task_type::TaskType;
|
||||
use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
|
||||
use tokens::tokens::batch_generations::BatchGenerationToken;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
use tokens::tokens::sqlite::tasks::TaskId;
|
||||
|
||||
pub struct TaskList {
|
||||
@@ -22,6 +25,11 @@ pub struct TaskItem {
|
||||
pub frontend_caller: Option<TauriCommandCaller>,
|
||||
pub frontend_subscriber_id: Option<String>,
|
||||
pub frontend_subscriber_payload: Option<String>,
|
||||
pub on_complete_primary_media_file_token: Option<MediaFileToken>,
|
||||
pub on_complete_primary_media_file_class: Option<TaskMediaFileClass>,
|
||||
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 created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
@@ -43,6 +51,11 @@ pub async fn list_tasks_for_frontend(
|
||||
frontend_caller,
|
||||
frontend_subscriber_id,
|
||||
frontend_subscriber_payload,
|
||||
on_complete_primary_media_file_token,
|
||||
on_complete_primary_media_file_class,
|
||||
on_complete_batch_token,
|
||||
on_complete_primary_media_file_cdn_url,
|
||||
on_complete_primary_media_file_thumbnail_url_template,
|
||||
created_at as "created_at: DateTime<Utc>",
|
||||
updated_at as "updated_at: DateTime<Utc>",
|
||||
completed_at as "completed_at: DateTime<Utc>"
|
||||
@@ -72,6 +85,13 @@ pub async fn list_tasks_for_frontend(
|
||||
.transpose()?,
|
||||
frontend_subscriber_id: raw.frontend_subscriber_id,
|
||||
frontend_subscriber_payload: raw.frontend_subscriber_payload,
|
||||
on_complete_primary_media_file_token: raw.on_complete_primary_media_file_token.map(|t| MediaFileToken::new_from_str(&t)),
|
||||
on_complete_primary_media_file_class: raw.on_complete_primary_media_file_class
|
||||
.map(|c| TaskMediaFileClass::from_str(&c))
|
||||
.transpose()?,
|
||||
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,
|
||||
created_at: raw.created_at,
|
||||
updated_at: raw.updated_at,
|
||||
completed_at: raw.completed_at,
|
||||
@@ -93,6 +113,11 @@ struct TaskItemRaw {
|
||||
frontend_caller: Option<String>,
|
||||
frontend_subscriber_id: Option<String>,
|
||||
frontend_subscriber_payload: Option<String>,
|
||||
on_complete_primary_media_file_token: Option<String>,
|
||||
on_complete_primary_media_file_class: Option<String>,
|
||||
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>,
|
||||
created_at: DateTime<Utc>,
|
||||
updated_at: DateTime<Utc>,
|
||||
completed_at: Option<DateTime<Utc>>,
|
||||
|
||||
@@ -3,4 +3,5 @@ pub mod list_tasks_by_provider_and_status;
|
||||
pub mod list_tasks_by_provider_and_tokens;
|
||||
pub mod list_tasks_for_frontend;
|
||||
pub mod mark_task_as_dismissed;
|
||||
pub mod update_successful_task_status_with_metadata;
|
||||
pub mod update_task_status;
|
||||
+59
@@ -0,0 +1,59 @@
|
||||
use crate::connection::TaskDbConnection;
|
||||
use crate::error::SqliteTasksError;
|
||||
use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass;
|
||||
use enums::tauri::tasks::task_status::TaskStatus;
|
||||
use tokens::tokens::batch_generations::BatchGenerationToken;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
use tokens::tokens::sqlite::tasks::TaskId;
|
||||
|
||||
const SUCCESSFUL_STATUS : &str = TaskStatus::CompleteSuccess.to_str();
|
||||
|
||||
pub struct UpdateSuccessfulTaskArgs<'a> {
|
||||
pub db: &'a TaskDbConnection,
|
||||
pub task_id: &'a TaskId,
|
||||
pub maybe_batch_token: Option<&'a BatchGenerationToken>,
|
||||
pub maybe_primary_media_file_token: Option<&'a MediaFileToken>,
|
||||
pub maybe_primary_media_file_class: Option<TaskMediaFileClass>,
|
||||
pub maybe_primary_media_file_cdn_url: Option<&'a str>,
|
||||
pub maybe_primary_media_file_thumbnail_url_template: Option<&'a str>,
|
||||
}
|
||||
|
||||
/// Returns true if rows were updated.
|
||||
pub async fn update_successful_task_status_with_metadata(
|
||||
args: UpdateSuccessfulTaskArgs<'_>,
|
||||
) -> Result<bool, SqliteTasksError> {
|
||||
|
||||
// TODO(bt,2025-07-12): Fix this. The sqlx mysql queries never required temporaries
|
||||
let task_id_temp = args.task_id.as_str();
|
||||
let maybe_batch_token_temp = args.maybe_batch_token.map(|t| t.as_str());
|
||||
let maybe_primary_media_token_temp = args.maybe_primary_media_file_token.map(|t| t.as_str());
|
||||
let maybe_primary_media_class = args.maybe_primary_media_file_class.map(|c| c.to_str());
|
||||
|
||||
// TODO(bt,2025-07-15): We can't set a LIMIT without a certain compiler flag for SQLite ?
|
||||
let query = sqlx::query!(r#"
|
||||
UPDATE tasks
|
||||
SET
|
||||
task_status = ?,
|
||||
on_complete_batch_token = ?,
|
||||
on_complete_primary_media_file_token = ?,
|
||||
on_complete_primary_media_file_class = ?,
|
||||
on_complete_primary_media_file_cdn_url = ?,
|
||||
on_complete_primary_media_file_thumbnail_url_template = ?
|
||||
WHERE id = ?
|
||||
"#,
|
||||
SUCCESSFUL_STATUS,
|
||||
maybe_batch_token_temp,
|
||||
maybe_primary_media_token_temp,
|
||||
maybe_primary_media_class,
|
||||
args.maybe_primary_media_file_cdn_url,
|
||||
args.maybe_primary_media_file_thumbnail_url_template,
|
||||
task_id_temp,
|
||||
);
|
||||
|
||||
// info!("query: {:?}", query.sql());
|
||||
|
||||
let res = query.execute(args.db.get_pool()).await?;
|
||||
let rows_updated = res.rows_affected() > 0;
|
||||
|
||||
Ok(rows_updated)
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod task_media_file_class;
|
||||
pub mod task_model_type;
|
||||
pub mod task_status;
|
||||
pub mod task_type;
|
||||
@@ -0,0 +1,140 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use crate::error::enum_error::EnumError;
|
||||
#[cfg(test)]
|
||||
use strum::EnumCount;
|
||||
#[cfg(test)]
|
||||
use strum::EnumIter;
|
||||
use utoipa::ToSchema;
|
||||
|
||||
#[cfg_attr(test, derive(EnumIter, EnumCount))]
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd, Serialize, Deserialize, ToSchema)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TaskMediaFileClass {
|
||||
/// Audio files: wav, mp3, etc.
|
||||
Audio,
|
||||
|
||||
/// Image files: png, jpeg, etc.
|
||||
Image,
|
||||
|
||||
/// Video files: mp4, etc.
|
||||
Video,
|
||||
|
||||
/// 3D engine data: glb, gltf, etc.
|
||||
Dimensional,
|
||||
}
|
||||
|
||||
impl_enum_display_and_debug_using_to_str!(TaskMediaFileClass);
|
||||
//impl_mysql_enum_coders!(TaskMediaFileType);
|
||||
//impl_mysql_from_row!(TaskMediaFileType);
|
||||
|
||||
// NB: We can derive `sqlx::Type` instead of using `impl_mysql_enum_coders`
|
||||
|
||||
impl TaskMediaFileClass {
|
||||
pub fn to_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Audio => "audio",
|
||||
Self::Image => "image",
|
||||
Self::Video => "video",
|
||||
Self::Dimensional => "dimensional",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_str(value: &str) -> Result<Self, EnumError> {
|
||||
match value {
|
||||
"audio" => Ok(Self::Audio),
|
||||
"image" => Ok(Self::Image),
|
||||
"video" => Ok(Self::Video),
|
||||
"dimensional" => Ok(Self::Dimensional),
|
||||
_ => Err(EnumError::CouldNotConvertFromString(value.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn all_variants() -> BTreeSet<Self> {
|
||||
// NB: BTreeSet is sorted
|
||||
// NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc.
|
||||
BTreeSet::from([
|
||||
Self::Audio,
|
||||
Self::Image,
|
||||
Self::Video,
|
||||
Self::Dimensional,
|
||||
])
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::tauri::tasks::task_media_file_class::TaskMediaFileClass;
|
||||
use crate::test_helpers::assert_serialization;
|
||||
|
||||
mod explicit_checks {
|
||||
use super::*;
|
||||
use crate::error::enum_error::EnumError;
|
||||
|
||||
#[test]
|
||||
fn test_serialization() {
|
||||
assert_serialization(TaskMediaFileClass::Audio, "audio");
|
||||
assert_serialization(TaskMediaFileClass::Image, "image");
|
||||
assert_serialization(TaskMediaFileClass::Video, "video");
|
||||
assert_serialization(TaskMediaFileClass::Dimensional, "dimensional");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn to_str() {
|
||||
assert_eq!(TaskMediaFileClass::Audio.to_str(), "audio");
|
||||
assert_eq!(TaskMediaFileClass::Image.to_str(), "image");
|
||||
assert_eq!(TaskMediaFileClass::Video.to_str(), "video");
|
||||
assert_eq!(TaskMediaFileClass::Dimensional.to_str(), "dimensional");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_str() {
|
||||
assert_eq!(TaskMediaFileClass::from_str("audio").unwrap(), TaskMediaFileClass::Audio);
|
||||
assert_eq!(TaskMediaFileClass::from_str("image").unwrap(), TaskMediaFileClass::Image);
|
||||
assert_eq!(TaskMediaFileClass::from_str("video").unwrap(), TaskMediaFileClass::Video);
|
||||
assert_eq!(TaskMediaFileClass::from_str("dimensional").unwrap(), TaskMediaFileClass::Dimensional);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_str_err() {
|
||||
let result = TaskMediaFileClass::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 = TaskMediaFileClass::all_variants();
|
||||
assert_eq!(variants.len(), 4);
|
||||
assert_eq!(variants.pop_first(), Some(TaskMediaFileClass::Audio));
|
||||
assert_eq!(variants.pop_first(), Some(TaskMediaFileClass::Image));
|
||||
assert_eq!(variants.pop_first(), Some(TaskMediaFileClass::Video));
|
||||
assert_eq!(variants.pop_first(), Some(TaskMediaFileClass::Dimensional));
|
||||
assert_eq!(variants.pop_first(), None);
|
||||
}
|
||||
}
|
||||
|
||||
mod mechanical_checks {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn variant_length() {
|
||||
use strum::IntoEnumIterator;
|
||||
assert_eq!(TaskMediaFileClass::all_variants().len(), TaskMediaFileClass::iter().len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round_trip() {
|
||||
for variant in TaskMediaFileClass::all_variants() {
|
||||
// Test to_str(), from_str(), Display, and Debug.
|
||||
assert_eq!(variant, TaskMediaFileClass::from_str(variant.to_str()).unwrap());
|
||||
assert_eq!(variant, TaskMediaFileClass::from_str(&format!("{}", variant)).unwrap());
|
||||
assert_eq!(variant, TaskMediaFileClass::from_str(&format!("{:?}", variant)).unwrap());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -35,7 +35,7 @@ impl Default for TaskStatus {
|
||||
}
|
||||
|
||||
impl TaskStatus {
|
||||
pub fn to_str(&self) -> &'static str {
|
||||
pub const fn to_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Pending => "pending",
|
||||
Self::Started => "started",
|
||||
|
||||
@@ -29,16 +29,16 @@ import {
|
||||
getProviderDisplayName,
|
||||
getModelDisplayName,
|
||||
} from "@storyteller/model-list";
|
||||
import { MediaFilesApi } from "@storyteller/api";
|
||||
import { CloseButton } from "@storyteller/ui-close-button";
|
||||
import { ActionReminderModal } from "@storyteller/ui-action-reminder-modal";
|
||||
import { TaskMediaFileClass } from "@storyteller/api-enums";
|
||||
|
||||
type InProgressTask = {
|
||||
id: string;
|
||||
title: string;
|
||||
subtitle?: string;
|
||||
progress: number;
|
||||
updatedAt?: string;
|
||||
updatedAt?: Date;
|
||||
canDismiss?: boolean;
|
||||
};
|
||||
|
||||
@@ -47,10 +47,11 @@ type CompletedTask = {
|
||||
title: string;
|
||||
subtitle?: string;
|
||||
thumbnailUrl?: string;
|
||||
completedAt?: string;
|
||||
updatedAt?: string;
|
||||
completedAt?: Date;
|
||||
updatedAt?: Date;
|
||||
imageUrls?: string[];
|
||||
mediaTokens?: string[];
|
||||
mediaFileClass?: TaskMediaFileClass;
|
||||
};
|
||||
|
||||
const InProgressCard = ({
|
||||
@@ -129,6 +130,16 @@ const CompletedCard = ({
|
||||
<img
|
||||
src={task.thumbnailUrl}
|
||||
alt={task.title}
|
||||
onError={(e) => {
|
||||
console.log("Failed to load thumbnail", e);
|
||||
let errorPlaceholder = "/resources/placeholders/placeholder.png";
|
||||
switch (task.mediaFileClass) {
|
||||
case TaskMediaFileClass.Video:
|
||||
errorPlaceholder = "/resources/placeholders/placeholder_play.png";
|
||||
break;
|
||||
}
|
||||
e.currentTarget.src = errorPlaceholder;
|
||||
}}
|
||||
className="h-full w-full object-cover"
|
||||
/>
|
||||
) : (
|
||||
@@ -148,7 +159,7 @@ const CompletedCard = ({
|
||||
)}
|
||||
{task.completedAt && (
|
||||
<div className="text-base-fg text-xs opacity-60">
|
||||
{task.completedAt}
|
||||
{task.completedAt.toISOString()}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
@@ -219,43 +230,11 @@ export const TaskQueue = () => {
|
||||
const load = async () => {
|
||||
try {
|
||||
const result = await GetTaskQueue();
|
||||
|
||||
if (cancelled) return;
|
||||
console.log("TaskQueue:GetTaskQueue result", result);
|
||||
const { tasks } = result;
|
||||
|
||||
// Fetch media files for thumbnails, using dummy tokens from API for now
|
||||
const tokens = result.results as string[] | undefined;
|
||||
let imageUrls: string[] = [];
|
||||
let mediaTokens: string[] | undefined = undefined;
|
||||
if (tokens && tokens.length > 0) {
|
||||
try {
|
||||
const api = new MediaFilesApi();
|
||||
const batchToken = tokens.find((t) => t.startsWith("batch_"));
|
||||
if (batchToken) {
|
||||
const resp = await api.GetMediaFilesByBatchToken({
|
||||
batchToken,
|
||||
});
|
||||
if (resp.success && resp.data) {
|
||||
imageUrls = resp.data
|
||||
.filter((m) => !!m?.media_links?.cdn_url)
|
||||
.map((m) => m.media_links.cdn_url);
|
||||
mediaTokens = resp.data.map((m) => m.token);
|
||||
}
|
||||
} else {
|
||||
const resp = await api.ListMediaFilesByTokens({
|
||||
mediaTokens: tokens,
|
||||
});
|
||||
if (resp.success && resp.data) {
|
||||
imageUrls = resp.data
|
||||
.filter((m) => !!m?.media_links?.cdn_url)
|
||||
.map((m) => m.media_links.cdn_url);
|
||||
mediaTokens = resp.data.map((m) => m.token);
|
||||
}
|
||||
}
|
||||
} catch (_) {
|
||||
// ignore
|
||||
}
|
||||
}
|
||||
const { tasks } = result;
|
||||
|
||||
const now = Date.now();
|
||||
const inProg = tasks
|
||||
@@ -286,7 +265,7 @@ export const TaskQueue = () => {
|
||||
title: `Generating ${parts.kind || "Task"}...`,
|
||||
subtitle: parts.subtitle,
|
||||
progress,
|
||||
updatedAt: t.updated_at?.toISOString(),
|
||||
updatedAt: t.updated_at,
|
||||
canDismiss,
|
||||
};
|
||||
});
|
||||
@@ -309,11 +288,17 @@ export const TaskQueue = () => {
|
||||
.map((t: TaskQueueItem) => ({
|
||||
id: t.id,
|
||||
...formatTitleParts(t),
|
||||
thumbnailUrl: imageUrls?.[0],
|
||||
completedAt: t.completed_at?.toISOString(),
|
||||
updatedAt: t.updated_at?.toISOString(),
|
||||
imageUrls: imageUrls?.length ? imageUrls : undefined,
|
||||
mediaTokens: mediaTokens?.length ? mediaTokens : undefined,
|
||||
thumbnailUrl: t.completed_item?.primary_media_file?.maybe_thumbnail_url_template ?
|
||||
t.completed_item?.primary_media_file?.maybe_thumbnail_url_template.replace("{WIDTH}", "250") :
|
||||
undefined,
|
||||
imageUrls: t.completed_item?.primary_media_file?.cdn_url ?
|
||||
[t.completed_item?.primary_media_file?.cdn_url] :
|
||||
[],
|
||||
mediaTokens: t.completed_item?.primary_media_file?.token ?
|
||||
[t.completed_item?.primary_media_file?.token] :
|
||||
[],
|
||||
completedAt: t.completed_at,
|
||||
updatedAt: t.updated_at,
|
||||
}));
|
||||
|
||||
setInProgress(inProg);
|
||||
@@ -478,8 +463,8 @@ export const TaskQueue = () => {
|
||||
id: firstMediaToken,
|
||||
label: t.title,
|
||||
thumbnail: t.thumbnailUrl || null,
|
||||
fullImage: t.thumbnailUrl || null,
|
||||
createdAt: new Date().toISOString(),
|
||||
fullImage: t.imageUrls?.[0] || null,
|
||||
createdAt: (t.completedAt || new Date()).toISOString(),
|
||||
mediaClass: "image",
|
||||
} as GalleryItem;
|
||||
galleryModalLightboxMediaId.value = item.id;
|
||||
@@ -600,9 +585,9 @@ export const TaskQueue = () => {
|
||||
id: t.id,
|
||||
label: t.title,
|
||||
thumbnail: t.thumbnailUrl || null,
|
||||
fullImage: t.thumbnailUrl || null,
|
||||
createdAt: new Date().toISOString(),
|
||||
mediaClass: "image",
|
||||
fullImage: t.imageUrls?.[0] || null,
|
||||
createdAt: (t.completedAt || new Date()).toISOString(),
|
||||
mediaClass: t.mediaFileClass,
|
||||
} as GalleryItem;
|
||||
galleryModalLightboxMediaId.value = item.id;
|
||||
galleryModalLightboxImage.value = item as GalleryItem;
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
export * from "./task_media_file_class.js";
|
||||
export * from "./task_model_type.js";
|
||||
export * from "./task_status.js";
|
||||
export * from "./task_type.js";
|
||||
@@ -0,0 +1,10 @@
|
||||
|
||||
// NOTE: These are defined in Rust (as the source of truth) and duplicated in the frontend.
|
||||
// In the future, we should use code gen (protobufs or similar) to keep the two sides in sync.
|
||||
|
||||
export enum TaskMediaFileClass {
|
||||
Audio = "audio",
|
||||
Image = "image",
|
||||
Video = "video",
|
||||
Dimensional = "dimensional",
|
||||
}
|
||||
@@ -26,7 +26,7 @@ interface GalleryDraggableItemProps {
|
||||
activeFilter: string;
|
||||
selected: boolean;
|
||||
onClick: () => void;
|
||||
onImageError: () => void;
|
||||
onImageError: (e: React.SyntheticEvent<HTMLImageElement>) => void;
|
||||
disableTooltipAndBadge?: boolean;
|
||||
imageFit?: "cover" | "contain";
|
||||
onDeleted?: (id: string) => void;
|
||||
|
||||
@@ -637,9 +637,11 @@ export const GalleryModal = React.memo(
|
||||
activeFilter={activeFilter}
|
||||
selected={selectedItemIds.includes(item.id)}
|
||||
onClick={() => handleItemClick(item)}
|
||||
onImageError={() =>
|
||||
handleImageError(item.thumbnail!)
|
||||
}
|
||||
onImageError={(e) => {
|
||||
// NB: Replace the broken thumbnail with a placeholder.
|
||||
e.currentTarget.src = "/resources/placeholders/placeholder.png";
|
||||
handleImageError(item.thumbnail!);
|
||||
}}
|
||||
disableTooltipAndBadge={mode === "select"}
|
||||
imageFit={imageFit}
|
||||
onDeleted={handleItemDeleted}
|
||||
|
||||
@@ -4,12 +4,12 @@ import {
|
||||
TaskStatus,
|
||||
TaskType,
|
||||
TaskModelType,
|
||||
TaskMediaFileClass,
|
||||
GenerationProvider,
|
||||
} from "@storyteller/api-enums";
|
||||
|
||||
interface GetTaskQueueResponse {
|
||||
tasks: TaskQueueItem[];
|
||||
results?: string[];
|
||||
}
|
||||
|
||||
export interface TaskQueueItem {
|
||||
@@ -19,11 +19,25 @@ export interface TaskQueueItem {
|
||||
model_type?: TaskModelType;
|
||||
provider?: GenerationProvider;
|
||||
provider_job_id?: string;
|
||||
completed_item?: TaskQueueCompletedItem;
|
||||
created_at: Date;
|
||||
updated_at: Date;
|
||||
completed_at?: Date;
|
||||
}
|
||||
|
||||
export interface TaskQueueCompletedItem {
|
||||
primary_media_file: MediaFileData,
|
||||
media_file_class?: TaskMediaFileClass,
|
||||
maybe_batch_token?: string,
|
||||
}
|
||||
|
||||
export interface MediaFileData {
|
||||
token: string,
|
||||
cdn_url: string,
|
||||
maybe_thumbnail_url_template?: string,
|
||||
created_at: Date,
|
||||
}
|
||||
|
||||
interface GetTaskQueueSuccess extends CommandResult {
|
||||
payload: GetTaskQueueResponse;
|
||||
}
|
||||
@@ -37,6 +51,12 @@ export const GetTaskQueue = async (): Promise<GetTaskQueueResponse> => {
|
||||
|
||||
// Convert timestamps to Date objects
|
||||
const newTasks: TaskQueueItem[] = tasks.map((task) => {
|
||||
let completed_item = task.completed_item;
|
||||
|
||||
if (!!completed_item) {
|
||||
completed_item.primary_media_file.created_at = new Date(completed_item.primary_media_file.created_at);
|
||||
}
|
||||
|
||||
return {
|
||||
id: task.id,
|
||||
task_status: task.task_status,
|
||||
@@ -44,18 +64,15 @@ export const GetTaskQueue = async (): Promise<GetTaskQueueResponse> => {
|
||||
model_type: task.model_type,
|
||||
provider: task.provider,
|
||||
provider_job_id: task.provider_job_id,
|
||||
completed_item: completed_item,
|
||||
created_at: new Date(task.created_at),
|
||||
updated_at: new Date(task.updated_at),
|
||||
completed_at: task.completed_at ? new Date(task.completed_at) : undefined,
|
||||
};
|
||||
});
|
||||
|
||||
// Dummy tokens until backend provides real ones
|
||||
const dummyTokens: string[] = ["batch_g_4r8m6kx41yxm7yvs0kqapm3e"];
|
||||
|
||||
return {
|
||||
tasks: newTasks,
|
||||
results: dummyTokens,
|
||||
};
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user