SupplyAgent/find_agent_v2/graph.py

"""One autonomous but policy-guarded business round compiled as LangGraph."""

from __future__ import annotations

import asyncio
import json
import re
from typing import Any, Protocol

from langgraph.graph import END, START, StateGraph

from find_agent_v2.context import build_node_slots, render_node_context
from find_agent_v2.observability import InputSlot, NullObserver, graph_spec_for
from find_agent_v2.prompts import EVALUATOR_PROMPT, EVIDENCE_PROMPT, SEARCH_PROMPT, SUPERVISOR_PROMPT
from find_agent_v2.service import FindAgentV2Service
from find_agent_v2.state import FindAgentGraphState, FindAgentState, NodeRun
from find_agent_v2.tools import (
    EVALUATION_TOOLS,
    EVIDENCE_TOOLS,
    SEARCH_TOOLS,
    ToolFn,
    bound_candidate_tools,
)


class NodeRunner(Protocol):
    async def run_node(
        self,
        *,
        node: str,
        round_index: int,
        system_prompt: str,
        user_content: str,
        tools: tuple[ToolFn, ...] = (),
        max_iterations: int = 12,
        slots: tuple[InputSlot, ...] = (),
        branch_key: str = "",
        allow_delegation: bool = True,
    ) -> NodeRun: ...


