src/tools/contracts.py

"""Canonical tool contracts (KM-465 — owned by the tool team).

These are the single source of truth for the tool <-> agent interface:

  - `ToolSpec` / `ToolRegistry` — the registry contract (§9.2). The concrete v1
    analytics registry instance is built on top of these.
  - `ToolOutput` — the tool -> agent output envelope (§8.1). Tools return this at
    TaskRunner time; the agent layer plans and degrades against it.

Ownership note (2026-06-08): the tool team owns these definitions outright; the
agent team's earlier copies in `src/agents/planner/contracts.py` are now thin
re-exports of this module so there is exactly ONE definition across the codebase.
The shapes here are kept byte-for-byte identical to those original stubs so the
already-landed planner / TaskRunner / Assembler keep working unchanged.

Data-flow decision (KM-465, agreed with the agent team): the analytics tools use
**Pattern A** — `analyze_*` tools do NOT self-fetch by `source_id`; each takes a
`data` argument that is a `"${t<id>}"` placeholder resolved to a DataFrame at
execution time. `analyze_comparison.output_kind` is `"stats"` (a labelled-metric
dict), aligning with the planner registry.

See AGENT_ARCHITECTURE_CONTEXT_new.md §8.1 / §9.2.
"""

from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, Field

# --------------------------------------------------------------------------- #
# Tool registry (§9.2)
# --------------------------------------------------------------------------- #


class ToolSpec(BaseModel):
    name: str
    category: str  # analytics.query | .aggregation | .timeseries | ...
    # JSON-schema-ish dict: {"required": [...], "properties": {arg: {"type": ...}}}.
    # VALIDATION CONTRACT — presence only: TaskRunner._validate_args enforces just
    # `required` (each must resolve to a non-None arg). The `properties` types are
    # DOCUMENTATION for the planner prompt, NOT checked at runtime — a wrong-typed
    # arg passes validation and only surfaces (if at all) inside the compute fn.
    # Do not assume type-safety here.
    input_schema: dict[str, Any]
    output_kind: str  # the ToolOutput.kind it returns
    description: str  # prompt-style: what it does, edge cases, what NOT to use it for
    phase: Literal["P0", "P1", "P2"] = "P0"


class ToolRegistry(BaseModel):
    tools: list[ToolSpec] = Field(default_factory=list)

    def names(self) -> set[str]:
        return {t.name for t in self.tools}

    def get(self, name: str) -> ToolSpec | None:
        for t in self.tools:
            if t.name == name:
                return t
        return None


# --------------------------------------------------------------------------- #
# Tool output envelope (§8.1)
# --------------------------------------------------------------------------- #


class ToolOutput(BaseModel):
    tool: str
    kind: Literal["scalar", "table", "stats", "series", "documents", "error"]
    value: Any | None = None
    columns: list[str] | None = None
    rows: list[list[Any]] | None = None
    meta: dict[str, Any] = Field(default_factory=dict)
    error: str | None = None