SupplyAgent/find_agent_v2/state.py

"""Thin workflow state.

Search pages, candidates, evidence and final buckets stay in the existing database.
This state only contains orchestration pointers and audit summaries.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Literal, TypedDict

from supply_agent.types import AgentResult

Phase = Literal["planning", "searching", "evidence", "evaluating", "done"]
SupervisorAction = Literal["search", "evidence", "evaluator", "finish"]
EndKind = Literal["goal_met", "partial", "no_match", "failed", "stopped"]


@dataclass(frozen=True)
class DiscoverySnapshot:
    """Small deterministic projection of the database state."""

    status: str
    search_count: int
    candidate_count: int
    pending_count: int
    primary_count: int
    valid_primary_count: int
    rejected_count: int
    outcome_status: str = ""


@dataclass(frozen=True)
class NodeRun:
    """One node's observable result."""

    node: str
    round_index: int
    content: str
    iterations: int
    tool_calls_made: int

    @classmethod
    def from_agent_result(
        cls,
        node: str,
        round_index: int,
        result: AgentResult,
    ) -> "NodeRun":
        return cls(
            node=node,
            round_index=round_index,
            content=result.content or "",
            iterations=int(result.iterations or 0),
            tool_calls_made=int(result.tool_calls_made or 0),
        )


@dataclass
class FindAgentState:
    """Control state passed through the single-round graph."""

    run_id: str
    user_input: str
    round_index: int = 0
    phase: Phase = "planning"
    plan: str = ""
    stop: bool = False
    stop_reason: str = ""
    previous_snapshot: DiscoverySnapshot | None = None
    snapshot: DiscoverySnapshot | None = None
    node_runs: list[NodeRun] = field(default_factory=list)
    failures: list[dict[str, Any]] = field(default_factory=list)


class FindAgentGraphState(TypedDict, total=False):
    """Serializable state used by the real one-round LangGraph."""

    run_id: str
    user_input: str
    round_index: int
    plan: str
    phase: Phase
    node_runs: list[NodeRun]
    snapshot: DiscoverySnapshot | None
    supervisor_step: int
    approved_action: SupervisorAction
    action_count: int
    search_actions: int
    worker_count: int
    evidence_scope: str
    decision_history: list[dict[str, Any]]
    evaluator_stagnation: int


@dataclass(frozen=True)
class FindAgentResult:
    """Workflow result with technical and business status kept separate."""

    run_id: str
    status: EndKind
    succeeded: bool
    business_outcome: str
    valid_primary_count: int
    rounds: int
    final_output: str
    node_runs: tuple[NodeRun, ...]
    stop_reason: str = ""

    @property
    def iterations(self) -> int:
        return sum(item.iterations for item in self.node_runs)

    @property
    def tool_calls_made(self) -> int:
        return sum(item.tool_calls_made for item in self.node_runs)