LQConstPlanWriterAgent/core/construction_write/component/state_models.py

# -*- coding: utf-8 -*-
"""
施工方案编写应用 - 状态定义模块
定义 LangGraph 工作流使用的状态类和任务信息类

该模块位于 component 层，为 workflows 层提供状态定义支持
"""

from typing import TypedDict, Optional, List, Dict, Any
from datetime import datetime
from langchain_core.messages import BaseMessage, HumanMessage


class OutlineGenerationState(TypedDict):
    """
    大纲生成 LangGraph 状态定义

    用于在 LangGraph 工作流节点之间传递状态数据
    注意：所有字段必须是可序列化的（不能包含 ProgressManager 等对象）

    Attributes:
        # 输入参数
        callback_task_id: 任务回调ID
        user_id: 用户ID
        project_info: 项目基础信息
        template_id: 模板ID
        generation_chapterenum: 生成章节代码列表
        generation_template: 模板结构
        similarity_config: 相似推荐配置（可选）
        knowledge_config: 知识点配置（可选）

        # 中间状态
        template: 加载的模板内容
        outline_structure: 生成的大纲结构
        key_points: 核心要点列表
        similar_cases: 相似案例推荐结果
        similar_fragments: 相似片段推荐结果
        knowledge_bases: 关联的知识点列表

        # 流程控制
        current_stage: 当前执行阶段
        overall_task_status: 整体任务状态
        error_message: 错误信息
        messages: LangChain 消息列表（用于状态追踪）
        # 注意：progress_manager 和 task_info 存储在类实例中，不在状态中
    """
    # 输入参数
    callback_task_id: str
    user_id: str
    project_info: Dict[str, Any]
    template_id: str
    generation_chapterenum: List[str]
    generation_template: List[Dict[str, Any]]
    similarity_config: Optional[Dict[str, Any]]
    knowledge_config: Optional[Dict[str, Any]]

    # 中间状态
    template: Optional[Dict[str, Any]]
    outline_structure: Optional[Dict[str, Any]]
    key_points: Optional[List[Dict[str, Any]]]
    similar_cases: Optional[List[Dict[str, Any]]]
    similar_fragments: Optional[List[Dict[str, Any]]]
    knowledge_bases: Optional[List[Dict[str, Any]]]

    # 流程控制
    current_stage: str
    overall_task_status: str
    error_message: Optional[str]
    messages: List[BaseMessage]


class OutlineTaskInfo:
    """
    大纲生成任务信息类

    封装大纲生成任务的完整信息，与 outline_views.py 传入的参数保持一致

    Attributes:
        callback_task_id: 任务回调ID，用于唯一标识任务
        user_id: 用户ID
        project_info: 项目基础信息字典（包含 project_name, construct_location, engineering_type, selectable）
        template_id: 大纲模板ID（对应 generation_template.alias）
        generation_chapterenum: 生成章节代码列表
        generation_template: 模板结构（对应 generation_template.structure）
        similarity_config: 相似推荐配置（包含 topk_plans, topk_fragments, threshold）
        knowledge_config: 知识点配置（包含 topk, threshold）
        status: 任务状态（pending/processing/completed/failed/cancelled）
        start_time: 任务开始时间戳
        end_time: 任务结束时间戳
        results: 任务执行结果
    """

    def __init__(
        self,
        callback_task_id: str,
        user_id: str,
        project_info: Dict[str, Any],
        template_id: str,
        generation_chapterenum: List[str],
        generation_template: List[Dict[str, Any]],
        similarity_config: Optional[Dict[str, Any]] = None,
        knowledge_config: Optional[Dict[str, Any]] = None
    ):
        """
        初始化大纲生成任务信息

        Args:
            callback_task_id: 任务回调ID
            user_id: 用户ID
            project_info: 项目基础信息（包含 project_name, construct_location, engineering_type, selectable）
            template_id: 模板ID（对应 generation_template.alias）
            generation_chapterenum: 生成章节代码列表
            generation_template: 模板结构（对应 generation_template.structure）
            similarity_config: 相似推荐配置（包含 topk_plans, topk_fragments, threshold）
            knowledge_config: 知识点配置（包含 topk, threshold）
        """
        self.callback_task_id = callback_task_id
        self.user_id = user_id
        self.project_info = project_info
        # 从 project_info 中提取常用字段（便于快速访问）
        self.project_name = project_info.get("project_name", "")
        self.construct_location = project_info.get("construct_location", "")
        self.engineering_type = project_info.get("engineering_type", "")
        self.selectable = project_info.get("selectable", "")

        # 模板相关
        self.template_id = template_id
        self.generation_chapterenum = generation_chapterenum
        self.generation_template = generation_template

        # 配置（使用 outline_views.py 中定义的结构）
        self.similarity_config = similarity_config or {
            "topk_plans": 3,
            "topk_fragments": 10,
            "threshold": 0.75
        }
        self.knowledge_config = knowledge_config or {
            "topk": 3,
            "threshold": 0.75
        }

        # 任务状态
        self.status = "pending"
        self.start_time: Optional[float] = None
        self.end_time: Optional[float] = None
        self.results: Dict[str, Any] = {}

    def start_processing(self):
        """标记任务开始处理"""
        import time
        self.status = "processing"
        self.start_time = time.time()

    def complete_processing(self, results: Dict[str, Any]):
        """
        标记任务完成

        Args:
            results: 任务执行结果
        """
        import time
        self.status = "completed"
        self.end_time = time.time()
        self.results = results

    def fail_processing(self, error: str):
        """
        标记任务失败

        Args:
            error: 错误信息
        """
        import time
        self.status = "failed"
        self.end_time = time.time()
        self.results = {"error": error}

    def cancel_processing(self):
        """标记任务被取消"""
        import time
        self.status = "cancelled"
        self.end_time = time.time()

    def __repr__(self) -> str:
        return (
            f"OutlineTaskInfo("
            f"callback_task_id={self.callback_task_id}, "
            f"user_id={self.user_id}, "
            f"project_name={self.project_name}, "
            f"status={self.status}"
            f")"
        )

    def to_dict(self) -> Dict[str, Any]:
        """
        转换为字典格式（用于序列化）

        Returns:
            Dict: 任务信息的字典表示
        """
        return {
            "callback_task_id": self.callback_task_id,
            "user_id": self.user_id,
            "project_info": self.project_info,
            "project_name": self.project_name,
            "construct_location": self.construct_location,
            "engineering_type": self.engineering_type,
            "selectable": self.selectable,
            "template_id": self.template_id,
            "generation_chapterenum": self.generation_chapterenum,
            "generation_template": self.generation_template,
            "similarity_config": self.similarity_config,
            "knowledge_config": self.knowledge_config,
            "status": self.status,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "results": self.results
        }


class ContentGenerationState(TypedDict):
    """
    内容生成 LangGraph 状态定义

    用于内容补全/续写工作流

    Attributes:
        task_id: 任务ID
        user_id: 用户ID
        section_path: 章节路径
        current_content: 当前内容
        context_window: 上下文窗口大小
        completion_mode: 补全模式（continue/expand/rewrite/optimize）
        target_length: 目标长度
        generated_content: 生成的内容
        references: 参考知识点
    """
    task_id: str
    user_id: str
    section_path: str
    current_content: str
    context_window: int
    completion_mode: str
    target_length: int
    generated_content: Optional[str]
    references: Optional[List[Dict[str, Any]]]
    overall_task_status: str
    error_message: Optional[str]


__all__ = [
    "OutlineGenerationState",
    "OutlineTaskInfo",
    "ContentGenerationState"
]