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

from collections.abc import Callable
from typing import Any

import pandas as pd

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

logger = get_logger("analytics_invoker")

# 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"),
    "analyze_merge": (merge.analyze_merge, "table"),
    "render_chart": (visualization.render_chart, "chart"),
}


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:
            logger.warning("tool returned error", tool=tool_name, error="unknown tool")
            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:
            logger.warning("tool returned error", tool=tool_name, error=err)
            return ToolOutput(tool=tool_name, kind="error", error=err)

        kwargs = {k: v for k, v in args.items() if k != "data"}

        # Second data input (Pattern A, extended): analyze_merge takes a
        # `data_right` placeholder the TaskRunner has resolved to another upstream
        # ToolOutput. Materialize it the same way as `data` and hand the compute fn
        # a DataFrame, not the raw envelope.
        if "data_right" in kwargs:
            df_right, err = _materialize(kwargs["data_right"])
            if err is not None:
                err = f"data_right: {err}"
                logger.warning("tool returned error", tool=tool_name, error=err)
                return ToolOutput(tool=tool_name, kind="error", error=err)
            kwargs["data_right"] = df_right

        try:
            result = fn(df, **kwargs)
        except Exception as exc:  # noqa: BLE001 — never-throw seam (§8.4)
            error = f"{type(exc).__name__}: {exc}"
            # Never-throw is intentional (§8.4), but a swallowed failure was
            # invisible: log it so a failed analysis step is diagnosable instead
            # of only surfacing as a vague "could not compute" in the answer.
            logger.warning("tool returned error", tool=tool_name, error=error)
            return ToolOutput(tool=tool_name, kind="error", error=error)

        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 normalized (see `_normalize_numeric`): DB NUMERIC values
    arrive as Python `Decimal`, and tabular sources sometimes store numbers as
    text — both break the float math in the `analyze_*` compute functions (or make
    a numeric column invisible to `is_numeric_dtype`). 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 _normalize_numeric(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 _normalize_numeric(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 _normalize_numeric(df), None

    # A {table_id/source_id} dict is a raw catalog reference the planner inlined
    # instead of chaining a retrieve_data output (Pattern A). analyze_* tools never
    # self-fetch, so give an actionable message rather than the opaque type name.
    # The planner validator should reject this upstream (Check 9); this is the
    # defensive net if a bad plan still reaches execution.
    if isinstance(data, dict) and ("table_id" in data or "source_id" in data):
        return None, (
            "'data' is a table reference (table_id/source_id), not a retrieved "
            "table — analyze_* consumes the output of retrieve_data, it does not "
            "fetch data by id (Pattern A)"
        )

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


def _normalize_numeric(df: pd.DataFrame) -> pd.DataFrame:
    """Coerce object-columns that are really numeric into numeric dtype in place.

    Two sources of "numbers hiding in object columns" break the analyze_* tools:
      - DB drivers (asyncpg) return NUMERIC/DECIMAL as Python `Decimal`, which
        raises `TypeError` on `float + Decimal` in share-of-total / cumulative math.
      - Tabular files (CSV/XLSX, or a stale Parquet) sometimes store numbers as
        text, so a numeric column is invisible to `pd.api.types.is_numeric_dtype`
        and tools like `analyze_correlation` see "0 numeric columns".

    Both are fixed by converting only the columns that are *entirely* numeric to a
    numeric dtype. A column with any non-numeric value (e.g. a category like
    "Online"/"Offline") fails the all-parseable check and is left untouched, so
    genuine categoricals are never mangled. Empty/None cells become NaN, which the
    compute functions already handle.

    Caveat: all-digit identifier columns stored as text (e.g. a zero-padded code
    "007") are treated as numeric — acceptable for an analytics data path.
    """
    for col in df.columns:
        if df[col].dtype != object:
            continue
        converted = pd.to_numeric(df[col], errors="coerce")
        # Convert only when every originally-present value parsed as a number, so
        # a single non-numeric value keeps the column as-is.
        if converted.notna().sum() == df[col].notna().sum():
            df[col] = converted
    return df