hello-agents/Co-creation-projects/lgs-only-NovelGenerator/agents/chapter_generate_agent.py

from dotenv import load_dotenv
load_dotenv()
import re
import os
import json
from pydantic import BaseModel
from typing import List, Dict, Any
from datetime import datetime
from hello_agents import SimpleAgent, HelloAgentsLLM
from hello_agents.tools import NoteTool
from prompt import CHAPTER_PROMPT, CHAPTER_REVIEW_PROMPT, CHAPTER_START_PROMPT


def extract_note_id(output: str) -> str:
    """从 NoteTool 的输出文本中提取 note_id"""
    match = re.search(r"ID:\s*(note_[0-9_]+)", output)
    if not match:
        raise ValueError(f"无法从输出解析 note_id:\n{output}")
    return match.group(1)


class MemoryItem(BaseModel):
    """记忆项数据结构"""
    node_id: str
    novel_id: str
    title: str
    content: str
    summary: str
    timestamp: datetime
    metadata: Dict[str, Any] = {}
    next_chapter_prediction: str = ""


class ChapterGenerateAgent:
    """具有上下文感知能力的 Agent"""

    def __init__(self, name: str, llm: HelloAgentsLLM = HelloAgentsLLM(), max_steps: int = 5, chapter_length: int = 3000, **kwargs):

        self.chapter_length = chapter_length
        self.max_steps = max_steps

        self.num_chapter_memories = kwargs.get("num_chapter_memories", 5)
        self.workspace = kwargs.get("workspace", "./outputs")
        self.note_tools: Dict[str, NoteTool] = {}
        
        self.generate_agent = SimpleAgent(name="章节生成助手", llm=llm, system_prompt='你是一位擅长长篇小说结构与文本细化的专业作者助理。')
        self.review_agent = SimpleAgent(name="章节审核助手", llm=llm, system_prompt='你是一位专业的小说审核助手，负责检查章节是否符合小说的结构和风格。')

        # 内存存储
        self.memories: Dict[str, List[MemoryItem]] = {}

    @staticmethod
    def extract_json_from_response(response: str) -> dict:
        """从模型输出中提取并解析 JSON"""
        # 尝试清理 Markdown 代码块标记
        clean_response = re.sub(r"```json\s*", "", response)
        clean_response = re.sub(r"```\s*$", "", clean_response)
        clean_response = clean_response.strip()
        
        try:
            return json.loads(clean_response)
        except json.JSONDecodeError as e:
            # 如果直接解析失败，尝试在文本中寻找第一个 { 和最后一个 }
            try:
                start = clean_response.find("{")
                end = clean_response.rfind("}")
                if start != -1 and end != -1:
                    json_str = clean_response[start : end + 1]
                    return json.loads(json_str)
            except Exception:
                pass
            raise ValueError(f"无法解析 JSON 响应: {response}") from e

    def _ensure_tool(self, novel_id: str, novel_title: str = None):
        if not self.note_tools.get(novel_id):
            if not novel_title:
                raise ValueError(f"Tool for novel_id {novel_id} not initialized and novel_title not provided.")
            self.note_tools[novel_id] = NoteTool(workspace=os.path.join(self.workspace, f"{novel_title}-{novel_id}", 'chapters'))

    def get_content_from_note(self, content: str) -> str:
        try:
            # 去除 YAML 前置元数据
            frontmatter_match = re.match(r'^---\s*\n(.*?)\n---\s*\n', content, re.DOTALL)
            if frontmatter_match:
                content = content[frontmatter_match.end():].strip()
            
            # 去除标题（第一行如果是标题）
            lines = content.split('\n')
            if lines and lines[0].startswith('# '):
                content = '\n'.join(lines[1:]).strip()
            
            return content
        except:
            return content

    def get_memories(self, novel_id: str):
        """获取最近章节记忆"""
        if not hasattr(self.note_tools[novel_id], "notes_index"):
            self.note_tools[novel_id]._load_index()

        notes = self.note_tools[novel_id].notes_index.get("notes", [])

        # 筛选相关章节笔记
        chapter_notes = [
            n for n in notes
            if n.get("note_type") == "chapter" and str(novel_id) in n.get("title", "")
        ]

        # 获取最后 N 章
        recent_notes = chapter_notes[-self.num_chapter_memories:]

        for note in recent_notes:
            note_id = note.get("id")
            file_path = os.path.join(self.workspace, f"{note_id}.md")

            if os.path.exists(file_path):
                with open(file_path, "r", encoding="utf-8") as f:
                    content = f.read()

                content = self.get_content_from_note(content)
                self.memories[novel_id].append(MemoryItem(
                    node_id=str(note_id),
                    title=note.get("title", "未知章节").strip(),
                    content=content,
                    novel_id=str(novel_id),
                    summary=note['tags'][0]if note.get("tags") and note['tags'] else '',
                    timestamp=datetime.fromisoformat(note.get("created_at", datetime.now().isoformat()))
                ))

    def run(self, user_input: str, **kwargs) -> str:
        """运行 Agent"""
        # 小说id用来区分小说，命名可能会重复
        novel_id = kwargs.pop("novel_id", None)
        assert novel_id, "请提供小说ID"

        novel_title = kwargs.pop("novel_title", None)
        assert novel_title, "请提供小说标题"

        self._ensure_tool(novel_id, novel_title)

        if not self.memories.get(novel_id):
            self.memories[novel_id] = []
            self.get_memories(novel_id)

        # 1. 构建上下文
        outline = self.get_outline(novel_id)
        prev_chapter = self.get_prev_chapter(novel_id)
        prev_summaries = self.get_prev_summaries(novel_id)
        chapter_length = kwargs.get("chapter_length", self.chapter_length)
        context = self.get_prompt(outline, prev_chapter, prev_summaries, user_input, novel_id, chapter_length=chapter_length)
        
        # 2. 使用上下文调用 LLM
        steps = 0
        while steps < self.max_steps:
            steps += 1

            # 生成章节内容
            response = self.generate_agent.run(context)
            try:
                response_data = self.extract_json_from_response(response)
                # 检查是否包含必要字段
                if 'title' not in response_data or 'content' not in response_data or 'next_chapter_prediction' not in response_data or 'summary' not in response_data:
                    raise ValueError("JSON 响应缺少必要字段 'title' 或 'content' 或 'next_chapter_prediction' 或 'summary'")
            except ValueError as e:
                print(f"步骤 {steps} 生成的 JSON 解析错误：{e}")
                continue
            
            # 审核章节内容
            review_context = CHAPTER_REVIEW_PROMPT.format(
                outline=outline,
                prev_chapter=prev_chapter,
                prev_summaries=prev_summaries,
                chapter_content=response_data.get('content', '')
            )
            review_response = self.review_agent.run(review_context)

            # 检查审核结果
            if "【通过】" in review_response:
                break
            
            context = self.get_prompt(outline, prev_chapter, prev_summaries, user_input, novel_id, response_data, review_response, chapter_length=chapter_length)

        # 3. 保存章节到笔记
        create_output = self.note_tools[novel_id].run({
            "action": "create",
            "title": f"{response_data.get('title', '未知章节')}",
            "content": response_data.get('content', ''),
            "note_type": "chapter",
            "tags": [response_data.get('summary', '')]
        })

        # 获取章节笔记ID，保存记忆，并建立与小说ID的关联
        note_id = extract_note_id(create_output)

        self.memories[novel_id].append(MemoryItem(
            node_id=note_id,
            title=response_data.get('title', '未知章节'),
            content=response_data.get('content', ''),
            novel_id=novel_id,
            summary=response_data.get('summary', ''),
            timestamp=datetime.now().isoformat(),
            next_chapter_prediction=response_data.get('next_chapter_prediction', '')
        ))

        return response_data, note_id

    def get_prompt(self, outline: str, prev_chapter: str, prev_summaries: str, user_input: str, novel_id: str, response_data: dict = None, review_response: str = None, chapter_length: int = None) -> str:
        """获取章节生成提示"""
        if chapter_length is None:
            chapter_length = self.chapter_length
        is_first_chapter = (prev_chapter == '无' and prev_summaries == '无')

        if is_first_chapter:
            prompt_template = CHAPTER_START_PROMPT
            context = prompt_template.format(
                outline=outline,
                chapter_history='无' if response_data is None else response_data.get('content', '无'),
                evaluation=review_response or "无",
                user_input=user_input,
                chapter_length=chapter_length
            )
        else:
            prompt_template = CHAPTER_PROMPT
            context = prompt_template.format(
                outline=outline,
                prev_chapter=prev_chapter,
                prev_summaries=prev_summaries,
                chapter_history='无' if response_data is None else response_data.get('content', '无'),
                evaluation=review_response or "无",
                user_input=user_input or [self.memories[novel_id][-1].next_chapter_prediction if self.memories[novel_id] else "无"][0],
                chapter_length=chapter_length
            )
        return context

    def get_outline(self, novel_id: str) -> str:    
        """获取大纲"""
        dir_path = f"{os.path.dirname(self.note_tools[novel_id].workspace)}/outline"
        paths = os.listdir(dir_path)
        assert len(paths) >= 1, f"目录 {dir_path} 下应该有大纲文件"
        # 简单取第一个文件，实际可能需要更精确的逻辑
        path = f"{dir_path}/{paths[0]}"
        with open(path, "r", encoding='utf-8') as f:
            outline = f.read()
        return self.get_content_from_note(outline)

    def get_prev_chapter(self, novel_id: str):
        """获取前一章内容"""
        if self.memories.get(novel_id):
            last_mem = self.memories[novel_id][-1]
            return f"【{last_mem.metadata.get('title', '未知')}】\n...{last_mem.content[-800:]}"
        return "无"

    def get_prev_summaries(self, novel_id: str):
        if self.memories.get(novel_id):
            return "\n".join([f"【{mem.title}】\n{mem.summary}" for mem in self.memories[novel_id][-self.num_chapter_memories:]])
        return "无"
    
    def del_chapter(self, novel_id:str, note_id: str, novel_title: str = None):
        """删除章节"""
        if novel_title:
            self._ensure_tool(novel_id, novel_title)
        self.note_tools[novel_id].run({
            "action": "delete",
            "note_id": note_id
        })
        # 从记忆中删除该章节
        if self.memories.get(novel_id):
            self.memories[novel_id] = [mem for mem in self.memories[novel_id] if mem.node_id != note_id]

    def update_chapter(self, novel_id:str, note_id: str, novel_title: str = None, **kwargs):
        """更新章节"""
        if novel_title:
            self._ensure_tool(novel_id, novel_title)
        self.note_tools[novel_id].run({
            "action": "update",
            "note_id": note_id,
            **kwargs
        })
        # 更新记忆中的章节内容
        if self.memories.get(novel_id):
            for mem in self.memories[novel_id]:
                if mem.node_id == note_id:
                    mem.title = kwargs.get('title', mem.title)
                    mem.content = kwargs.get('content', mem.content)
                    mem.summary = kwargs.get('summary', mem.summary)
                    mem.next_chapter_prediction = kwargs.get('next_chapter_prediction', mem.next_chapter_prediction)
                    mem.timestamp = datetime.now().isoformat()
                    break

