src/agents/planner/schemas.py

"""Planner output schemas — the `TaskList` contract.

The planner emits exactly one `TaskList`: a DAG of typed tasks, each an ordered
chain of fully-specified tool calls. This is the static plan the TaskRunner
executes verbatim. There is no replanning (INV-6), so there are no
`ReplanRequest` / `ReplanResponse` schemas.

See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3.
"""

from __future__ import annotations

import re
from typing import Any, Literal

from pydantic import BaseModel, Field

# The "${t<id>}" placeholder convention (Pattern A): a ToolCall arg whose value
# matches this pattern refers to an upstream task's output, resolved by the
# TaskRunner at execution time. Single definition — the planner validator and
# the TaskRunner both import it (R11).
PLACEHOLDER_RE = re.compile(r"\$\{(t[^}]+)\}")

CrispStage = Literal[
    "data_understanding",
    "data_preparation",
    "modeling",  # no tools in v1; the planner does not emit modeling tasks
    "evaluation",
]


class ToolCall(BaseModel):
    """One call to a registry tool with concrete, fully-specified arguments.

    `tool` must exist in the ToolRegistry. `args` is validated against the
    tool's input_schema; it may contain "${t<id>}" placeholders that the
    TaskRunner resolves from an upstream task's output at execution time.
    """

    tool: str
    args: dict[str, Any] = Field(default_factory=dict)


class Task(BaseModel):
    id: str  # "t1", "t2", ...
    stage: CrispStage
    objective: str  # plain-language intent for this step
    tool_calls: list[ToolCall] = Field(..., min_length=1)  # ordered chain
    expected_output: str  # named result this task produces
    success_criteria: str  # REPORTING signal, not a control trigger
    depends_on: list[str] = Field(default_factory=list)  # task ids
    estimated_cost: Literal["low", "medium", "high"] = "low"


class TaskList(BaseModel):
    plan_id: str
    goal_restated: str
    assumptions: list[str] = Field(default_factory=list)
    open_questions: list[str] = Field(default_factory=list)
    tasks: list[Task] = Field(default_factory=list)
    # Infeasible sentinel (planner.md "When the catalog cannot answer"): set with
    # an EMPTY `tasks` list when no catalog column plausibly holds the requested
    # measure/entity. Explains what is missing and names the nearest available
    # data. The coordinator renders it as an honest data-gap answer instead of
    # running the pipeline — the alternative was the planner force-mapping
    # unrelated columns (observed: `pa` aliased as "revenue").
    infeasible_reason: str | None = Field(
        None,
        description=(
            "Set ONLY when the question cannot be answered from the catalog: no "
            "column plausibly holds the requested measure or entity. State what "
            "is missing and the nearest data that IS available. Leave tasks "
            "empty when set."
        ),
    )