jaspersoft-agent-learn/step_02_state/concept.py

"""
Step 02: 理解 State - 状态管理

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

🎓 本节内容：
    1. 什么是 Agent State？
    2. 如何设计 State 结构？
    3. State 如何在多步骤任务中传递？

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
"""

from typing import TypedDict, List, Dict, Any, Optional
from dataclasses import dataclass, field
from datetime import datetime
import json


# ═══════════════════════════════════════════════════════════════════════════════
# 第一部分：理解为什么需要 State
# ═══════════════════════════════════════════════════════════════════════════════

"""
在开始写代码之前，我们先理解 State 的本质。

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
场景：用户想生成一个报表，然后修改它
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

❌ 没有 State 的实现（错误）：
    def handle_request(user_input):
        if "生成" in user_input:
            jrxml = generate_jrxml(user_input)
            return jrxml  # 生成后就"丢"了

        if "修改" in user_input:
            # 糟糕！我不知道当前报表是什么
            # 只能让用户重新描述
            return "请重新描述你想要修改的报表"

✅ 有 State 的实现（正确）：
    class Agent:
        def __init__(self):
            self.state = {}  # 用一个字典存储状态

        def handle_request(self, user_input):
            if "生成" in user_input:
                jrxml = generate_jrxml(user_input)
                self.state["current_jrxml"] = jrxml  # 保存到状态
                return jrxml

            if "修改" in user_input:
                current = self.state.get("current_jrxml")  # 从状态读取
                if not current:
                    return "没有可修改的报表"
                modified = modify_jrxml(current, user_input)
                self.state["current_jrxml"] = modified  # 更新状态
                return modified

这就是 State 的作用：在多次交互中保持信息！
"""

# ═══════════════════════════════════════════════════════════════════════════════
# 第二部分：设计 State 的数据结构
# ═══════════════════════════════════════════════════════════════════════════════

"""
设计 State 时，我们使用 Python 的 TypedDict
这是因为：
    1. 有类型提示，IDE 能帮你检查错误
    2. 有代码补全，写代码更方便
    3. 文档化，其他人知道 State 里有什么
"""

class AgentState(TypedDict, total=False):
    """
    Agent 的状态定义

    为什么用 TypedDict 而不是 dataclass？
        因为 TypedDict 更直观，看起来就像一个字典
        而且 LangGraph 直接支持 TypedDict

    total=False 的含义：
        所有字段都是可选的
        这样初始化时可以只填需要的字段

    每个字段的用途：
        - user_input: 当前用户的输入
        - current_jrxml: 当前正在编辑的报表代码
        - conversation_history: 对话历史
        - status: 当前状态（处理中/完成/错误）
        - error_msg: 错误信息（如果有）
    """
    # === 核心工作字段 ===
    user_input: str                    # 用户当前输入
    current_jrxml: str                # 当前 JRXML 代码
    status: str                       # 处理状态: "processing" / "success" / "error"
    error_msg: str                    # 错误信息

    # === 对话相关 ===
    conversation_history: List[dict]  # 对话历史 [{"role": "user", "content": "..."}]
    full_conversation_history: List[dict]  # 完整的对话历史（含时间戳）

    # === 生成相关 ===
    stage: str                       # 当前阶段: "initial" / "refine" / "mapping"
    generated_jrxml: str            # 生成的完整 JRXML
    is_modified: bool                # 是否有未保存的修改

    # === 验证相关 ===
    validation_result: dict          # 验证结果
    retry_count: int                 # 重试次数

    # === 元信息 ===
    session_id: str                  # 会话 ID
    created_at: str                  # 创建时间
    updated_at: str                  # 更新时间


# ═══════════════════════════════════════════════════════════════════════════════
# 第三部分：实际应用 - Jaspersoft 报表生成的完整状态
# ═══════════════════════════════════════════════════════════════════════════════

