src/tools/invoker.py

"""AnalyticsToolInvoker — the runtime seam implementation (KM-465).

Implements the `ToolInvoker` Protocol the slow-path TaskRunner calls
(src/agents/slow_path/invoker.py). One method, `invoke(tool_name, args)`, does the
whole job for the `analyze_*` family:

1. Look the tool up in a name -> (compute fn, output_kind) dispatch map; an unknown
   name returns an error envelope (never an exception).
2. Materialize the Pattern A `data` argument — which the TaskRunner has already
   resolved to the upstream task's `ToolOutput` (kind="table") — into a DataFrame.
3. Call the pure compute function with the remaining args as keyword arguments
   (their names match the compute signatures one-to-one).
4. Wrap the result in a `ToolOutput` with the tool's declared `kind`.

Frozen guarantee (§8.4): **never throws.** Any failure — unknown tool, bad data,
or an exception from compute (e.g. GroupNotFoundError) — comes back as
`ToolOutput(kind="error", error=...)`, so the TaskRunner's degrade-and-continue
keeps working.
"""

from __future__ import annotations

import decimal
from collections.abc import Callable
from typing import Any

import pandas as pd

from src.tools.analytics import (
    aggregation,
    comparison,
    decomposition,
    descriptive,
    quality,
    relationship,
    segmentation,
    temporal,
)
from src.tools.contracts import ToolOutput
from src.tools.data_access import DATA_ACCESS_TOOLS, DataAccessToolInvoker

# tool name -> (compute callable, ToolOutput.kind it produces). Kept in lockstep
# with src/tools/registry.py output_kind values.
_DISPATCH: dict[str, tuple[Callable[..., Any], str]] = {
    "analyze_descriptive": (descriptive.analyze_descriptive, "stats"),
    "analyze_aggregate": (aggregation.analyze_aggregate, "table"),
    "analyze_comparison": (comparison.analyze_comparison, "stats"),
    "analyze_contribution": (decomposition.analyze_contribution, "table"),
    "analyze_profile": (quality.analyze_profile, "stats"),
    "analyze_correlation": (relationship.analyze_correlation, "stats"),
    "analyze_segment": (segmentation.analyze_segment, "table"),
    "analyze_trend": (temporal.analyze_trend, "series"),
}


class AnalyticsToolInvoker:
    """Never-throwing invoker for the `analyze_*` tools (implements ToolInvoker)."""

    async def invoke(self, tool_name: str, args: dict[str, Any]) -> ToolOutput:
        entry = _DISPATCH.get(tool_name)
        if entry is None:
            return ToolOutput(
                tool=tool_name, kind="error", error=f"unknown tool {tool_name!r}"
            )
        fn, kind = entry

        df, err = _materialize(args.get("data"))
        if err is not None:
            return ToolOutput(tool=tool_name, kind="error", error=err)

        kwargs = {k: v for k, v in args.items() if k != "data"}
        try:
            result = fn(df, **kwargs)
        except Exception as exc:  # noqa: BLE001 — never-throw seam (§8.4)
            return ToolOutput(
                tool=tool_name,
                kind="error",
                error=f"{type(exc).__name__}: {exc}",
            )

        return ToolOutput(tool=tool_name, kind=kind, value=result)


class CompositeToolInvoker:
    """One `invoke()` for the whole tool surface (KM-465 #4).

    The TaskRunner only ever calls one `ToolInvoker`. This composes the two
    families behind a single dispatch: the stateless `AnalyticsToolInvoker`
    (`analyze_*`) and the per-request stateful `DataAccessToolInvoker`
    (catalog/query/retrieval, which need the authenticated `user_id`). Routing
    is by tool name; an unknown name falls through to the analytics invoker,
    which returns the standard unknown-tool error envelope.

    Constructed per-request — the Coordinator injects the request's `user_id`
    and `CatalogReader` into the data-access invoker (INV-7: the agent layer
    stays tool-agnostic).

    Frozen guarantee (§8.4): **never throws** — both sub-invokers return
    `ToolOutput(kind="error", ...)` on any failure.
    """

    def __init__(
        self,
        data_access: DataAccessToolInvoker,
        analytics: AnalyticsToolInvoker | None = None,
    ) -> None:
        self._data_access = data_access
        self._analytics = analytics or AnalyticsToolInvoker()

    async def invoke(self, tool_name: str, args: dict[str, Any]) -> ToolOutput:
        if tool_name in DATA_ACCESS_TOOLS:
            return await self._data_access.invoke(tool_name, args)
        return await self._analytics.invoke(tool_name, args)


def _materialize(data: Any) -> tuple[pd.DataFrame, None] | tuple[None, str]:
    """Turn the resolved `data` argument into a DataFrame.

    Accepts the upstream `ToolOutput` (kind="table"), a raw DataFrame, or a
    {"columns", "rows"} dict (a serialized table). Returns (df, None) on success
    or (None, error_message) on failure — the caller wraps the message.

    Numeric columns are coerced to float (see `_coerce_decimals`): DB NUMERIC
    columns arrive as Python `Decimal`, which mixes badly with the float math in
    the `analyze_*` compute functions (e.g. `float + Decimal` -> TypeError).
    Normalizing here fixes the whole tool family in one place.
    """
    if data is None:
        return None, "missing 'data' argument (no upstream table to analyze)"

    if isinstance(data, pd.DataFrame):
        return _coerce_decimals(data), None

    if isinstance(data, ToolOutput):
        if data.kind == "error":
            return None, f"upstream data unavailable: {data.error}"
        if data.kind != "table" or data.columns is None:
            return None, f"cannot materialize 'data' of kind {data.kind!r}"
        return _coerce_decimals(pd.DataFrame(data.rows or [], columns=data.columns)), None

    if isinstance(data, dict) and "columns" in data:
        df = pd.DataFrame(data.get("rows") or [], columns=data["columns"])
        return _coerce_decimals(df), None

    return None, f"unsupported 'data' type: {type(data).__name__}"


def _coerce_decimals(df: pd.DataFrame) -> pd.DataFrame:
    """Convert `decimal.Decimal` object-columns to float64 in place.

    DB drivers (asyncpg) return NUMERIC/DECIMAL values as Python `Decimal`, which
    land in object-dtype columns. The `analyze_*` compute functions do float math
    on these (share-of-total, cumulative sums), and `float + Decimal` raises
    TypeError. We only touch columns that actually contain a `Decimal`, so real
    string/categorical columns are left untouched. `None`/missing values become
    NaN, which the compute functions already handle.
    """
    for col in df.columns:
        if df[col].dtype == object and df[col].map(lambda v: isinstance(v, decimal.Decimal)).any():
            df[col] = df[col].astype(float)
    return df