bidmaster-cli/src/bidmaster/agents/builders/toc_builder.py

"""目录生成Agent构建器

用于构建TocGenerator Agent的完整工作流。
"""

import logging
from typing import Dict, Any, Optional


from ..base import AgentBuilder, BaseAgent, BaseAgentFactory
from ...nodes.toc import (
    GroupCriteriaNode,
    GenerateFirstLevelNode,
    GenerateSubChaptersNode,
    ReviewStructureNode,
    ApplyReviewSuggestionsNode,
    AdjustChaptersNode,
    FinalizeChaptersNode
)
from ...nodes.toc.user_feedback import UserFeedbackNode
from ...nodes.toc.optimize_with_feedback import OptimizeWithFeedbackNode
from ...nodes.toc.base_mixins import WorkflowUtilsMixin

logger = logging.getLogger(__name__)


# 使用通用的工作流条件判断函数
should_continue = WorkflowUtilsMixin.should_continue_workflow


class TocAgentBuilder(AgentBuilder):
    """目录生成Agent构建器"""

    @classmethod
    def create(cls, interaction_handler = None) -> 'TocAgentBuilder':
        """创建预配置的TocAgent构建器

        Args:
            interaction_handler: InteractionHandler实例或兼容的交互处理函数

        Returns:
            配置完成的构建器实例
        """
        builder = cls(interaction_handler)

        # 添加所有节点
        builder.add_node(GroupCriteriaNode()) \
               .add_node(GenerateFirstLevelNode()) \
               .add_node(GenerateSubChaptersNode()) \
               .add_node(ReviewStructureNode()) \
               .add_node(ApplyReviewSuggestionsNode()) \
               .add_node(AdjustChaptersNode()) \
               .add_node(FinalizeChaptersNode()) \
               .add_node(UserFeedbackNode()) \
               .add_node(OptimizeWithFeedbackNode())

        # 设置入口点
        builder.set_entry("group_criteria")

        # 配置工作流程
        builder._configure_workflow()

        logger.info("TocAgent构建器创建完成")
        return builder

    def _configure_workflow(self) -> None:
        """配置工作流程"""
        # 添加条件边 - 每个节点都检查是否应该继续
        self.add_conditional_edge(
            "group_criteria",
            should_continue,
            {"continue": "generate_first_level", "end": "END"}
        )

        self.add_conditional_edge(
            "generate_first_level",
            should_continue,
            {"continue": "generate_sub_chapters", "end": "END"}
        )

        self.add_conditional_edge(
            "generate_sub_chapters",
            should_continue,
            {"continue": "review_structure", "end": "END"}
        )

        self.add_conditional_edge(
            "review_structure",
            should_continue,
            {"continue": "apply_suggestions", "end": "END"}
        )

        self.add_conditional_edge(
            "apply_suggestions",
            should_continue,
            {"continue": "adjust_chapters", "end": "END"}
        )

        self.add_conditional_edge(
            "adjust_chapters",
            should_continue,
            {"continue": "finalize_chapters", "end": "END"}
        )

        # finalize_chapters连接到user_feedback
        self.add_conditional_edge(
            "finalize_chapters",
            should_continue,
            {"continue": "user_feedback", "end": "END"}
        )

        # user_feedback的条件边
        def check_needs_optimization(state: Dict[str, Any]) -> str:
            """检查是否需要根据用户反馈进行优化"""
            if state.get("needs_optimization", False):
                return "optimize"
            return "end"

        self.add_conditional_edge(
            "user_feedback",
            check_needs_optimization,
            {"optimize": "optimize_with_feedback", "end": "END"}
        )

        # 优化后循环回finalize_chapters
        self.add_edge("optimize_with_feedback", "finalize_chapters")


class TocAgent(BaseAgent, BaseAgentFactory):
    """目录生成Agent

    封装了目录生成的完整工作流程。
    """

    def __init__(self, interaction_handler = None):
        """初始化TocAgent

        Args:
            interaction_handler: InteractionHandler实例或兼容的交互处理函数
        """
        builder = TocAgentBuilder.create(interaction_handler)
        super().__init__(builder)

    async def generate_toc(self,
                          technical_criteria: list,
                          generation_mode: Optional[str] = None,
                          template_file: Optional[str] = None,
                          document_context: Optional[Any] = None,
                          auto_mode: Optional[bool] = None) -> Dict[str, Any]:
        """生成目录结构

        Args:
            technical_criteria: 技术评分项列表
            generation_mode: 生成模式（ai/template），可选
            template_file: 模板文件路径，可选

        Returns:
            包含生成结果的状态字典
        """
        initial_state = {
            "technical_criteria": technical_criteria,
            "should_continue": True,
            "current_step": "",
            "category_groups": {},
            "preliminary_chapters": [],
            "structure_review": {},
            "final_chapters": [],
            "error": "",
            "warnings": [],
            "document_context": document_context,
            "auto_mode": True if auto_mode is None else auto_mode,
            "user_feedback": ""
        }

        # 如果提供了预设参数，添加到状态中
        if generation_mode:
            initial_state["preset_generation_mode"] = generation_mode
        if template_file:
            initial_state["preset_template_file"] = template_file

        # 添加交互处理器到状态中
        initial_state["interaction_handler"] = self.builder.interaction_handler

        return await self.execute(initial_state)

    def generate_toc_sync(self,
                         technical_criteria: list,
                         generation_mode: Optional[str] = None,
                         template_file: Optional[str] = None,
                         document_context: Optional[Any] = None,
                         auto_mode: Optional[bool] = None) -> Dict[str, Any]:
        """同步生成目录结构

        Args:
            technical_criteria: 技术评分项列表
            generation_mode: 生成模式（ai/template），可选
            template_file: 模板文件路径，可选

        Returns:
            包含生成结果的状态字典
        """
        import asyncio
        return asyncio.run(
            self.generate_toc(technical_criteria, generation_mode, template_file, document_context, auto_mode)
        )