69 lines
2.0 KiB
Python
69 lines
2.0 KiB
Python
"""单章节处理Agent
|
||
|
||
负责处理单个章节的完整流程。
|
||
"""
|
||
|
||
import logging
|
||
from typing import Any, Dict
|
||
|
||
from .base import BaseAgent, AgentBuilder
|
||
from ..nodes.content import (
|
||
GenerateContentNode,
|
||
SaveToWordNode,
|
||
CleanupMarkdownInWordNode,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class SingleChapterAgentBuilder(AgentBuilder):
|
||
"""单章节Agent构建器"""
|
||
|
||
@classmethod
|
||
def create(cls, interaction_handler=None) -> "SingleChapterAgentBuilder":
|
||
"""创建单章节Agent构建器"""
|
||
builder = cls(interaction_handler)
|
||
|
||
# 添加节点（去掉PrepareChapterNode，因为外层已准备好）
|
||
builder.add_node(GenerateContentNode()) \
|
||
.add_node(SaveToWordNode()) \
|
||
.add_node(CleanupMarkdownInWordNode())
|
||
|
||
# 设置入口
|
||
builder.set_entry("generate_content")
|
||
|
||
# 配置流程：线性流程
|
||
builder.add_edge("generate_content", "save_to_word")
|
||
builder.add_edge("save_to_word", "cleanup_markdown_in_word")
|
||
builder.add_edge("cleanup_markdown_in_word", "END")
|
||
|
||
return builder
|
||
|
||
|
||
class SingleChapterAgent(BaseAgent):
|
||
"""单章节处理Agent"""
|
||
|
||
def __init__(self, interaction_handler=None):
|
||
builder = SingleChapterAgentBuilder.create(interaction_handler)
|
||
super().__init__(builder)
|
||
|
||
def process_chapter(self, chapter: Dict[str, Any], state: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""处理单个章节
|
||
|
||
Args:
|
||
chapter: 章节信息
|
||
state: 全局状态（包含word_file等）
|
||
|
||
Returns:
|
||
处理结果
|
||
"""
|
||
# 准备单章节状态（继承所有字段，避免遗漏rag_tool等初始化数据）
|
||
chapter_state = {
|
||
**state, # 继承所有字段（包括rag_tool、expanded_configs等）
|
||
"current_chapter": chapter, # 覆盖当前章节
|
||
}
|
||
|
||
# 执行
|
||
result_state = self.execute_sync(chapter_state)
|
||
|
||
return result_state |