| 1 | # Copyright (C) 2025 AIDC-AI |
| 2 | # |
| 3 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | # you may not use this file except in compliance with the License. |
| 5 | # You may obtain a copy of the License at |
| 6 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 7 | # Unless required by applicable law or agreed to in writing, software |
| 8 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 9 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 10 | # See the License for the specific language governing permissions and |
| 11 | # limitations under the License. |
| 12 | |
| 13 | """ |
| 14 | Task data models |
| 15 | """ |
| 16 | |
| 17 | from datetime import datetime |
| 18 | from enum import Enum |
| 19 | from typing import Any, Optional |
| 20 | from pydantic import BaseModel, Field |
| 21 | |
| 22 | |
| 23 | class TaskStatus(str, Enum): |
| 24 | """Task status""" |
| 25 | PENDING = "pending" |
| 26 | RUNNING = "running" |
| 27 | COMPLETED = "completed" |
| 28 | FAILED = "failed" |
| 29 | CANCELLED = "cancelled" |
| 30 | |
| 31 | |
| 32 | class TaskType(str, Enum): |
| 33 | """Task type""" |
| 34 | VIDEO_GENERATION = "video_generation" |
| 35 | |
| 36 | |
| 37 | class TaskProgress(BaseModel): |
| 38 | """Task progress information""" |
| 39 | current: int = 0 |
| 40 | total: int = 0 |
| 41 | percentage: float = 0.0 |
| 42 | message: str = "" |
| 43 | |
| 44 | |
| 45 | class Task(BaseModel): |
| 46 | """Task model""" |
| 47 | task_id: str |
| 48 | task_type: TaskType |
| 49 | status: TaskStatus = TaskStatus.PENDING |
| 50 | |
| 51 | # Progress tracking |
| 52 | progress: Optional[TaskProgress] = None |
| 53 | |
| 54 | # Result |
| 55 | result: Optional[Any] = None |
| 56 | error: Optional[str] = None |
| 57 | |
| 58 | # Metadata |
| 59 | created_at: datetime = Field(default_factory=datetime.now) |
| 60 | started_at: Optional[datetime] = None |
| 61 | completed_at: Optional[datetime] = None |
| 62 | |
| 63 | # Request parameters (for reference) |
| 64 | request_params: Optional[dict] = None |
| 65 | |
| 66 | class Config: |
| 67 | json_encoders = { |
| 68 | datetime: lambda v: v.isoformat() |
| 69 | } |
| 70 | |
| 71 |