def main():
    print("=" * 80)
    print("Novel ChapterGenerateAgent 示例")
    print("=" * 80 + "\n")

    # llm = HelloAgentsLLM(model="qwen3:0.6b", api_key="ollama", base_url="http://127.0.0.1:11434/v1", provider='ollama')
    llm = HelloAgentsLLM(provider='qwen')
    novel_id = "demo_novel_001"
    novel_title = "记忆之城"

    # 1. 模拟大纲文件存在
    # 因为 ChapterGenerateAgent.get_outline 依赖于文件系统查找大纲
    # 我们手动创建一个假的大纲文件用于测试
    workspace_root = "./outputs"
    # 注意：这里模拟 OutlineAgent 的输出路径结构
    outline_dir = os.path.join(workspace_root, f"{novel_title}-{novel_id}", "outline")
    if not os.path.exists(outline_dir):
        os.makedirs(outline_dir)
    
    # 清理旧文件以确保测试环境干净
    for f in os.listdir(outline_dir):
        try:
            os.remove(os.path.join(outline_dir, f))
        except Exception:
            pass
        
    dummy_outline_content = """---
tags: [outline]
created_at: 2025-01-27T10:00:00
---
# 记忆之城-大纲

## 核心梗概
一位能与城市记忆对话的年轻人，在拆迁浪潮中发现一段被刻意抹去的历史。

## 主要人物
- 李寻：主角，拥有"读取"物体记忆的能力。
- 陈叔：古董店老板，似乎知道李寻身世的秘密。

## 故事走向
1. 觉醒能力，卷入拆迁冲突。
2. 发现神秘物品，引出旧事。
3. ...
"""
    dummy_outline_path = os.path.join(outline_dir, f"{novel_id}-outline.md")
    with open(dummy_outline_path, "w", encoding="utf-8") as f:
        f.write(dummy_outline_content)

    print(f"已创建模拟大纲文件: {dummy_outline_path}")
    
    # 2. 初始化章节生成 Agent
    chapter_agent = ChapterGenerateAgent(
        name="小说章节助手",
        llm=llm,
        workspace=workspace_root,  # 使用与 OutlineAgent 一致的根目录
        chapter_length=1000 # 演示用，设短一点
    )

    # 3. 生成第一章
    print(f"\n正在生成第一章...")
    try:
        # run 方法需要 novel_title 来定位目录
        chapter_data_1, note_id_1 = chapter_agent.run(
            user_input="第一章需要通过一个具体的拆迁冲突场景，引出主角的能力。主角李寻在试图保护一家老店不被强拆时，无意中听到了推土机的'心声'。",
            novel_id=novel_id,
            novel_title=novel_title 
        )
        print(f"第一章生成完成，Note ID: {note_id_1}")
        print(f"标题: {chapter_data_1.get('title')}")
        print(f"摘要: {chapter_data_1.get('summary')}")
        print(f"下一章预测: {chapter_data_1.get('next_chapter_prediction')}")

        # 4. 生成第二章（会自动读取第一章作为上下文）
        print(f"\n正在生成第二章...")
        chapter_data_2, note_id_2 = chapter_agent.run(
            user_input="主角在废墟中发现了一个奇怪的物品，触发了回忆。那个物品似乎在呼唤他。",
            novel_id=novel_id,
            novel_title=novel_title
        )
        print(f"第二章生成完成，Note ID: {note_id_2}")
        print(f"标题: {chapter_data_2.get('title')}")
        print(f"摘要: {chapter_data_2.get('summary')}")
        
    except Exception as e:
        print(f"生成过程中出错: {e}")
        import traceback
        traceback.print_exc()



if __name__ == "__main__":
    main()