class JaspersoftAgentState(TypedDict, total=False):
    """
    Jaspersoft 报表生成 Agent 的完整状态

    这个状态设计对应了你实际项目中的需求
    我们逐个解释每个字段的作用
    """

    # ═══════════════════════════════════════════════════════════════════════
    # 1. 基础信息
    # ═══════════════════════════════════════════════════════════════════════

    session_id: str          # 会话唯一标识
    session_name: str        # 会话名称（用户友好）
    created_at: str          # 创建时间
    updated_at: str          # 最后更新时间

    # ═══════════════════════════════════════════════════════════════════════
    # 2. 用户输入
    # ═══════════════════════════════════════════════════════════════════════

    user_input: str         # 用户当前的输入
    uploaded_file_path: str  # 上传的文件路径（如果有）

    # ═══════════════════════════════════════════════════════════════════════
    # 3. 对话历史
    # ═══════════════════════════════════════════════════════════════════════

    """
    对话历史的设计考虑：

    为什么需要两种历史？
        1. conversation_history：精简版，用于发送给 LLM（节省 token）
        2. full_conversation_history：完整版，包含时间戳等元信息（用于审计）

    为什么不直接保存所有消息？
        因为 LLM 有上下文长度限制
        当对话很长时，我们只能发送最近的几轮
        所以需要"精简版"和"完整版"的区分
    """
    conversation_history: List[dict]         # 精简对话历史
    full_conversation_history: List[dict]    # 完整对话历史
    compressed_history: str                  # 压缩后的早期对话

    # ═══════════════════════════════════════════════════════════════════════
    # 4. 报表相关
    # ═══════════════════════════════════════════════════════════════════════

    """
    报表相关的状态字段是最核心的部分

    current_jrxml vs final_jrxml 的区别：
        - current_jrxml：正在编辑的版本，可能还没验证通过
        - final_jrxml：经过验证的最终版本，可以导出

    为什么需要版本历史？
        - 支持撤销操作
        - 用户可能想回退到之前的某个版本
        - 记录每次修改的轨迹
    """
    current_jrxml: str       # 当前正在编辑的 JRXML
    final_jrxml: str         # 最终确认的 JRXML（验证通过）

    # 版本管理
    jrxml_versions: List[dict]   # 历史版本列表
    history_states: List[dict]   # 历史状态快照（用于撤销）
    last_saved_version: int     # 最后保存的版本号

    # ═══════════════════════════════════════════════════════════════════════
    # 5. 生成过程
    # ═══════════════════════════════════════════════════════════════════════

    """
    生成过程的中间状态

    为什么要记录这些？
        1. 用户可能想知道生成到哪一步了
        2. 出错时可以定位问题在哪一步
        3. 方便调试和优化
    """
    stage: str              # 当前阶段
    """
    可能的阶段值：
        - "initial_generation": 初始生成
        - "layout_refine": 布局精调
        - "field_mapping": 字段映射
        - "validation": 验证
        - "correction": 修正
    """

    intent: str             # 用户意图
    """
    可能的意图值：
        - "initial_generation": 生成新报表
        - "modify_report": 修改现有报表
        - "preview_report": 预览报表
        - "consult_question": 咨询问题
    """

    # 生成相关的中间结果
    retrieved_context: str      # RAG 检索到的上下文
    layout_schema: dict         # OCR 分析出的布局信息
    ocr_extraction_result: dict # OCR 提取的字段

    # ═══════════════════════════════════════════════════════════════════════
    # 6. 验证和错误处理
    # ═══════════════════════════════════════════════════════════════════════

    """
    验证和错误处理是 Agent 可靠性的关键

    retry_count 的设计：
        - 每次生成失败后 +1
        - 达到上限后停止重试
        - 这样避免无限循环
    """
    status: str                    # 状态：success / error / processing
    error_msg: str                 # 错误信息
    retry_count: int               # 当前重试次数
    max_retries: int               # 最大重试次数

    # 错误处理
    pending_failure_context: dict   # 待处理的失败上下文
    """
    这个字段用于"失败恢复"
    当重试耗尽时，我们保存失败信息
    下次用户输入时，自动注入这个上下文
    """

    # ═══════════════════════════════════════════════════════════════════════
    # 7. 知识库相关
    # ═══════════════════════════════════════════════════════════════════════

    """
    知识库（KB）相关的状态

    多租户设计：
        - kb_id：当前会话绑定的知识库 ID
        - 不同用户/项目可以使用不同的知识库
    """
    kb_id: str               # 当前知识库 ID
    kb_fields: List[dict]   # 知识库中的字段定义
    kb_template_jrxml: str  # 知识库中的模板 JRXML

    # ═══════════════════════════════════════════════════════════════════════
    # 8. 用户解释（让用户理解 Agent 在做什么）
    # ═══════════════════════════════════════════════════════════════════════

    """
    natural_explanation 是给用户看的解释

    为什么需要这个？
        - 用户不只是想知道结果，还想知道 Agent 是怎么想的
        - 如果出错，用户想知道哪里出了问题
        - 这增加了透明度和信任
    """
    natural_explanation: str  # 对用户的自然语言解释


# ═══════════════════════════════════════════════════════════════════════════════
# 第四部分：State 的操作工具函数
# ═══════════════════════════════════════════════════════════════════════════════

def create_initial_state(session_id: str) -> JaspersoftAgentState:
    """
    创建初始状态

    这是一个工厂函数，用于生成新的状态实例
    所有必要的默认值在这里设置

    为什么用工厂函数而不是直接初始化？
        1. 确保所有必填字段有默认值
        2. 统一初始化逻辑
        3. 方便以后修改默认行为
    """
    now = datetime.now().isoformat()

    return JaspersoftAgentState(
        # 基础信息
        session_id=session_id,
        session_name="新会话",
        created_at=now,
        updated_at=now,

        # 对话历史
        conversation_history=[],
        full_conversation_history=[],
        compressed_history="",

        # 报表相关
        current_jrxml="",
        final_jrxml="",
        jrxml_versions=[],
        history_states=[],
        last_saved_version=0,

        # 生成过程
        stage="initial",
        intent="initial_generation",
        retrieved_context="",
        layout_schema={},
        ocr_extraction_result={},

        # 验证和错误
        status="processing",
        error_msg="",
        retry_count=0,
        max_retries=5,

        # 知识库
        kb_id="",
        kb_fields=[],
        kb_template_jrxml="",

        # 解释
        natural_explanation="",
    )


