use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[serde(transparent)]
pub struct TaskId(pub String);
impl TaskId {
#[must_use]
pub fn new(s: impl Into<String>) -> Self {
Self(s.into())
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for TaskId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
impl From<String> for TaskId {
fn from(s: String) -> Self {
Self(s)
}
}
impl From<&str> for TaskId {
fn from(s: &str) -> Self {
Self(s.to_string())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[serde(rename_all = "lowercase")]
pub enum TaskStatus {
Queued,
Running,
Success,
Failed,
Cancelled,
Unknown,
Banned,
Expired,
}
impl TaskStatus {
#[must_use]
pub fn is_terminal(self) -> bool {
matches!(
self,
Self::Success | Self::Failed | Self::Cancelled | Self::Banned | Self::Expired
)
}
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Balance {
pub balance: f64,
pub frozen: f64,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct UploadedFile {
pub file_token: uuid::Uuid,
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct TaskOutput {
#[serde(default)]
pub model: Option<String>,
#[serde(default)]
pub base_model: Option<String>,
#[serde(default)]
pub pbr_model: Option<String>,
#[serde(default)]
pub rendered_image: Option<String>,
#[serde(default)]
pub riggable: Option<bool>,
#[serde(default)]
pub rig_type: Option<crate::enums::RigTypeResponse>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Task {
pub task_id: TaskId,
#[serde(rename = "type")]
pub task_type: String,
pub status: TaskStatus,
#[serde(default)]
pub input: BTreeMap<String, serde_json::Value>,
#[serde(default)]
pub output: TaskOutput,
#[serde(default)]
pub progress: i32,
#[serde(default)]
pub create_time: i64,
#[serde(default)]
pub running_left_time: Option<i64>,
#[serde(default)]
pub queuing_num: Option<i32>,
#[serde(default)]
pub error_code: Option<i32>,
#[serde(default)]
pub error_msg: Option<String>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn task_status_terminality() {
assert!(TaskStatus::Success.is_terminal());
assert!(TaskStatus::Failed.is_terminal());
assert!(TaskStatus::Banned.is_terminal());
assert!(!TaskStatus::Queued.is_terminal());
assert!(!TaskStatus::Running.is_terminal());
assert!(!TaskStatus::Unknown.is_terminal());
}
#[test]
fn deserializes_task_with_minimal_body() {
let body = r#"{
"task_id":"abc123","type":"text_to_model","status":"running","progress":42
}"#;
let task: Task = serde_json::from_str(body).unwrap();
assert_eq!(task.task_id.as_str(), "abc123");
assert_eq!(task.status, TaskStatus::Running);
assert_eq!(task.progress, 42);
assert!(task.output.model.is_none());
}
}