SupplyAgent/find_agent_v2/graph.py

"""One complete business round as a deterministic DAG.

The project does not depend on LangGraph, so this module preserves the same boundary
without adding a second Agent runtime: one invocation is one acyclic business round.
"""

from __future__ import annotations

from typing import Protocol

from find_agent_v2.context import build_node_slots, render_node_context
from find_agent_v2.observability import InputSlot, NullObserver
from find_agent_v2.prompts import EVALUATOR_PROMPT, EVIDENCE_PROMPT, PLANNER_PROMPT, SEARCH_PROMPT
from find_agent_v2.service import FindAgentV2Service
from find_agent_v2.state import FindAgentState, NodeRun
from find_agent_v2.tools import EVALUATION_TOOLS, EVIDENCE_TOOLS, SEARCH_TOOLS, ToolFn


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, ...] = (),
    ) -> NodeRun: ...


class FindAgentRoundGraph:
    """planner -> search -> evidence -> evaluator -> END."""

    def __init__(
        self, *, service: FindAgentV2Service, runner: NodeRunner,
        observer=None,
    ) -> None:
        self.service = service
        self.runner = runner
        self.observer = observer or NullObserver()

    def _full_state(self, state: FindAgentState, *, pending_only: bool = False):
        return self.service.get_full_state(
            state.run_id, pending_only=pending_only,
        )

    def _context(self, state: FindAgentState, *, 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.plan,
        )

    def _slots(
        self, state: FindAgentState, *, 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.plan,
        )

    async def invoke(self, state: FindAgentState) -> FindAgentState:
        with self.observer.round(round_index=state.round_index) as round_observation:
            state = await self._invoke_nodes(state)
            round_observation.set_output({
                "状态快照": state.snapshot.__dict__ if state.snapshot else {},
                "本轮计划": state.plan,
            }, ok=True)
        return state

    async def _invoke_nodes(self, state: FindAgentState) -> FindAgentState:
        state.phase = "planning"
        plan = await self.runner.run_node(
            node="planner",
            round_index=state.round_index,
            system_prompt=PLANNER_PROMPT,
            user_content=self._context(state),
            tools=(),
            max_iterations=2,
            slots=self._slots(state),
        )
        state.node_runs.append(plan)
        state.plan = plan.content.strip()
        self.service.update_round(
            state.run_id, state.round_index, phase="searching", plan=state.plan,
        )

        state.phase = "searching"
        search = 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),
        )
        state.node_runs.append(search)
        after_search = self.service.snapshot(state.run_id)

        if after_search.pending_count:
            state.phase = "evidence"
            self.service.update_round(state.run_id, state.round_index, phase="evidence")
            evidence = await self.runner.run_node(
                node="evidence",
                round_index=state.round_index,
                system_prompt=EVIDENCE_PROMPT,
                user_content=self._context(state, pending_only=True),
                tools=EVIDENCE_TOOLS,
                max_iterations=12,
                slots=self._slots(state, pending_only=True),
            )
            state.node_runs.append(evidence)

            state.phase = "evaluating"
            self.service.update_round(state.run_id, state.round_index, phase="evaluating")
            # A model may intentionally keep one tool payload small (for example,
            # evaluate ten candidates at a time).  One evaluator invocation is
            # therefore not proof that the queue has been drained.  Re-enter the
            # node with a fresh DB snapshot until every candidate has a terminal
            # bucket, while failing fast if an invocation makes no progress.
            stagnant_attempts = 0
            for _ in range(64):
                before_evaluation = self.service.snapshot(state.run_id)
                if not before_evaluation.pending_count:
                    break
                evaluation = await self.runner.run_node(
                    node="evaluator",
                    round_index=state.round_index,
                    system_prompt=EVALUATOR_PROMPT,
                    user_content=self._context(state, pending_only=True),
                    tools=EVALUATION_TOOLS,
                    max_iterations=12,
                    slots=self._slots(state, pending_only=True),
                )
                state.node_runs.append(evaluation)
                after_evaluation = self.service.snapshot(state.run_id)
                if after_evaluation.pending_count >= before_evaluation.pending_count:
                    stagnant_attempts += 1
                    if stagnant_attempts >= 3:
                        raise RuntimeError(
                            "评估节点连续 3 次未消费 pending_evaluation 候选："
                            f"remaining={after_evaluation.pending_count}"
                        )
                else:
                    stagnant_attempts = 0
            else:
                raise RuntimeError("评估批次超过安全上限 64")

        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,
        )
        return state