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 pydantic import BaseModel, Field, model_validator

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


class DemandBrief(BaseModel):
    """Planner-owned compact interpretation of the immutable raw demand."""

    core_intent: str = Field(min_length=1)
    target_audience: list[str] = Field(default_factory=list)
    forwarding_audience: list[str] = Field(default_factory=list)
    key_scenarios: list[str] = Field(default_factory=list)
    reference_signals: list[str] = Field(default_factory=list)
    relevance_criteria: list[str] = Field(min_length=1)
    exclusion_criteria: list[str] = Field(default_factory=list)
    temporal_requirements: list[str] = Field(default_factory=list)


class SearchTask(BaseModel):
    """One independently executable search task produced by planning."""

    task_id: str = Field(min_length=1)
    keyword: str = Field(min_length=1)
    query_reason: str = Field(min_length=1)
    source_type: Literal["demand", "seed", "point", "mixed"] = "mixed"
    provider: Literal["internal_keyword", "tikhub"] | None = Field(
        default=None,
        description="不填则先内部搜索、无结果再 TikHub；internal_keyword 只用内部搜索；tikhub 只用 TikHub。",
    )
    max_pages: int = Field(default=1, ge=1, le=2)
    coverage_targets: list[str] = Field(default_factory=list)


class EvaluationBrief(BaseModel):
    """Demand semantics required by evaluators, without the original request."""

    relevance_criteria: list[str] = Field(min_length=1)
    audience_criteria: list[str] = Field(default_factory=list)
    temporal_requirements: list[str] = Field(default_factory=list)
    exclusion_criteria: list[str] = Field(default_factory=list)


class ExecutionPlan(BaseModel):
    """Validated contract passed from the planner to later stages."""

    schema_version: Literal["1.0"] = "1.0"
    demand_brief: DemandBrief
    search_tasks: list[SearchTask] = Field(min_length=1, max_length=24)
    evaluation_brief: EvaluationBrief

    @model_validator(mode="after")
    def validate_unique_task_ids(self):
        task_ids = [item.task_id for item in self.search_tasks]
        if len(task_ids) != len(set(task_ids)):
            raise ValueError("search_tasks.task_id 不能重复")
        return self


class SupervisorDecision(BaseModel):
    """One routing decision; subsequent decisions may add search tasks."""

    next_action: SupervisorAction
    reason: str = ""
    worker_count: int = Field(default=4, ge=1, le=8)
    evidence_scope: Literal["detail", "portrait", "both"] = "both"
    additional_search_tasks: list[SearchTask] = Field(default_factory=list, max_length=6)


class PlanningDecision(SupervisorDecision):
    execution_plan: ExecutionPlan


class PlannerAssignment(BaseModel):
    run_id: str
    round_index: int
    raw_demand: str


class SupervisorAssignment(BaseModel):
    run_id: str
    round_index: int
    execution_plan: dict[str, Any]
    execution_state: dict[str, Any]


class SearchAssignment(BaseModel):
    run_id: str
    round_index: int
    tasks: list[SearchTask]


class EvidenceAssignment(BaseModel):
    run_id: str
    round_index: int
    candidate_ids: list[int]
    evidence_type: Literal["detail", "portrait"]
    candidates: list[dict[str, Any]]


class EvaluationAssignment(BaseModel):
    run_id: str
    round_index: int
    candidate_ids: list[int]
    evaluation_brief: EvaluationBrief
    quality_gate_rules: dict[str, Any]
    current_datetime: str = ""
    timezone: str = ""
    video_understanding_candidate_ids: list[int] = Field(default_factory=list)
    candidates: list[dict[str, Any]]


class ReportAssignment(BaseModel):
    run_id: str
    final_state: dict[str, Any]


@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"
    execution_plan: ExecutionPlan | None = None
    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
    execution_plan: ExecutionPlan | None
    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)