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 assignment_slots, render_assignment
from find_agent_v2.gates import evaluate_candidate_gate
from find_agent_v2.observability import InputSlot, NullObserver, graph_spec_for
from find_agent_v2.prompts import (
    EVALUATOR_PROMPT,
    EVIDENCE_PROMPT,
    PLANNER_PROMPT,
    SEARCH_PROMPT,
    SUPERVISOR_PROMPT,
)
from find_agent_v2.service import FindAgentV2Service
from find_agent_v2.state import (
    EvidenceAssignment,
    EvaluationAssignment,
    FindAgentGraphState,
    FindAgentState,
    NodeRun,
    PlannerAssignment,
    PlanningDecision,
    SearchAssignment,
    SupervisorAssignment,
    SupervisorDecision,
)
from find_agent_v2.tools import (
    EVALUATION_TOOLS,
    EVIDENCE_TOOLS,
    SEARCH_TOOLS,
    ToolFn,
    bound_candidate_tools,
    normalize_evaluation_items,
)


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,
    ) -> 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.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,
        )

    @staticmethod
    def _assignment_input(assignment, *, source: str) -> tuple[str, tuple[InputSlot, ...]]:
        return render_assignment(assignment), assignment_slots(assignment, source=source)

    def _supervisor_state(self, state: FindAgentGraphState) -> dict[str, Any]:
        """Compact progress projection; routing never needs full candidate payloads."""
        full_state = self._full_state(state)
        run = full_state.get("run") or {}
        candidates = list(full_state.get("candidates") or [])
        pending_candidates = [
            item
            for item in candidates
            if item.get("decision_bucket") == "pending_evaluation"
        ]
        return {
            "run": {
                key: run.get(key)
                for key in (
                    "status", "outcome_status", "current_round", "search_count",
                    "candidate_count", "valid_primary_count",
                )
            },
            "searches": full_state.get("searches") or [],
            "candidate_progress": {
                "pending_count": sum(
                    1 for _item in pending_candidates
                ),
                "primary_count": sum(
                    item.get("decision_bucket") == "primary" for item in candidates
                ),
                "rejected_count": sum(
                    item.get("decision_bucket") == "rejected" for item in candidates
                ),
                "detail_pending_count": sum(
                    item.get("detail_status") == "pending"
                    for item in pending_candidates
                ),
                "detail_success_count": sum(
                    item.get("detail_status") == "success"
                    for item in pending_candidates
                ),
                "detail_failed_count": sum(
                    item.get("detail_status") == "failed"
                    for item in pending_candidates
                ),
                "portrait_pending_count": sum(
                    item.get("portrait_status") == "pending"
                    for item in pending_candidates
                ),
                "portrait_success_count": sum(
                    item.get("portrait_status") == "success"
                    for item in pending_candidates
                ),
                "portrait_failed_count": sum(
                    item.get("portrait_status") == "failed"
                    for item in pending_candidates
                ),
                "evidence_completed_count": sum(
                    item.get("detail_status") != "pending"
                    and item.get("portrait_status") != "pending"
                    for item in pending_candidates
                ),
                "evidence_success_count": sum(
                    item.get("detail_status") == "success"
                    and item.get("portrait_status") == "success"
                    for item in pending_candidates
                ),
            },
        }

    def _video_understanding_ids(
        self,
        state: FindAgentGraphState,
        items: list[dict[str, Any]],
    ) -> list[int]:
        """Only candidates passing every deterministic hard gate may use video understanding."""
        run = self._full_state(state).get("run") or {}
        rules = run.get("rule_config") or {}
        selected: list[int] = []
        for item in items:
            gate = evaluate_candidate_gate(item, rules)
            checks = list(gate.get("checks") or [])
            hard_gate_passed = gate.get("status") == "pass" and all(
                check.get("status") == "pass" and not check.get("compensated") for check in checks
            )
            if str(item.get("video_url") or "").strip() and hard_gate_passed:
                selected.append(int(item["candidate_id"]))
        return selected

    @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"

        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]:
        execution_plan = state.get("execution_plan")
        observed_decision: dict[str, Any] = {}
        if execution_plan is None:
            assignment = PlannerAssignment(
                run_id=state["run_id"],
                round_index=state["round_index"],
                raw_demand=state["user_input"],
            )
            user_content, slots = self._assignment_input(
                assignment,
                source="find_agent_v2_run.input_json",
            )
            structured_runner = getattr(self.runner, "run_planning", None)
            if callable(structured_runner):
                run, planning = await structured_runner(
                    round_index=state["round_index"],
                    system_prompt=PLANNER_PROMPT,
                    user_content=user_content,
                    slots=slots,
                )
            else:
                run = await self.runner.run_node(
                    node="supervisor",
                    round_index=state["round_index"],
                    system_prompt=PLANNER_PROMPT,
                    user_content=user_content,
                    tools=(),
                    max_iterations=2,
                    slots=slots,
                    allow_delegation=False,
                )
                planning = PlanningDecision.model_validate(self._parse_supervisor(run.content))
            execution_plan = planning.execution_plan
            proposal = planning.model_dump(mode="json", exclude={"execution_plan"})
        else:
            assignment = SupervisorAssignment(
                run_id=state["run_id"],
                round_index=state["round_index"],
                execution_plan=execution_plan,
                execution_state=self._supervisor_state(state),
            )
            user_content, slots = self._assignment_input(
                assignment,
                source="ExecutionPlan + FindAgentV2Service.get_full_state",
            )
            structured_runner = getattr(self.runner, "run_supervision", None)
            def approve_for_observation(supervision: SupervisorDecision) -> dict[str, Any]:
                proposal_payload = supervision.model_dump(mode="json")
                approved, approved_reason, approved_workers, approved_scope = (
                    self._approve_action(state, proposal_payload)
                )
                decision_payload = {
                    "step": int(state.get("supervisor_step") or 0) + 1,
                    "proposed_action": proposal_payload.get("next_action"),
                    "approved_action": approved,
                    "overridden": proposal_payload.get("next_action") != approved,
                    "reason": approved_reason,
                    "worker_count": approved_workers,
                    "evidence_scope": approved_scope,
                }
                observed_decision.update(decision_payload)
                return {"Supervisor决策校验": decision_payload}

            if callable(structured_runner):
                run, supervision = await structured_runner(
                    round_index=state["round_index"],
                    system_prompt=SUPERVISOR_PROMPT,
                    user_content=user_content,
                    slots=slots,
                    output_enricher=approve_for_observation,
                )
            else:
                run = await self.runner.run_node(
                    node="supervisor",
                    round_index=state["round_index"],
                    system_prompt=SUPERVISOR_PROMPT,
                    user_content=user_content,
                    tools=(),
                    max_iterations=2,
                    slots=slots,
                    allow_delegation=False,
                )
                supervision = SupervisorDecision.model_validate(
                    self._parse_supervisor(run.content),
                )
            proposal = supervision.model_dump(mode="json")
            if supervision.additional_search_tasks:
                known = {item.task_id for item in execution_plan.search_tasks}
                merged = list(execution_plan.search_tasks)
                for item in supervision.additional_search_tasks:
                    if item.task_id not in known and len(merged) < 24:
                        merged.append(item)
                        known.add(item.task_id)
                if len(merged) != len(execution_plan.search_tasks):
                    execution_plan = type(execution_plan).model_validate(
                        {
                            **execution_plan.model_dump(),
                            "search_tasks": merged,
                        }
                    )
        if observed_decision:
            decision = observed_decision
            action = str(decision["approved_action"])
            reason = str(decision["reason"])
            workers = int(decision["worker_count"])
            scope = str(decision["evidence_scope"])
        else:
            action, reason, workers, scope = self._approve_action(state, proposal)
            decision = {
                "step": int(state.get("supervisor_step") or 0) + 1,
                "proposed_action": proposal.get("next_action"),
                "approved_action": action,
                "overridden": proposal.get("next_action") != action,
                "reason": reason,
                "worker_count": workers,
                "evidence_scope": scope,
            }
        self.service.update_round(
            state["run_id"],
            state["round_index"],
            phase="planning",
            plan=execution_plan.model_dump_json(exclude_none=True),
        )
        return {
            "phase": "planning",
            "execution_plan": execution_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")
        execution_plan = state.get("execution_plan")
        if execution_plan is None:
            raise RuntimeError("Search 缺少已校验的 ExecutionPlan")
        existing = {
            (
                str(item.get("keyword") or ""),
                str(item.get("provider") or ""),
                str(item.get("query_reason") or ""),
            )
            for item in self._full_state(state).get("searches") or []
        }
        tasks = [
            item
            for item in execution_plan.search_tasks
            if (item.keyword, item.provider, item.query_reason) not in existing
        ][:6]
        assignment = SearchAssignment(
            run_id=state["run_id"],
            round_index=state["round_index"],
            tasks=tasks,
        )
        if not tasks:
            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", []),
                    NodeRun("search", state["round_index"], "没有未执行的搜索任务", 0, 0),
                ],
            }
        user_content, slots = self._assignment_input(
            assignment,
            source="ExecutionPlan.search_tasks",
        )
        run = await self.runner.run_node(
            node="search",
            round_index=state["round_index"],
            system_prompt=SEARCH_PROMPT,
            user_content=user_content,
            tools=SEARCH_TOOLS,
            max_iterations=10,
            slots=slots,
        )
        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]
            assignment = EvidenceAssignment(
                run_id=state["run_id"],
                round_index=state["round_index"],
                candidate_ids=candidate_ids,
                evidence_type=evidence_type,
                candidates=items,
            )
            user_content, slots = self._assignment_input(
                assignment,
                source="host evidence shard",
            )
            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=user_content,
                    tools=bound_candidate_tools(
                        selected,
                        run_id=state["run_id"],
                        candidate_ids=candidate_ids,
                    ),
                    max_iterations=12,
                    slots=slots,
                    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"]
        run = self._full_state(state).get("run") or {}
        rules = run.get("rule_config") or {}
        eligible: list[dict[str, Any]] = []
        gate_failures: list[tuple[int, dict[str, Any]]] = []
        for item in pending:
            gate = evaluate_candidate_gate(item, rules)
            if gate.get("status") == "pass":
                eligible.append(item)
            else:
                gate_failures.append((int(item["candidate_id"]), gate))
        self.service.reject_failed_gates(state["run_id"], gate_failures)
        shards = [eligible[index : index + 8] for index in range(0, len(eligible), 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]
            video_understanding_ids = self._video_understanding_ids(state, items)
            execution_plan = state.get("execution_plan")
            if execution_plan is None:
                raise RuntimeError("Evaluator 缺少已校验的 ExecutionPlan")
            assignment = EvaluationAssignment(
                run_id=state["run_id"],
                round_index=state["round_index"],
                candidate_ids=candidate_ids,
                evaluation_brief=execution_plan.evaluation_brief,
                quality_gate_rules=rules,
                current_datetime=str(rules.get("current_datetime") or ""),
                timezone=str(rules.get("timezone") or ""),
                video_understanding_candidate_ids=video_understanding_ids,
                candidates=items,
            )
            user_content, slots = self._assignment_input(
                assignment,
                source="ExecutionPlan.evaluation_brief + host candidate shard",
            )
            video_tools = (
                bound_candidate_tools(
                    (EVALUATION_TOOLS[0],),
                    run_id=state["run_id"],
                    candidate_ids=video_understanding_ids,
                )
                if video_understanding_ids
                else ()
            )
            async with semaphore:
                structured_runner = getattr(self.runner, "run_evaluation", None)
                if callable(structured_runner):
                    run, proposed = await structured_runner(
                        round_index=state["round_index"],
                        system_prompt=EVALUATOR_PROMPT,
                        user_content=user_content,
                        slots=slots,
                        branch_key=f"step-{state.get('supervisor_step', 0)}-shard-{index}",
                        tools=video_tools,
                    )
                    normalized = normalize_evaluation_items(
                        proposed,
                        allowed_candidates=items,
                    )
                    updated = self.service.evaluate(state["run_id"], normalized)
                    updated_ids = {int(item["candidate_id"]) for item in updated}
                    if updated_ids != set(candidate_ids):
                        raise RuntimeError(
                            "评估写入结果不完整："
                            f"expected={candidate_ids}, updated={sorted(updated_ids)}"
                        )
                    return run
                return await self.runner.run_node(
                    node="evaluator",
                    round_index=state["round_index"],
                    system_prompt=EVALUATOR_PROMPT,
                    user_content=user_content,
                    tools=(
                        *video_tools,
                        *bound_candidate_tools(
                            EVALUATION_TOOLS[1:],
                            run_id=state["run_id"],
                            candidate_ids=candidate_ids,
                        ),
                    ),
                    max_iterations=12,
                    slots=slots,
                    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(
                f"评估节点连续 3 次未消费 pending_evaluation 候选：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,
            "execution_plan": state.execution_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.execution_plan = output.get("execution_plan")
            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.execution_plan.model_dump(mode="json")
                        if state.execution_plan is not None
                        else {}
                    ),
                    "Supervisor决策轨迹": output.get("decision_history") or [],
                },
                ok=True,
            )
        return state