agent_jrxml/agent/graph.py

"""LangGraph JRXML 生成代理的状态图定义。"""

import os
from typing import Literal

from dotenv import load_dotenv
from langgraph.graph import StateGraph, END

from agent.state import AgentState
from agent.nodes import (
    load_session_node,
    process_input,
    manage_context,
    save_state_snapshot,
    classify_intent,
    retrieve,
    generate,
    modify_jrxml,
    handle_consult,
    handle_undo,
    handle_reset,
    save_session_node,
    validate,
    explain_error,
    correct_jrxml,
    finalize,
)

load_dotenv()
MAX_RETRY = int(os.getenv("MAX_RETRY", "3"))

# ============================================================
# 路由函数
# ============================================================

def route_by_intent(state: AgentState) -> Literal[
    "retrieve", "modify_jrxml", "save_session",
    "handle_consult", "handle_undo", "handle_reset"
]:
    """根据 classify_intent 的结果路由到对应的处理节点。"""
    intent = state.get("intent", "initial_generation")

    if intent == "initial_generation":
        return "retrieve"
    elif intent == "modify_report":
        return "modify_jrxml"
    elif intent in ("preview_report", "export_pdf", "export_jrxml"):
        return "save_session"
    elif intent == "consult_question":
        return "handle_consult"
    elif intent == "undo_modification":
        return "handle_undo"
    elif intent == "reset_session":
        return "handle_reset"
    else:
        # 兜底：根据是否有报表判断
        if state.get("current_jrxml"):
            return "modify_jrxml"
        return "retrieve"


def route_after_generate(state: AgentState) -> Literal["save_session"]:
    return "save_session"


def route_after_modify(state: AgentState) -> Literal["save_session"]:
    return "save_session"


def route_after_undo(state: AgentState) -> Literal["save_session"]:
    return "save_session"


def route_after_save(state: AgentState) -> Literal["validate", "finalize"]:
    # 预览/导出意图跳过验证，直接完成
    intent = state.get("intent", "")
    if intent in ("preview_report", "export_pdf", "export_jrxml"):
        return "finalize"
    return "validate"


def route_after_validate(state: AgentState) -> Literal["finalize", "explain_error"]:
    if state.get("status") == "pass":
        return "finalize"
    return "explain_error"


def route_after_explain(state: AgentState) -> Literal["correct_jrxml"]:
    return "correct_jrxml"


def route_after_correct(state: AgentState) -> Literal["validate", "finalize"]:
    retry = state.get("retry_count", 0)
    if retry >= MAX_RETRY:
        return "finalize"
    return "validate"


# ============================================================
# 图构建
# ============================================================

def build_graph() -> StateGraph:
    workflow = StateGraph(AgentState)

    # 现有节点
    workflow.add_node("load_session", load_session_node)
    workflow.add_node("process_input", process_input)
    workflow.add_node("manage_context", manage_context)
    workflow.add_node("save_session", save_session_node)
    workflow.add_node("retrieve", retrieve)
    workflow.add_node("generate", generate)
    workflow.add_node("modify_jrxml", modify_jrxml)
    workflow.add_node("validate", validate)
    workflow.add_node("explain_error", explain_error)
    workflow.add_node("correct_jrxml", correct_jrxml)
    workflow.add_node("finalize", finalize)

    # 新增节点：意图识别
    workflow.add_node("save_state_snapshot", save_state_snapshot)
    workflow.add_node("classify_intent", classify_intent)
    workflow.add_node("handle_consult", handle_consult)
    workflow.add_node("handle_undo", handle_undo)
    workflow.add_node("handle_reset", handle_reset)

    # ---- 入口和前置流程 ----
    workflow.set_entry_point("load_session")
    workflow.add_edge("load_session", "process_input")
    workflow.add_edge("process_input", "manage_context")
    workflow.add_edge("manage_context", "save_state_snapshot")
    workflow.add_edge("save_state_snapshot", "classify_intent")

    # ---- 意图路由 ----
    workflow.add_conditional_edges(
        "classify_intent",
        route_by_intent,
        {
            "retrieve": "retrieve",
            "modify_jrxml": "modify_jrxml",
            "save_session": "save_session",
            "handle_consult": "handle_consult",
            "handle_undo": "handle_undo",
            "handle_reset": "handle_reset",
        },
    )

    # ---- 初始生成分支 ----
    workflow.add_edge("retrieve", "generate")
    workflow.add_conditional_edges(
        "generate",
        route_after_generate,
        {"save_session": "save_session"},
    )

    # ---- 修改分支 ----
    workflow.add_conditional_edges(
        "modify_jrxml",
        route_after_modify,
        {"save_session": "save_session"},
    )

    # ---- 撤销分支 ----
    workflow.add_conditional_edges(
        "handle_undo",
        route_after_undo,
        {"save_session": "save_session"},
    )

    # ---- 保存后进入验证 ----
    workflow.add_conditional_edges(
        "save_session",
        route_after_save,
        {"validate": "validate"},
    )

    # ---- 验证 → 修正循环 ----
    workflow.add_conditional_edges(
        "validate",
        route_after_validate,
        {"finalize": "finalize", "explain_error": "explain_error"},
    )
    workflow.add_conditional_edges(
        "explain_error",
        route_after_explain,
        {"correct_jrxml": "correct_jrxml"},
    )
    workflow.add_conditional_edges(
        "correct_jrxml",
        route_after_correct,
        {"validate": "validate", "finalize": "finalize"},
    )

    # ---- 咨询 / 重置 → 直接结束 ----
    workflow.add_edge("handle_consult", "finalize")
    workflow.add_edge("handle_reset", "finalize")

    # ---- 结束 ----
    workflow.add_edge("finalize", END)

    return workflow.compile()


# ============================================================
# 初始状态
# ============================================================

def create_initial_state() -> AgentState:
    return AgentState(
        conversation_history=[],
        current_jrxml="",
        user_input="",
        status="",
        error_msg="",
        natural_explanation="",
        retry_count=0,
        user_modification_request="",
        final_jrxml="",
        stage="initial_generation",
        retrieved_context="",
        full_conversation_history=[],
        compressed_history="",
        current_token_count=0,
        session_id="",
        session_name="",
        created_at="",
        updated_at="",
        intent="",
        history_states=[],
        jrxml_versions=[],
        last_error_case={},
        pending_failure_context={},
    )