SupplyAgent/find_agent_v2/agent.py

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

from __future__ import annotations

import asyncio
from dataclasses import dataclass

from find_agent_v2.context import assignment_slots, render_assignment
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 (
    RUN_TIMEOUT_REASON,
    RUN_TIMEOUT_SECONDS,
    FindAgentV2Service,
    get_find_agent_v2_service,
)
from find_agent_v2.state import (
    DiscoverySnapshot,
    FindAgentResult,
    FindAgentState,
    ReportAssignment,
)
from find_agent_v2.tools import REPORT_TOOLS
from supply_agent.config import Settings


@dataclass(frozen=True)
class ExplorationDecision:
    continue_exploring: bool
    reason: str
    new_candidate_count: int
    new_evaluated_count: int
    round_pass_rate: float
    cumulative_pass_rate: float


def decide_continued_exploration(
    previous: DiscoverySnapshot, current: DiscoverySnapshot,
) -> ExplorationDecision:
    """Decide whether to continue from discovery volume and evaluation yield."""
    new_candidates = max(0, current.candidate_count - previous.candidate_count)
    previous_evaluated = previous.primary_count + previous.rejected_count
    current_evaluated = current.primary_count + current.rejected_count
    new_evaluated = max(0, current_evaluated - previous_evaluated)
    new_passed = max(0, current.valid_primary_count - previous.valid_primary_count)
    round_rate = new_passed / new_evaluated if new_evaluated else 0.0
    cumulative_rate = (
        current.valid_primary_count / current_evaluated if current_evaluated else 0.0
    )
    if new_candidates == 0:
        keep_going = False
        reason = "本轮没有新增候选，搜索前沿已无信息增益"
    elif current_evaluated < 8:
        keep_going = True
        reason = "已评估样本不足 8 条，继续探索以形成可靠通过率"
    elif new_evaluated and round_rate == 0:
        keep_going = False
        reason = "本轮有足够评估样本但通过率为 0，继续搜索的预期收益低"
    elif new_candidates >= 2 and (round_rate >= 0.10 or cumulative_rate >= 0.10):
        keep_going = True
        reason = "本轮仍有候选增量且评估保持正向通过率，继续探索"
    else:
        keep_going = False
        reason = "候选增量或评估通过率不足，停止探索"

    return ExplorationDecision(
        continue_exploring=keep_going,
        reason=reason,
        new_candidate_count=new_candidates,
        new_evaluated_count=new_evaluated,
        round_pass_rate=round(round_rate, 4),
        cumulative_pass_rate=round(cumulative_rate, 4),
    )


class FindAgentV2:
    """Python outer loop + guarded Supervisor graph + 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,
        max_actions_per_round: int = 16,
        max_search_actions_per_round: int = 3,
        max_runtime_seconds: float = RUN_TIMEOUT_SECONDS,
        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.max_actions_per_round = max(4, int(max_actions_per_round))
        self.max_search_actions_per_round = max(1, int(max_search_actions_per_round))
        self.max_runtime_seconds = max(0.01, float(max_runtime_seconds))

    async def arun(
        self, *, run_id: str, user_input: str, resume: bool = False,
    ) -> FindAgentResult:
        run = self.service.require_run(run_id)
        if resume and str(run.get("status") or "") != "running":
            run = self.service.prepare_resume(run_id)
        if str(run.get("status") or "") != "running":
            raise ValueError(f"run_id={run_id} 当前状态不可执行: {run.get('status')}")
        load_plan = getattr(self.service, "get_latest_execution_plan", None)
        restored_plan = load_plan(run_id) if callable(load_plan) else None
        state = FindAgentState(
            run_id=run_id, user_input=user_input, execution_plan=restored_plan,
        )
        graph = FindAgentRoundGraph(
            service=self.service, runner=self.node_runner, observer=self.observer,
            max_actions=self.max_actions_per_round,
            max_search_actions=self.max_search_actions_per_round,
        )
        reset_usage = getattr(self.node_runner, "reset_usage", None)
        if callable(reset_usage):
            reset_usage()
        default_model = self.models_by_role.get("supervisor", "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),
            )
            try:
                result = await asyncio.wait_for(
                    self._arun_inner(
                        state=state, graph=graph,
                        start_round=int(run.get("current_round") or 0) + 1,
                    ),
                    timeout=self.max_runtime_seconds,
                )
            except TimeoutError:
                self.service.fail_run(run_id, RUN_TIMEOUT_REASON)
                final_run = self.service.require_run(run_id)
                result = FindAgentResult(
                    run_id=run_id,
                    status="failed",
                    succeeded=False,
                    business_outcome="failed",
                    valid_primary_count=int(final_run.get("valid_primary_count") or 0),
                    rounds=state.round_index,
                    final_output=f"find_agent_v2 failed：{RUN_TIMEOUT_REASON}",
                    node_runs=tuple(state.node_runs),
                    stop_reason=RUN_TIMEOUT_REASON,
                )
            usage = getattr(self.node_runner, "usage", None)
            if isinstance(usage, dict):
                self.service.add_usage(run_id, usage)
            observation_run.finish(final_output=result.final_output)
            return result

    async def _arun_inner(
        self, *, state: FindAgentState, graph: FindAgentRoundGraph, start_round: int = 1,
    ) -> FindAgentResult:
        run_id = state.run_id
        end_reason = ""
        failed = False
        try:
            end_round = start_round + self.max_rounds
            for round_index in range(start_round, end_round):
                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)
                pending_count = self.service.count_pending_candidates(run_id)
                if pending_count:
                    failed = True
                    end_reason = f"第 {round_index} 轮结束仍有 {pending_count} 条待评估候选"
                    break
                exploration = decide_continued_exploration(
                    state.previous_snapshot, current,
                )
                if not exploration.continue_exploring:
                    end_reason = exploration.reason
                    break
                if round_index == end_round - 1:
                    end_reason = (
                        f"达到最大业务轮数 {self.max_rounds}；{exploration.reason}；"
                        f"本轮新增={exploration.new_candidate_count}，"
                        f"本轮通过率={exploration.round_pass_rate:.1%}，"
                        f"累计通过率={exploration.cumulative_pass_rate:.1%}"
                    )
        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_assignment = ReportAssignment(
                    run_id=run_id,
                    final_state=self.service.get_report_state(run_id),
                )
                report = await self.node_runner.run_node(
                    node="report",
                    round_index=state.round_index,
                    system_prompt=REPORT_PROMPT,
                    user_content=render_assignment(report_assignment),
                    tools=REPORT_TOOLS,
                    max_iterations=4,
                    slots=assignment_slots(
                        report_assignment,
                        source="FindAgentV2Service.get_report_state",
                    ),
                )
                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,
    max_actions_per_round: int = 16,
    max_search_actions_per_round: int = 3,
) -> 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,
        max_actions_per_round=max_actions_per_round,
        max_search_actions_per_round=max_search_actions_per_round,
    )