src/tools/registry.py

"""Analytics tool registry

The real registry of the `analyze_*` family, built on the canonical `ToolSpec`
(src/tools/contracts.py) and the prompt-style `DESCRIPTION` constants the Planner
reads to choose a tool (KM-625). This replaces the agent team's local stub in
`src/agents/planner/registry.py` for the analytics slice.

Conventions (decided with the agent team, KM-465):
- **Pattern A** — `analyze_*` tools do NOT self-fetch by `source_id`. Each takes a
  `data` argument that is a `"${t<id>}"` placeholder pointing at an upstream
  `retrieve_data` table output, resolved to a DataFrame at execution time.
  Column arguments reference the aliases that upstream query produced.
- `input_schema` is the lightweight JSON-schema-ish dict the planner validator
  consumes: `required` (arg names with no default) + `properties` (allowed args).
  `required` mirrors each compute function's no-default parameters; value-typing of
  placeholder args is deferred to execution time.
- `output_kind` is the `ToolOutput.kind` each tool returns: stats (labelled-metric
  dict) | table (rows×cols) | series (ordered periods).

The four data-access tools (retrieve_data / retrieve_knowledge / check_data /
check_knowledge) are registered separately once their wrappers land (KM-465 #4);
`default_registry()` composes both slices.
"""

from __future__ import annotations

from src.tools.analytics import (
    aggregation,
    comparison,
    decomposition,
    descriptive,
    quality,
    relationship,
    segmentation,
    temporal,
)
from src.tools.contracts import ToolRegistry, ToolSpec

# Active this round — the four analytics tools the Planner may select.
ACTIVE_ANALYTICS_TOOLS: list[ToolSpec] = [
    ToolSpec(
        name="analyze_descriptive",
        category="analytics.descriptive",
        input_schema={
            "required": ["data", "column_ids"],
            "properties": {
                "data": {"type": "string"},
                "column_ids": {"type": "array"},
                "metrics": {"type": "array"},
            },
        },
        output_kind="stats",
        description=descriptive.DESCRIPTION,
    ),
    ToolSpec(
        name="analyze_aggregate",
        category="analytics.aggregation",
        input_schema={
            "required": ["data", "aggregations"],
            "properties": {
                "data": {"type": "string"},
                "aggregations": {"type": "object"},
                "group_by": {"type": "array"},
            },
        },
        output_kind="table",
        description=aggregation.DESCRIPTION,
    ),
    ToolSpec(
        name="analyze_correlation",
        category="analytics.relationship",
        input_schema={
            "required": ["data"],
            "properties": {
                "data": {"type": "string"},
                "column_ids": {"type": "array"},
                "method": {"type": "string"},
            },
        },
        output_kind="stats",
        description=relationship.DESCRIPTION,
    ),
    ToolSpec(
        name="analyze_trend",
        category="analytics.timeseries",
        input_schema={
            "required": ["data", "date_column", "value_column"],
            "properties": {
                "data": {"type": "string"},
                "date_column": {"type": "string"},
                "value_column": {"type": "string"},
                "freq": {"type": "string"},
                "agg": {"type": "string"},
            },
        },
        output_kind="series",
        description=temporal.DESCRIPTION,
    ),
]

# Deferred this round — specs kept intact for easy re-activation, NOT exposed to
# the Planner. The compute fns still exist (src/tools/analytics/*) and the invoker
# still maps them (src/tools/invoker.py); only registry exposure is withheld.
# To re-activate, move a spec back into ACTIVE_ANALYTICS_TOOLS. NOTE: a deferred
# tool re-activated here must also be re-added to the Planner few-shots
# (src/agents/planner/examples.py) — keep the two in sync.
DEFERRED_ANALYTICS_TOOLS: list[ToolSpec] = [
    ToolSpec(
        name="analyze_comparison",
        category="analytics.comparison",
        input_schema={
            "required": ["data", "dimension", "value_column", "group_a", "group_b"],
            "properties": {
                "data": {"type": "string"},
                "dimension": {"type": "string"},
                "value_column": {"type": "string"},
                "group_a": {},
                "group_b": {},
                "agg": {"type": "string"},
            },
        },
        output_kind="stats",
        description=comparison.DESCRIPTION,
    ),
    ToolSpec(
        name="analyze_contribution",
        category="analytics.decomposition",
        input_schema={
            "required": ["data", "dimension", "value_column"],
            "properties": {
                "data": {"type": "string"},
                "dimension": {"type": "string"},
                "value_column": {"type": "string"},
                "agg": {"type": "string"},
                "top_n": {"type": "integer"},
            },
        },
        output_kind="table",
        description=decomposition.DESCRIPTION,
    ),
    ToolSpec(
        name="analyze_profile",
        category="analytics.quality",
        input_schema={
            "required": ["data"],
            "properties": {
                "data": {"type": "string"},
                "column_ids": {"type": "array"},
            },
        },
        output_kind="stats",
        description=quality.DESCRIPTION,
    ),
    ToolSpec(
        name="analyze_segment",
        category="analytics.segmentation",
        input_schema={
            "required": ["data", "column", "bins"],
            "properties": {
                "data": {"type": "string"},
                "column": {"type": "string"},
                "bins": {},
                "method": {"type": "string"},
                "labels": {"type": "array"},
                "value_column": {"type": "string"},
                "agg": {"type": "string"},
            },
        },
        output_kind="table",
        description=segmentation.DESCRIPTION,
    ),
]

# Full set (active + deferred) — kept for callers that need every spec, e.g. tests
# or the invoker's name checks. The Planner-visible registry uses ACTIVE only.
ANALYTICS_TOOLS: list[ToolSpec] = [*ACTIVE_ANALYTICS_TOOLS, *DEFERRED_ANALYTICS_TOOLS]


def analytics_registry() -> ToolRegistry:
    """The analytics (`analyze_*`) slice of the tool registry (fresh instance).

    Exposes only `ACTIVE_ANALYTICS_TOOLS`; deferred specs are withheld from the
    Planner (see `DEFERRED_ANALYTICS_TOOLS`).
    """
    return ToolRegistry(tools=list(ACTIVE_ANALYTICS_TOOLS))