def update_state(state: JaspersoftAgentState, **updates) -> JaspersoftAgentState:
    """
    更新状态

    这是一个辅助函数，用于安全地更新状态字段

    为什么需要这个函数？
        1. 自动更新时间戳
        2. 类型检查（确保字段存在）
        3. 记录更新历史（可选）

    用法：
        state = update_state(state, current_jrxml="new content", status="success")
    """
    # 更新指定字段
    for key, value in updates.items():
        if key in state:
            state[key] = value
        else:
            raise KeyError(f"State 没有字段: {key}")

    # 自动更新时间戳
    state["updated_at"] = datetime.now().isoformat()

    return state


def save_state_snapshot(state: JaspersoftAgentState) -> dict:
    """
    保存状态快照

    这用于"撤销"功能
    在执行重要操作前，保存当前状态的快照
    如果操作失败，可以回滚到这个快照

    返回的快照包含：
        - 报表内容
        - 对话历史
        - 意图
        - 用户请求
    """
    return {
        "current_jrxml": state.get("current_jrxml", ""),
        "final_jrxml": state.get("final_jrxml", ""),
        "status": state.get("status", ""),
        "conversation_history": list(state.get("conversation_history", [])),
        "user_input": state.get("user_input", ""),
        "intent": state.get("intent", ""),
        "timestamp": datetime.now().isoformat(),
    }


def restore_state_snapshot(state: JaspersoftAgentState, snapshot: dict) -> JaspersoftAgentState:
    """
    从快照恢复状态

    用于"撤销"操作
    从历史快照中恢复之前保存的状态
    """
    state["current_jrxml"] = snapshot.get("current_jrxml", "")
    state["final_jrxml"] = snapshot.get("final_jrxml", "")
    state["status"] = snapshot.get("status", "")
    state["conversation_history"] = snapshot.get("conversation_history", [])
    state["updated_at"] = datetime.now().isoformat()

    return state


# ═══════════════════════════════════════════════════════════════════════════════
# 第五部分：演示代码
# ═══════════════════════════════════════════════════════════════════════════════

def demo():
    """
    演示 State 的使用
    """
    print("=" * 60)
    print("Step 02: 理解 State - 状态管理演示")
    print("=" * 60)

    # 1. 创建初始状态
    print("\n📦 步骤 1: 创建初始状态")
    state = create_initial_state("session_001")
    print(f"  会话 ID: {state['session_id']}")
    print(f"  创建时间: {state['created_at']}")
    print(f"  当前状态: {state['status']}")

    # 2. 更新状态
    print("\n🔄 步骤 2: 更新状态")
    state = update_state(
        state,
        user_input="生成一个销售报表",
        current_jrxml='<?xml version="1.0"?><jasperReport/>',
        status="success",
    )
    print(f"  用户输入: {state['user_input']}")
    print(f"  生成状态: {state['status']}")
    print(f"  JRXML 长度: {len(state['current_jrxml'])} 字符")

    # 3. 保存快照
    print("\n📸 步骤 3: 保存状态快照")
    snapshot = save_state_snapshot(state)
    print(f"  快照时间: {snapshot['timestamp']}")
    print(f"  快照内容: JRXML ({len(snapshot['current_jrxml'])} 字符)")

    # 4. 修改状态（模拟用户修改）
    print("\n✏️ 步骤 4: 修改报表（模拟）")
    state["current_jrxml"] = '<?xml version="1.0"?><jasperReport modified="true"/>'
    state["status"] = "modified"
    print(f"  新状态: {state['status']}")
    print(f"  新 JRXML: {state['current_jrxml']}")

    # 5. 撤销（恢复到快照）
    print("\n↩️ 步骤 5: 撤销操作")
    state = restore_state_snapshot(state, snapshot)
    print(f"  恢复状态: {state['status']}")
    print(f"  恢复 JRXML: {state['current_jrxml']}")

    # 6. 模拟完整的生成流程
    print("\n🔄 步骤 6: 模拟完整生成流程")
    state = create_initial_state("session_002")

    # 阶段 1: 用户输入
    state["user_input"] = "生成一个采购单报表"
    state["intent"] = "initial_generation"
    print(f"  [阶段1] 用户输入: {state['user_input']}")

    # 阶段 2: 生成
    state["current_jrxml"] = "<?xml>...生成的 JRXML...</xml>"
    state["stage"] = "initial_generation"
    print(f"  [阶段2] 生成完成，长度: {len(state['current_jrxml'])}")

    # 阶段 3: 验证
    state["status"] = "success"
    state["final_jrxml"] = state["current_jrxml"]
    print(f"  [阶段3] 验证通过，状态: {state['status']}")

    print("\n" + "=" * 60)
    print("✅ State 管理演示完成")
    print("=" * 60)


if __name__ == "__main__":
    demo()