"""Core Pipeline Schemas 定义工作流管道的核心 Schema,包括 StepInput、StepOutput、StepResult、 StepStatus、StepRecord、WorkflowState 和 StepDefinition。 使用 Pydantic v2 确保类型安全和自动验证。 """ from pydantic import BaseModel, Field from typing import Optional, Dict, Literal, Any from datetime import datetime from enum import Enum class StepStatus(str, Enum): """步骤执行状态枚举""" PENDING = "pending" RUNNING = "running" COMPLETED = "completed" FAILED = "failed" SKIPPED = "skipped" class StepInput(BaseModel): """所有步骤输入的基类 Attributes: content_id: 内容唯一标识,如 B站 BV号、知乎回答ID等 platform: 平台名称,限定为 bilibili 或 zhihu """ content_id: str = Field(..., min_length=1, description="内容唯一标识") platform: Literal["bilibili", "zhihu"] = Field(..., description="平台名称") class StepOutput(BaseModel): """所有步骤输出的基类 Attributes: success: 步骤是否成功执行 error_message: 错误信息,成功时为 None completed_at: 完成时间,默认为当前时间 """ success: bool = Field(..., description="步骤是否成功") error_message: Optional[str] = Field(None, description="错误信息") completed_at: datetime = Field(default_factory=datetime.now, description="完成时间") class StepResult(BaseModel): """步骤执行结果 包含步骤的输入、输出和执行状态。 Attributes: step_name: 步骤名称 input_data: 步骤输入数据 output_data: 步骤输出数据,可能为 None status: 步骤执行状态 """ step_name: str = Field(..., description="步骤名称") input_data: StepInput = Field(..., description="步骤输入数据") output_data: Optional[StepOutput] = Field(None, description="步骤输出数据") status: StepStatus = Field(default=StepStatus.PENDING, description="步骤执行状态") class StepRecord(BaseModel): """单个步骤的执行记录 记录步骤执行的状态、时间和结果摘要。 Attributes: step_name: 步骤名称 status: 步骤执行状态 started_at: 开始时间 completed_at: 完成时间 error_message: 错误信息 output_summary: 输出摘要 """ step_name: str = Field(..., description="步骤名称") status: StepStatus = Field(..., description="步骤执行状态") started_at: Optional[datetime] = Field(None, description="开始时间") completed_at: Optional[datetime] = Field(None, description="完成时间") error_message: Optional[str] = Field(None, description="错误信息") output_summary: Optional[str] = Field(None, description="输出摘要") class WorkflowState(BaseModel): """工作流状态 - 存储在 workflow_state.json 管理整个工作流的状态,包括内容信息、执行历史和步骤状态。 Attributes: schema_version: Schema 版本号 content_id: 内容唯一标识 platform: 平台名称 source_url: 原始链接 created_at: 创建时间 updated_at: 更新时间 steps: 步骤执行记录字典 """ schema_version: str = Field(default="1.0", description="Schema 版本号") content_id: str = Field(..., description="内容唯一标识") platform: str = Field(..., description="平台名称") source_url: str = Field(..., description="原始链接") created_at: datetime = Field(default_factory=datetime.now, description="创建时间") updated_at: datetime = Field(default_factory=datetime.now, description="更新时间") steps: Dict[str, StepRecord] = Field(default_factory=dict, description="步骤执行记录") def is_step_completed(self, step_name: str) -> bool: """检查指定步骤是否已完成 Args: step_name: 步骤名称 Returns: 如果步骤存在且状态为 COMPLETED 返回 True,否则返回 False """ record = self.steps.get(step_name) return record is not None and record.status == StepStatus.COMPLETED def mark_step_started(self, step_name: str) -> None: """标记步骤开始执行 Args: step_name: 步骤名称 """ self.steps[step_name] = StepRecord( step_name=step_name, status=StepStatus.RUNNING, started_at=datetime.now() ) self.updated_at = datetime.now() def mark_step_completed(self, step_name: str, output_summary: Optional[str] = None) -> None: """标记步骤完成 Args: step_name: 步骤名称 output_summary: 输出摘要(可选) """ if step_name in self.steps: self.steps[step_name].status = StepStatus.COMPLETED self.steps[step_name].completed_at = datetime.now() self.steps[step_name].output_summary = output_summary else: self.steps[step_name] = StepRecord( step_name=step_name, status=StepStatus.COMPLETED, completed_at=datetime.now(), output_summary=output_summary ) self.updated_at = datetime.now() def mark_step_failed(self, step_name: str, error_message: str) -> None: """标记步骤失败 Args: step_name: 步骤名称 error_message: 错误信息 """ if step_name in self.steps: self.steps[step_name].status = StepStatus.FAILED self.steps[step_name].error_message = error_message self.steps[step_name].completed_at = datetime.now() else: self.steps[step_name] = StepRecord( step_name=step_name, status=StepStatus.FAILED, error_message=error_message, completed_at=datetime.now() ) self.updated_at = datetime.now() class StepDefinition(BaseModel): """步骤定义 - 描述步骤的前置条件、输出格式、成功标准 Attributes: step_name: 步骤名称 description: 步骤描述 input_schema: 输入 Schema 类型 output_schema: 输出 Schema 类型 dependencies: 依赖的步骤列表 """ step_name: str = Field(..., description="步骤名称") description: str = Field(..., description="步骤描述") input_schema: type[StepInput] = Field(..., description="输入 Schema 类型") output_schema: type[StepOutput] = Field(..., description="输出 Schema 类型") dependencies: list[str] = Field(default_factory=list, description="依赖的步骤列表")