agent_jrxml/agent/graph.py

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

import functools
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,
)
from backend.logger import get_logger

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

_graph_log = get_logger("agent")


def _log_route(route_name: str):
    """装饰器：自动记录路由决策。"""
    def decorator(func):
        @functools.wraps(func)
        def wrapper(state: AgentState, *args, **kwargs):
            target = func(state, *args, **kwargs)
            _graph_log.info(
                f"[路由] {route_name} → {target}",
                extra={
                    "route": route_name,
                    "target": target,
                    "session_id": state.get("session_id", ""),
                    "intent": state.get("intent", ""),
                    "status": state.get("status", ""),
                    "has_jrxml": bool(state.get("current_jrxml", "").strip()),
                    "retry_count": state.get("retry_count", 0),
                },
            )
            return target
        return wrapper
    return decorator

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

@_log_route("route_by_intent")
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"


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


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


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


@_log_route("route_after_save")
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"


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


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


@_log_route("route_after_correct")
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={},
    )