agent_deply/pam_deploy_graph/langgraph_runtime.py

"""chat 人工确认点的 LangGraph interrupt 运行器。"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any
from uuid import uuid4

from .agent import PamDeployAgent
from .models import AgentState


@dataclass(slots=True)
class LangGraphRunResult:
    """一次 LangGraph 执行或恢复后的结果摘要。"""

    state: AgentState | None = None
    report: str = ""
    confirmation: dict[str, Any] = field(default_factory=dict)
    interrupted: bool = False
    chunks: list[dict[str, Any]] = field(default_factory=list)


class LangGraphDeploymentRuntime:
    """用 LangGraph interrupt/checkpointer 托管 chat 中的人工确认流程。"""

    def __init__(self, *, agent: PamDeployAgent, thread_id: str | None = None) -> None:
        """初始化图实例和会话线程 ID。"""
        self.agent = agent
        self.thread_id = thread_id or str(uuid4())
        self._waiting_confirmation = False
        self._graph = self._build_graph()

    @property
    def waiting_confirmation(self) -> bool:
        """返回当前 LangGraph 会话是否停在 interrupt 确认点。"""
        return self._waiting_confirmation

    def start(self, state: AgentState) -> LangGraphRunResult:
        """从给定 AgentState 开始执行，直到结束或遇到人工确认点。"""
        self._waiting_confirmation = False
        return self._consume(self._graph.stream({"agent_state": state}, self._config()))

    def resume(self, *, approved: bool, note: str = "") -> LangGraphRunResult:
        """把人工确认结果交回 LangGraph，并继续执行。"""
        try:
            from langgraph.types import Command
        except ImportError as exc:  # pragma: no cover - 依赖缺失时由调用方降级
            raise RuntimeError("未安装 langgraph，无法恢复 interrupt。") from exc

        decision = {"approved": approved, "note": note}
        return self._consume(self._graph.stream(Command(resume=decision), self._config()))

    def _build_graph(self):
        """构建 deploy -> confirm interrupt -> deploy 的循环图。"""
        try:
            from langgraph.checkpoint.memory import InMemorySaver
            from langgraph.graph import END, START, StateGraph
            from langgraph.types import interrupt
        except ImportError as exc:  # pragma: no cover - 依赖缺失时由调用方降级
            raise RuntimeError("未安装 langgraph，无法启用 chat interrupt。") from exc

        def deploy_node(state: dict[str, Any]) -> dict[str, Any]:
            """执行部署流，遇到 pending_confirmation 时由路由转入确认节点。"""
            agent_state = self.agent.run_deploy_flow(state["agent_state"])
            return {"agent_state": agent_state}

        def confirm_node(state: dict[str, Any]) -> dict[str, Any]:
            """把确认请求交给 LangGraph interrupt，并在恢复后执行确认动作。"""
            agent_state = state["agent_state"]
            request = self.agent.build_confirmation_request(agent_state)
            decision = interrupt(request)
            approved, note = _parse_confirmation_decision(decision)
            agent_state = self.agent.confirm_pending(
                agent_state,
                approved=approved,
                operator_note=note,
            )
            return {"agent_state": agent_state}

        def report_node(state: dict[str, Any]) -> dict[str, Any]:
            """渲染当前状态报告。"""
            return {"report": self.agent.render_report(state["agent_state"])}

        def route_after_deploy(state: dict[str, Any]) -> str:
            """根据是否存在 pending_confirmation 决定下一步。"""
            agent_state = state["agent_state"]
            return "confirm" if agent_state.pending_confirmation else "report"

        graph = StateGraph(dict)
        graph.add_node("deploy", deploy_node)
        graph.add_node("confirm", confirm_node)
        graph.add_node("report", report_node)
        graph.add_edge(START, "deploy")
        graph.add_conditional_edges(
            "deploy",
            route_after_deploy,
            {"confirm": "confirm", "report": "report"},
        )
        graph.add_edge("confirm", "deploy")
        graph.add_edge("report", END)
        return graph.compile(checkpointer=InMemorySaver())

    def _config(self) -> dict[str, Any]:
        """生成 LangGraph checkpointer 使用的线程配置。"""
        return {"configurable": {"thread_id": self.thread_id}}

    def _consume(self, chunks: Any) -> LangGraphRunResult:
        """消费 LangGraph stream 输出，提取状态、报告和 interrupt 请求。"""
        result = LangGraphRunResult()
        for chunk in chunks:
            result.chunks.append(chunk)
            if "__interrupt__" in chunk:
                result.interrupted = True
                result.confirmation = _extract_interrupt_value(chunk["__interrupt__"])
                continue

            for value in chunk.values():
                if not isinstance(value, dict):
                    continue
                if isinstance(value.get("agent_state"), AgentState):
                    result.state = value["agent_state"]
                if isinstance(value.get("report"), str):
                    result.report = value["report"]

        self._waiting_confirmation = result.interrupted
        return result


def _extract_interrupt_value(interrupts: Any) -> dict[str, Any]:
    """从 LangGraph interrupt 对象中提取确认请求字典。"""
    if not interrupts:
        return {}
    first = interrupts[0]
    value = getattr(first, "value", first)
    return value if isinstance(value, dict) else {"value": value}


def _parse_confirmation_decision(value: Any) -> tuple[bool, str]:
    """把 interrupt resume 值解析为 approved/note。"""
    if isinstance(value, dict):
        return bool(value.get("approved", False)), str(value.get("note", ""))
    if isinstance(value, bool):
        return value, ""
    if isinstance(value, str):
        normalized = value.strip().lower()
        return normalized in ("approve", "approved", "yes", "y", "true"), value
    return False, str(value)