76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
from typing import Optional
|
||
from typing import Dict, Any, Union
|
||
from datetime import datetime
|
||
|
||
from pydantic import BaseModel, Field, ConfigDict
|
||
from app.models.enums import TaskStatus, TaskType
|
||
|
||
class DecodeTaskParams(BaseModel):
|
||
filename: str = Field(..., description="文件名")
|
||
sheet_name: str = Field(..., description="工作表名称")
|
||
decode_type: str = Field(..., description="解码类型")
|
||
header: str = Field(..., description="解析列")
|
||
model: Optional[str] = Field(None, description="模型名称")
|
||
|
||
class DecodeTaskResult(BaseModel):
|
||
filename: str = Field(..., description="结果保存文件名")
|
||
spend: float = Field(..., description="任务耗时")
|
||
rows: int = Field(..., description="解析行数")
|
||
|
||
class TaskResponse(BaseModel):
|
||
id: int
|
||
name: str
|
||
type: TaskType # ← Pydantic 会自动转为字符串
|
||
params: dict
|
||
status: TaskStatus # ← 同上
|
||
result: Optional[dict]
|
||
|
||
class Config:
|
||
from_attributes=True
|
||
use_enum_values = True # ← 关键！将 Enum 转为值（如 'pending'）
|
||
# 或者在 Pydantic v2 中：
|
||
# model_config = ConfigDict(use_enum_values=True)
|
||
|
||
class TaskModel(BaseModel):
|
||
name: str
|
||
type: TaskType # Enum: e.g., "decode", "upload", etc.
|
||
params: Dict[str, Any] # 存储原始参数字典
|
||
status: TaskStatus
|
||
result: Optional[Dict[str, Any]] = None
|
||
|
||
@classmethod
|
||
def create_decode_task(
|
||
cls,
|
||
name: str,
|
||
params: DecodeTaskParams,
|
||
status: TaskStatus = TaskStatus.PENDING
|
||
) -> "TaskModel":
|
||
"""创建解码任务的工厂方法"""
|
||
return cls(
|
||
name=name,
|
||
type=TaskType.DECODE_ORDER,
|
||
params=params.dict(),
|
||
status=status
|
||
)
|
||
|
||
def get_params(self) -> BaseModel:
|
||
"""反序列化 params 为具体模型"""
|
||
if self.type == TaskType.DECODE_ORDER:
|
||
return DecodeTaskParams(**self.params)
|
||
raise ValueError(f"Unknown task type: {self.type}")
|
||
|
||
def set_result(self, result: Union[DecodeTaskResult]) -> None:
|
||
"""设置结果（自动序列化）"""
|
||
if self.type == TaskType.DECODE_ORDER:
|
||
self.result = result.dict()
|
||
else:
|
||
raise ValueError(f"Unsupported result for task type: {self.type}")
|
||
|
||
def get_result(self) -> Optional[BaseModel]:
|
||
"""反序列化 result 为具体模型"""
|
||
if self.result is None:
|
||
return None
|
||
if self.type == TaskType.DECODE_ORDER:
|
||
return DecodeTaskResult(**self.result)
|
||
raise ValueError(f"Unknown task type: {self.type}")
|