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