SupplyAgent/find_agent_v2/agent.py

"""Outer round orchestrator for the isolated find agent."""

from __future__ import annotations

import json

from find_agent_v2.context import build_node_slots
from find_agent_v2.graph import FindAgentRoundGraph, NodeRunner
from find_agent_v2.observability import ObagentObserver
from find_agent_v2.prompts import REPORT_PROMPT
from find_agent_v2.runtime import FindAgentNodeHost, normalize_models
from find_agent_v2.service import FindAgentV2Service, get_find_agent_v2_service
from find_agent_v2.state import FindAgentResult, FindAgentState
from find_agent_v2.tools import REPORT_TOOLS
from supply_agent.config import Settings


class FindAgentV2:
    """Python outer loop + one-round DAG + node-local ReAct."""

    def __init__(
        self,
        *,
        service: FindAgentV2Service | None = None,
        node_runner: NodeRunner | None = None,
        settings: Settings | None = None,
        models_by_role: dict[str, str] | None = None,
        max_rounds: int = 2,
        target_primary_count: int = 5,
        observer: ObagentObserver | None = None,
    ) -> None:
        self.service = service or get_find_agent_v2_service()
        self.observer = observer or ObagentObserver()
        self.node_runner = node_runner or FindAgentNodeHost(
            settings=settings,
            models_by_role=models_by_role,
            observer=self.observer,
        )
        self.models_by_role = dict(models_by_role or {})
        self.max_rounds = max(1, int(max_rounds))
        self.target_primary_count = max(1, int(target_primary_count))

    async def arun(self, *, run_id: str, user_input: str) -> FindAgentResult:
        run = self.service.require_run(run_id)
        if str(run.get("status") or "") != "running":
            raise ValueError(f"run_id={run_id} 当前状态不可执行: {run.get('status')}")
        state = FindAgentState(run_id=run_id, user_input=user_input)
        graph = FindAgentRoundGraph(
            service=self.service, runner=self.node_runner, observer=self.observer,
        )
        default_model = self.models_by_role.get("planner", "google/gemini-3-flash-preview")
        with self.observer.run(
            run_id=run_id,
            demand_word=str(run.get("demand_word") or ""),
            model=default_model,
            models_by_role=self.models_by_role,
        ) as observation_run:
            self.service.set_obagent_run_uid(
                run_id, getattr(observation_run, "run_uid", None),
            )
            result = await self._arun_inner(state=state, graph=graph)
            observation_run.finish(final_output=result.final_output)
            return result

    async def _arun_inner(
        self, *, state: FindAgentState, graph: FindAgentRoundGraph,
    ) -> FindAgentResult:
        run_id = state.run_id
        end_reason = ""
        failed = False
        try:
            for round_index in range(1, self.max_rounds + 1):
                state.round_index = round_index
                state.previous_snapshot = self.service.snapshot(run_id)
                self.service.begin_round(run_id, round_index, state.previous_snapshot)
                await graph.invoke(state)
                current = state.snapshot or self.service.snapshot(run_id)
                if current.valid_primary_count >= self.target_primary_count:
                    end_reason = f"已获得 {current.valid_primary_count} 条有效 primary"
                    break
                if current.pending_count:
                    failed = True
                    end_reason = f"第 {round_index} 轮结束仍有 {current.pending_count} 条待评估候选"
                    break
                if current.candidate_count <= state.previous_snapshot.candidate_count:
                    end_reason = "本轮未发现新增候选，搜索前沿已无信息增益"
                    break
                if round_index == self.max_rounds:
                    end_reason = f"达到最大业务轮数 {self.max_rounds}"
        except Exception as exc:
            failed = True
            end_reason = f"{type(exc).__name__}: {exc}"
            state.failures.append({"round": state.round_index, "error": end_reason})
            if state.round_index:
                try:
                    self.service.update_round(
                        run_id, state.round_index, status="failed", error=end_reason,
                    )
                except Exception:
                    pass

        final_run = self.service.finalize(run_id, failed=failed, reason=end_reason)
        final_output = (
            f"find_agent_v2 {final_run['outcome_status']}，"
            f"有效 primary={final_run['valid_primary_count']}。"
        )
        if not failed:
            try:
                report = await self.node_runner.run_node(
                    node="report",
                    round_index=state.round_index,
                    system_prompt=REPORT_PROMPT,
                    user_content=json.dumps(
                        self.service.get_full_state(run_id),
                        ensure_ascii=False,
                        default=str,
                    ),
                    tools=REPORT_TOOLS,
                    max_iterations=4,
                    slots=build_node_slots(
                        user_input=state.user_input,
                        full_state=self.service.get_full_state(run_id),
                        round_index=state.round_index,
                        plan=state.plan,
                    ),
                )
                state.node_runs.append(report)
                final_output = report.content or final_output
            except Exception as exc:
                state.failures.append({"round": state.round_index, "node": "report", "error": str(exc)})

        outcome = str(final_run.get("outcome_status") or "failed")
        return FindAgentResult(
            run_id=run_id,
            status=outcome if outcome in {"goal_met", "partial", "no_match", "failed"} else "failed",  # type: ignore[arg-type]
            succeeded=not failed and final_run.get("status") == "finished",
            business_outcome=outcome,
            valid_primary_count=int(final_run.get("valid_primary_count") or 0),
            rounds=state.round_index,
            final_output=final_output,
            node_runs=tuple(state.node_runs),
            stop_reason=end_reason,
        )


def create_find_agent_v2(
    settings: Settings | None = None,
    *,
    model: str | None = None,
    planning_model: str | None = None,
    search_model: str | None = None,
    evidence_model: str | None = None,
    evaluation_model: str | None = None,
    report_model: str | None = None,
    max_rounds: int = 2,
) -> FindAgentV2:
    return FindAgentV2(
        settings=settings,
        models_by_role=normalize_models(
            model=model,
            planning=planning_model,
            search=search_model,
            evidence=evidence_model,
            evaluation=evaluation_model,
            report=report_model,
        ),
        max_rounds=max_rounds,
    )