bidmaster-cli/src/bidmaster/nodes/toc/generate_sub_chapters.py

"""生成二三级子标题的节点

通过AI智能生成或模板生成二三级子标题。
"""

import json
import logging
from typing import Dict, List, Any, Optional

from ..base import BaseNode, NodeContext
from ...tools.parser import ScoringCriteria, DocumentChapter
from ...tools.llm import llm_service

logger = logging.getLogger(__name__)


class GenerateSubChaptersNode(BaseNode):
    """生成二三级子标题的节点"""

    @property
    def name(self) -> str:
        return "generate_sub_chapters"

    @property
    def description(self) -> str:
        return "生成二三级子标题"

    def execute(self, state: Dict[str, Any], context: NodeContext) -> Dict[str, Any]:
        """执行子标题生成"""
        self._log_start()

        try:
            preliminary_chapters = state.get("preliminary_chapters", [])
            technical_criteria = state.get("technical_criteria", [])

            if not preliminary_chapters:
                raise ValueError("缺少初步章节")

            # 只使用AI生成，不再需要交互选择
            generation_mode = "ai"
            template_file = None

            # 为每个章节生成子标题
            enhanced_chapters = []
            for chapter in preliminary_chapters:
                enhanced_chapter = self._enhance_chapter_with_subs(
                    chapter, technical_criteria, generation_mode, template_file
                )
                enhanced_chapters.append(enhanced_chapter)

            logger.info(f"完成{len(enhanced_chapters)}个章节的子标题生成")

            # 更新状态
            state["preliminary_chapters"] = enhanced_chapters
            state["current_step"] = self.name

            self._log_success()
            return state

        except Exception as e:
            self._log_error(e)
            state["error"] = str(e)
            state["should_continue"] = False
            raise

    def _enhance_chapter_with_subs(self,
                                  chapter: DocumentChapter,
                                  technical_criteria: List[ScoringCriteria],
                                  generation_mode: str,
                                  template_file: Optional[str]) -> DocumentChapter:
        """为章节增强子标题

        Args:
            chapter: 原始章节
            technical_criteria: 技术评分项
            generation_mode: 生成模式
            template_file: 模板文件（仅template模式使用）

        Returns:
            增强后的章节
        """
        # 找到该章节对应的评分项
        corresponding_criteria = self._find_corresponding_criteria(chapter, technical_criteria)

        if not corresponding_criteria:
            logger.warning(f"章节 {chapter.title} 没有找到对应的评分项")
            return chapter

        # 根据模式生成子标题
        if generation_mode == "ai":
            sub_chapters = self._generate_ai_sub_chapters(corresponding_criteria, chapter)
        else:
            sub_chapters = self._generate_template_sub_chapters(
                corresponding_criteria[0], chapter, template_file
            )

        # 设置子章节
        chapter.children = sub_chapters
        logger.info(f"章节 {chapter.title} 生成了 {len(sub_chapters)} 个子标题")

        return chapter

    def _find_corresponding_criteria(self,
                                   chapter: DocumentChapter,
                                   technical_criteria: List[ScoringCriteria]) -> List[ScoringCriteria]:
        """查找章节对应的评分项

        Args:
            chapter: 章节
            technical_criteria: 技术评分项列表

        Returns:
            对应的评分项列表
        """
        # 从章节ID提取类别
        if "_" not in chapter.id:
            return []

        parts = chapter.id.split("_")
        if len(parts) < 3:
            return []

        category = "_".join(parts[2:])

        # 找到该类别的所有评分项
        return [
            c for c in technical_criteria
            if c.category.value == category
        ]

    def _generate_ai_sub_chapters(self,
                                criteria_list: List[ScoringCriteria],
                                parent_chapter: DocumentChapter) -> List[DocumentChapter]:
        """AI生成子标题

        Args:
            criteria_list: 对应的评分项列表
            parent_chapter: 父章节

        Returns:
            子章节列表
        """
        try:
            # 构建评分项信息
            criteria_info = []
            for criteria in criteria_list:
                criteria_info.append(f"- {criteria.item_name} ({criteria.max_score}分)")

            prompt = f"""
为以下大类别生成专业的标书子标题：

【大类别】: {parent_chapter.title}
【评分项】:
{chr(10).join(criteria_info)}

生成要求：
1. 为每个评分项生成对应的子标题名称（不要包含编号）
2. 重要评分项可添加三级子标题（不要包含编号）
3. 只返回标题文本，编号由Word自动管理

返回JSON格式：
{{
  "sub_chapters": [
    {{"title": "技术架构设计", "level": 2, "score": 5, "children": []}}
  ]
}}

只返回JSON："""

            # 使用统一的LLM服务
            response = llm_service.call(prompt)
            if not response:
                raise ValueError("AI生成子标题失败: API无响应")

            # 解析响应
            result_data = self._parse_ai_response(response)
            sub_chapters_data = result_data.get("sub_chapters", [])

            # 构建子章节对象
            sub_chapters = []
            for i, sub_data in enumerate(sub_chapters_data, 1):
                title = sub_data.get("title", f"子标题{i}")

                sub_chapter = DocumentChapter(
                    id=f"{parent_chapter.id}_sub_{i:02d}",
                    title=title,  # 不添加编号
                    level=sub_data.get("level", 2),
                    score=sub_data.get("score", 0),
                    template_placeholder=f"{{{{{parent_chapter.id}_sub_{i:02d}_content}}}}"
                )

                # 处理三级标题
                for j, child_data in enumerate(sub_data.get("children", []), 1):
                    child_title = child_data.get("title", f"三级标题{j}")

                    child_chapter = DocumentChapter(
                        id=f"{parent_chapter.id}_sub_{i:02d}_{j:02d}",
                        title=child_title,  # 不添加编号
                        level=child_data.get("level", 3),
                        template_placeholder=f"{{{{{parent_chapter.id}_sub_{i:02d}_{j:02d}_content}}}}"
                    )
                    sub_chapter.children.append(child_chapter)

                sub_chapters.append(sub_chapter)

            return sub_chapters

        except Exception as e:
            logger.error(f"AI生成子标题失败: {e}")
            # 立即失败，不提供后备方案
            raise

    def _parse_ai_response(self, response: str) -> Dict[str, Any]:
        """解析AI响应

        Args:
            response: AI响应文本

        Returns:
            解析后的数据

        Raises:
            ValueError: 解析失败时抛出
        """
        try:
            clean_response = response.strip()
            if clean_response.startswith("```json"):
                clean_response = clean_response[7:]
            if clean_response.endswith("```"):
                clean_response = clean_response[:-3]

            return json.loads(clean_response.strip())

        except (json.JSONDecodeError, KeyError) as e:
            raise ValueError(f"解析AI响应失败: {e}")

    def _generate_template_sub_chapters(self,
                                      criteria: ScoringCriteria,
                                      parent_chapter: DocumentChapter,
                                      template_file: Optional[str]) -> List[DocumentChapter]:
        """基于模板生成子标题

        Args:
            criteria: 评分项（仅作参考）
            parent_chapter: 父章节
            template_file: 模板文件路径

        Returns:
            子章节列表
        """
        # 提供默认结构（不包含编号）
        default_sub_chapters = [
            DocumentChapter(
                id=f"{parent_chapter.id}_def_01",
                title="方案概述",  # 不包含编号
                level=2,
                template_placeholder=f"{{{{{parent_chapter.id}_def_01_content}}}}"
            ),
            DocumentChapter(
                id=f"{parent_chapter.id}_def_02",
                title="具体实施",  # 不包含编号
                level=2,
                template_placeholder=f"{{{{{parent_chapter.id}_def_02_content}}}}"
            ),
            DocumentChapter(
                id=f"{parent_chapter.id}_def_03",
                title="保障措施",  # 不包含编号
                level=2,
                template_placeholder=f"{{{{{parent_chapter.id}_def_03_content}}}}"
            )
        ]

        logger.info(f"使用默认模板为 {parent_chapter.title} 生成子标题")
        return default_sub_chapters