diff --git a/kb_pipeline/schemas.py b/kb_pipeline/schemas.py new file mode 100644 index 0000000..4504a6f --- /dev/null +++ b/kb_pipeline/schemas.py @@ -0,0 +1,187 @@ +"""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="依赖的步骤列表") diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..c7b23ec --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +pythonpath = . +testpaths = tests