class FindAgentRoundGraph:
    """Supervisor-directed graph with deterministic policy approval."""

    def __init__(
        self, *, service: FindAgentV2Service, runner: NodeRunner, observer=None,
        max_actions: int = 16, max_search_actions: int = 3,
        target_primary_count: int = 5,
    ) -> None:
        self.service = service
        self.runner = runner
        self.observer = observer or NullObserver()
        self.max_actions = max(4, int(max_actions))
        self.max_search_actions = max(1, int(max_search_actions))
        self.target_primary_count = max(1, int(target_primary_count))
        self.app = self._build_graph()
        self.obagent_spec = graph_spec_for(self.app)

    def _full_state(self, state: FindAgentGraphState, *, pending_only: bool = False):
        return self.service.get_full_state(
            state["run_id"], pending_only=pending_only,
        )

    def _context(self, state: FindAgentGraphState, *, pending_only: bool = False) -> str:
        return render_node_context(
            user_input=state["user_input"],
            full_state=self._full_state(state, pending_only=pending_only),
            round_index=state["round_index"],
            plan=state.get("plan", ""),
        )

    def _slots(
        self, state: FindAgentGraphState, *, pending_only: bool = False,
    ) -> tuple[InputSlot, ...]:
        return build_node_slots(
            user_input=state["user_input"],
            full_state=self._full_state(state, pending_only=pending_only),
            round_index=state["round_index"],
            plan=state.get("plan", ""),
        )

    def _shard_state(
        self, state: FindAgentGraphState, items: list[dict[str, Any]],
    ) -> dict[str, Any]:
        full_state = self._full_state(state, pending_only=True)
        return {**full_state, "candidates": items}

    def _shard_context(
        self, state: FindAgentGraphState, items: list[dict[str, Any]],
    ) -> str:
        return render_node_context(
            user_input=state["user_input"],
            full_state=self._shard_state(state, items),
            round_index=state["round_index"],
            plan=state.get("plan", ""),
        )

    def _shard_slots(
        self, state: FindAgentGraphState, items: list[dict[str, Any]],
    ) -> tuple[InputSlot, ...]:
        return build_node_slots(
            user_input=state["user_input"],
            full_state=self._shard_state(state, items),
            round_index=state["round_index"],
            plan=state.get("plan", ""),
        )

    @staticmethod
    def _parse_supervisor(content: str) -> dict[str, Any]:
        text = content.strip()
        fenced = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.S)
        if fenced:
            text = fenced.group(1)
        else:
            start, end = text.find("{"), text.rfind("}")
            if start >= 0 and end > start:
                text = text[start:end + 1]
        try:
            value = json.loads(text)
            return value if isinstance(value, dict) else {}
        except (TypeError, ValueError):
            return {}

    def _approve_action(
        self, state: FindAgentGraphState, proposal: dict[str, Any],
    ) -> tuple[str, str, int, str]:
        """Turn an LLM proposal into a safe, executable transition."""
        pending = self._full_state(state, pending_only=True).get("candidates") or []
        missing_detail = any(item.get("detail_status") == "pending" for item in pending)
        missing_portrait = any(item.get("portrait_status") == "pending" for item in pending)
        proposed = str(proposal.get("next_action") or "").lower()
        reason = str(proposal.get("reason") or "")
        searches = int(state.get("search_actions") or 0)
        actions = int(state.get("action_count") or 0)
        try:
            worker_count = max(1, min(8, int(proposal.get("worker_count") or 4)))
        except (TypeError, ValueError):
            worker_count = 4
        scope = str(proposal.get("evidence_scope") or "both").lower()
        if scope not in {"detail", "portrait", "both"}:
            scope = "both"

        snapshot = self.service.snapshot(state["run_id"])
        if not pending and snapshot.valid_primary_count >= self.target_primary_count:
            return "finish", "有效 primary 已达到目标，结束本轮", worker_count, scope

        if actions > self.max_actions:
            if not pending:
                return "finish", "安全收敛动作已完成", worker_count, scope
            raise RuntimeError(
                f"Supervisor 安全收敛动作未产生进展：actions={actions}, pending={len(pending)}"
            )
        if actions == self.max_actions:
            if pending and not (missing_detail or missing_portrait):
                return "evaluator", "动作预算耗尽，强制消费待评估候选", worker_count, scope
            if pending:
                return "evidence", "动作预算耗尽，强制补齐缺失证据", worker_count, "both"
            return "finish", "达到单轮动作安全上限", worker_count, scope
        if pending:
            if missing_detail or missing_portrait:
                if scope == "detail" and not missing_detail:
                    scope = "portrait"
                elif scope == "portrait" and not missing_portrait:
                    scope = "detail"
                return "evidence", reason or "存在缺失证据，框架要求先补证", worker_count, scope
            return "evaluator", reason or "候选证据已处理，进入评估", worker_count, scope
        if searches == 0:
            return "search", reason or "本轮尚未搜索，框架要求先建立候选池", worker_count, scope
        if proposed == "search" and searches < self.max_search_actions:
            return "search", reason or "Supervisor 判断继续搜索仍有信息增益", worker_count, scope
        return "finish", reason or "没有待处理候选，结束本轮", worker_count, scope

    async def _supervisor(self, state: FindAgentGraphState) -> dict[str, Any]:
        run = await self.runner.run_node(
            node="supervisor",
            round_index=state["round_index"],
            system_prompt=SUPERVISOR_PROMPT,
            user_content=self._context(state),
            tools=(),
            max_iterations=2,
            slots=self._slots(state),
            allow_delegation=False,
        )
        proposal = self._parse_supervisor(run.content)
        action, reason, workers, scope = self._approve_action(state, proposal)
        proposed_plan = proposal.get("plan")
        plan = (
            json.dumps(proposed_plan, ensure_ascii=False)
            if isinstance(proposed_plan, dict)
            else state.get("plan", "")
        )
        decision = {
            "step": int(state.get("supervisor_step") or 0) + 1,
            "proposed_action": proposal.get("next_action"),
            "approved_action": action,
            "reason": reason,
            "worker_count": workers,
            "evidence_scope": scope,
        }
        self.service.update_round(
            state["run_id"], state["round_index"], phase="planning", plan=plan,
        )
        return {
            "phase": "planning", "plan": plan, "approved_action": action,
            "supervisor_step": decision["step"], "worker_count": workers,
            "evidence_scope": scope,
            "decision_history": [*state.get("decision_history", []), decision],
            "node_runs": [*state.get("node_runs", []), run],
        }

    async def _search(self, state: FindAgentGraphState) -> dict[str, Any]:
        self.service.update_round(state["run_id"], state["round_index"], phase="searching")
        run = await self.runner.run_node(
            node="search",
            round_index=state["round_index"],
            system_prompt=SEARCH_PROMPT,
            user_content=self._context(state),
            tools=SEARCH_TOOLS,
            max_iterations=10,
            slots=self._slots(state),
        )
        return {
            "phase": "searching",
            "search_actions": int(state.get("search_actions") or 0) + 1,
            "action_count": int(state.get("action_count") or 0) + 1,
            "node_runs": [*state.get("node_runs", []), run],
        }

    async def _evidence(self, state: FindAgentGraphState) -> dict[str, Any]:
        self.service.update_round(state["run_id"], state["round_index"], phase="evidence")
        pending = self._full_state(state, pending_only=True)["candidates"]
        detail_items = [item for item in pending if item.get("detail_status") == "pending"]
        portrait_items = [item for item in pending if item.get("portrait_status") == "pending"]
        scope = state.get("evidence_scope", "both")
        if scope == "detail":
            portrait_items = []
        elif scope == "portrait":
            detail_items = []
        jobs = [
            ("detail", detail_items[index:index + 8])
            for index in range(0, len(detail_items), 8)
        ] + [
            ("portrait", portrait_items[index:index + 8])
            for index in range(0, len(portrait_items), 8)
        ]

        semaphore = asyncio.Semaphore(max(1, min(8, int(state.get("worker_count") or 4))))

        async def run_shard(
            index: int, evidence_type: str, items: list[dict[str, Any]],
        ) -> NodeRun:
            candidate_ids = [int(item["candidate_id"]) for item in items]
            selected = (
                (EVIDENCE_TOOLS[0], EVIDENCE_TOOLS[2])
                if evidence_type == "detail"
                else (EVIDENCE_TOOLS[1], EVIDENCE_TOOLS[2])
            )
            async with semaphore:
                return await self.runner.run_node(
                    node="evidence",
                round_index=state["round_index"],
                system_prompt=EVIDENCE_PROMPT,
                user_content=(
                    f"你只负责当前 {evidence_type} 分片 candidate_ids={candidate_ids}。"
                    f"必须为这些候选补齐 {evidence_type}，不得访问其他候选。\n\n"
                    + self._shard_context(state, items)
                ),
                tools=bound_candidate_tools(
                    selected, run_id=state["run_id"], candidate_ids=candidate_ids,
                ),
                max_iterations=12,
                slots=self._shard_slots(state, items),
                branch_key=f"{evidence_type}-shard-{index}",
                allow_delegation=False,
                )

        runs = await asyncio.gather(*(
            run_shard(index, evidence_type, items)
            for index, (evidence_type, items) in enumerate(jobs, start=1)
        )) if jobs else []
        return {
            "phase": "evidence",
            "action_count": int(state.get("action_count") or 0) + 1,
            "node_runs": [*state.get("node_runs", []), *runs],
        }

    async def _evaluator(self, state: FindAgentGraphState) -> dict[str, Any]:
        self.service.update_round(state["run_id"], state["round_index"], phase="evaluating")
        node_runs = list(state.get("node_runs", []))
        before = self.service.snapshot(state["run_id"])
        pending = self._full_state(state, pending_only=True)["candidates"]
        shards = [pending[index:index + 8] for index in range(0, len(pending), 8)]
        semaphore = asyncio.Semaphore(max(1, min(8, int(state.get("worker_count") or 4))))

        async def run_shard(index: int, items: list[dict[str, Any]]) -> NodeRun:
            candidate_ids = [int(item["candidate_id"]) for item in items]
            async with semaphore:
                return await self.runner.run_node(
                    node="evaluator",
                    round_index=state["round_index"],
                    system_prompt=EVALUATOR_PROMPT,
                    user_content=(
                        f"你只负责当前分片 candidate_ids={candidate_ids}。"
                        "必须把这些候选全部分池，不得评估其他候选。\n\n"
                        + self._shard_context(state, items)
                    ),
                    tools=bound_candidate_tools(
                        EVALUATION_TOOLS, run_id=state["run_id"], candidate_ids=candidate_ids,
                    ),
                    max_iterations=12,
                    slots=self._shard_slots(state, items),
                    branch_key=f"step-{state.get('supervisor_step', 0)}-shard-{index}",
                    allow_delegation=False,
                )

        runs = await asyncio.gather(*(
            run_shard(index, items) for index, items in enumerate(shards, start=1)
        )) if shards else []
        node_runs.extend(runs)
        self.service.recount_valid_primary(state["run_id"])
        after = self.service.snapshot(state["run_id"])
        stagnant = int(state.get("evaluator_stagnation") or 0)
        stagnant = stagnant + 1 if after.pending_count >= before.pending_count else 0
        if stagnant >= 3:
            raise RuntimeError(
                "评估节点连续 3 次未消费 pending_evaluation 候选："
                f"remaining={after.pending_count}"
            )
        return {
            "phase": "evaluating", "node_runs": node_runs,
            "action_count": int(state.get("action_count") or 0) + 1,
            "evaluator_stagnation": stagnant,
        }

    @staticmethod
    def _route(state: FindAgentGraphState) -> str:
        return state.get("approved_action", "finish")

    def _build_graph(self):
        builder = StateGraph(FindAgentGraphState)
        builder.add_node("supervisor", self._supervisor)
        builder.add_node("search", self._search)
        builder.add_node("evidence", self._evidence)
        builder.add_node("evaluator", self._evaluator)
        builder.add_edge(START, "supervisor")
        builder.add_conditional_edges("supervisor", self._route, {
            "search": "search", "evidence": "evidence",
            "evaluator": "evaluator", "finish": END,
        })
        builder.add_edge("search", "supervisor")
        builder.add_edge("evidence", "supervisor")
        builder.add_edge("evaluator", "supervisor")
        return builder.compile()

    async def invoke(self, state: FindAgentState) -> FindAgentState:
        graph_state: FindAgentGraphState = {
            "run_id": state.run_id,
            "user_input": state.user_input,
            "round_index": state.round_index,
            "plan": state.plan,
            "phase": state.phase,
            "node_runs": [],
            "snapshot": state.snapshot,
            "supervisor_step": 0,
            "action_count": 0,
            "search_actions": 0,
            "worker_count": 4,
            "evidence_scope": "both",
            "decision_history": [],
            "evaluator_stagnation": 0,
        }
        with self.observer.round(
            round_index=state.round_index, spec=self.obagent_spec,
        ) as round_observation:
            output = await self.app.ainvoke(
                graph_state, config={"recursion_limit": max(64, self.max_actions * 4)},
            )
            state.plan = str(output.get("plan") or "")
            state.node_runs.extend(output.get("node_runs") or [])
            state.snapshot = self.service.snapshot(state.run_id)
            state.phase = "done"
            self.service.update_round(
                state.run_id,
                state.round_index,
                phase="done",
                status="done",
                snapshot=state.snapshot,
            )
            round_observation.set_output({
                "状态快照": state.snapshot.__dict__,
                "本轮计划": state.plan,
                "Supervisor决策轨迹": output.get("decision_history") or [],
            }, ok=True)
        return state