1use std::collections::BTreeMap;
4
5use serde::{Deserialize, Serialize};
6
7#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
9#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
10#[serde(transparent)]
11pub struct TaskId(pub String);
12
13impl TaskId {
14 #[must_use]
16 pub fn new(s: impl Into<String>) -> Self {
17 Self(s.into())
18 }
19 #[must_use]
21 pub fn as_str(&self) -> &str {
22 &self.0
23 }
24}
25
26impl std::fmt::Display for TaskId {
27 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28 f.write_str(&self.0)
29 }
30}
31
32impl From<String> for TaskId {
33 fn from(s: String) -> Self {
34 Self(s)
35 }
36}
37impl From<&str> for TaskId {
38 fn from(s: &str) -> Self {
39 Self(s.to_string())
40 }
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
45#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
46#[serde(rename_all = "lowercase")]
47pub enum TaskStatus {
48 Queued,
50 Running,
52 Success,
54 Failed,
56 Cancelled,
58 Unknown,
60 Banned,
62 Expired,
64}
65
66impl TaskStatus {
67 #[must_use]
69 pub fn is_terminal(self) -> bool {
70 matches!(
71 self,
72 Self::Success | Self::Failed | Self::Cancelled | Self::Banned | Self::Expired
73 )
74 }
75}
76
77#[derive(Debug, Clone, Deserialize, Serialize)]
79#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
80pub struct Balance {
81 pub balance: f64,
83 pub frozen: f64,
85}
86
87#[derive(Debug, Clone, Deserialize, Serialize)]
89#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
90pub struct UploadedFile {
91 pub file_token: uuid::Uuid,
93}
94
95#[derive(Debug, Clone, Default, Deserialize, Serialize)]
97#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
98pub struct TaskOutput {
99 #[serde(default)]
101 pub model: Option<String>,
102 #[serde(default)]
104 pub base_model: Option<String>,
105 #[serde(default)]
107 pub pbr_model: Option<String>,
108 #[serde(default)]
110 pub rendered_image: Option<String>,
111 #[serde(default)]
113 pub riggable: Option<bool>,
114 #[serde(default)]
116 pub rig_type: Option<crate::enums::RigTypeResponse>,
117}
118
119#[derive(Debug, Clone, Deserialize, Serialize)]
121#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
122pub struct Task {
123 pub task_id: TaskId,
125 #[serde(rename = "type")]
127 pub task_type: String,
128 pub status: TaskStatus,
130 #[serde(default)]
132 pub input: BTreeMap<String, serde_json::Value>,
133 #[serde(default)]
135 pub output: TaskOutput,
136 #[serde(default)]
138 pub progress: i32,
139 #[serde(default)]
141 pub create_time: i64,
142 #[serde(default)]
144 pub running_left_time: Option<i64>,
145 #[serde(default)]
147 pub queuing_num: Option<i32>,
148 #[serde(default)]
150 pub error_code: Option<i32>,
151 #[serde(default)]
153 pub error_msg: Option<String>,
154}
155
156#[cfg(test)]
157mod tests {
158 use super::*;
159
160 #[test]
161 fn task_status_terminality() {
162 assert!(TaskStatus::Success.is_terminal());
163 assert!(TaskStatus::Failed.is_terminal());
164 assert!(TaskStatus::Banned.is_terminal());
165 assert!(!TaskStatus::Queued.is_terminal());
166 assert!(!TaskStatus::Running.is_terminal());
167 assert!(!TaskStatus::Unknown.is_terminal());
168 }
169
170 #[test]
171 fn deserializes_task_with_minimal_body() {
172 let body = r#"{
173 "task_id":"abc123","type":"text_to_model","status":"running","progress":42
174 }"#;
175 let task: Task = serde_json::from_str(body).unwrap();
176 assert_eq!(task.task_id.as_str(), "abc123");
177 assert_eq!(task.status, TaskStatus::Running);
178 assert_eq!(task.progress, 42);
179 assert!(task.output.model.is_none());
180 }
181}