Blame for src/tools/analytics/aggregation.py

"""analyze_aggregate — group-by aggregation (KM-608).

An analytical "family" tool: in ONE call it groups rows by one or more keys
and computes aggregates (sum, mean, count, min, max, median, nunique) per
group. This is the deterministic compute twin of the Planner's
`retrieve_data` step — the wrapper layer later maps a QueryIR onto this.

STATUS: compute layer only — the function takes an already-materialized
DataFrame. The wrapper layer (fetching data from the catalog via source_id,
the ToolOutput envelope, ToolSpec registration) is added once the Planner
seam (KM-418) is settled. Keeping compute separate from data-fetching makes
this function easy to unit-test in isolation and stable when wrapped.
"""

from __future__ import annotations

import pandas as pd
from pydantic import BaseModel, Field

from src.tools.analytics.agg_vocab import normalise_agg
from src.tools.analytics.descriptive import ColumnNotFoundError

# Aggregation functions the tool understands. Whitelisted on purpose so an
# unknown function fails loudly instead of silently doing the wrong thing.
SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median", "nunique")


class UnsupportedAggregationError(ValueError):
    """Requested aggregation is not in SUPPORTED_AGGS (maps to UNSUPPORTED_AGG)."""


class UnsupportedGroupByError(ValueError):
    """A group_by entry is not a plain column name (maps to UNSUPPORTED_GROUP_BY).

    group_by keys must be column-name strings. A derived grouping — e.g. a CASE /
    binning expression the planner emits as a dict — is unhashable and would crash
    `df.groupby` with a bare "unhashable type: 'dict'"; raising here turns it into a
    clear, actionable error instead. Bucketing a numeric column into ranges belongs
    to analyze_segmentation, not analyze_aggregate.
    """


def _clean(value: object) -> object:
    """Convert numpy/pandas scalars to plain Python so the output is JSON-clean.

    A `group_by` over a datetime column yields `pandas.Timestamp` group keys,
    which (like numpy scalars) are not JSON-serializable — normalise those too.
    """
    if isinstance(value, pd.Timestamp):
        return value.isoformat()
    if hasattr(value, "item"):
        return value.item()
    return value


# Prompt-style description read by the Planner to decide WHEN to pick this tool.
# Final destination is ToolSpec.description once the wrapper layer is built.

class AggregateInput(BaseModel):
    """Planner-facing inputs for `analyze_aggregate` (spec_builder contract)."""

    data: str = Field(..., description="Placeholder ${t<id>} of the upstream table (Pattern A).")
    aggregations: dict[str, list[str]] = Field(
        ...,
        description="Column -> list of aggregation functions, e.g. "
        '{"revenue": ["sum", "mean"], "order_id": ["count"]}. '
        "`avg`/`count_distinct` are accepted as synonyms of `mean`/`nunique`.",
    )
    group_by: list[str] | None = Field(
        None,
        description="Column names to group by; omit to aggregate the whole table into one row.",
    )
    sort_by: str | None = Field(
        None,
        description='Output column to rank the groups by, e.g. "is_bad_actor_sum" '
        "(<column>_<func>); the aggregated column's own name works when it has one "
        "function. Omit to keep group-key order.",
    )
    ascending: bool = Field(
        False, description="Sort order for sort_by; default largest first."
    )
    top_n: int | None = Field(
        None, description="Keep only the first N groups after sorting; omit for all."
    )

DESCRIPTION = """\
Summary: Group-by aggregation. Splits rows by one or more key columns and \
computes aggregates per group. Returns one row per group.

SUPPORTED FUNCTIONS — use ONLY these: sum, mean, count, min, max, median, \
nunique. Standard deviation / variance are NOT available here: for the spread of \
a whole column use analyze_descriptive (whole-column only — spread PER GROUP is \
not available in v1). NEVER pass std / var / stdev to this tool; it will fail.

USE WHEN a per-group statistic needs a function the retrieve_data query \
cannot compute — `median` per group — on rows a retrieve_data task already \
returned. A plain sum / average / count / min / max per group is computed \
INSIDE the retrieve_data query (group_by + agg), over every matching row; it \
does not come here.

RANKING — for "which X has the most / least", set sort_by to the aggregate's \
output column (`<column>_<func>`, with ascending=true for "least") and top_n. \
Unsorted, the groups come back in key order and the top ones are not first.

DON'T USE WHEN:
  - it summarizes a column with no grouping -> analyze_descriptive
  - it compares two specific groups (A vs B) -> analyze_ab_test
  - it splits a single total into shares -> analyze_pareto
  - it asks for direction, growth or turning points over time -> analyze_trend \
(a plain total per month is a grouped retrieve_data)
  - a plain total / average / count per group -> group_by + agg in retrieve_data

GRACEFUL DEGRADE — if a tool named above (e.g. analyze_ab_test) is NOT present \
in the "Available tools" list, do NOT fake it with an unsupported aggregation \
(like std) and do NOT abandon the task. Degrade to an answer that stays on the \
SAME question: group the measure by the group column with `mean` (optionally \
min/max/median) so each group's level and their difference are still shown. Do \
NOT switch to an unrelated tool just to produce output. The Assembler will note \
that a significance test / per-group spread was not computed — a narrower, true \
answer beats a failed one.

Example questions:
  - "median delivery time per warehouse"
  - "median and mean basket size for each store"
"""


def analyze_aggregate(
    df: pd.DataFrame,
    aggregations: dict[str, list[str]],
    group_by: list[str] | None = None,
    sort_by: str | None = None,
    ascending: bool = False,
    top_n: int | None = None,
) -> list[dict[str, object]]:
    """Group-by aggregation over one or many keys.

    Args:
        df: already-materialized data (in the real system the wrapper fetches
            this from a source_id).
        aggregations: which columns to aggregate and how, e.g.
            {"revenue": ["sum", "mean"], "order_id": ["count"]}.
        group_by: grouping keys. If None/empty, the whole table is aggregated
            into a single row.
        sort_by: output column to order the groups by ("<column>_<func>" or a
            group key; a bare aggregated column with one function is accepted).
            None keeps group-key order.
        ascending: sort order for sort_by (default largest first).
        top_n: keep only the first N groups (after sorting); None keeps all.

    Returns:
        list[dict]: one row per group. Each row holds the group keys plus a
        column per aggregate, named "<column>_<func>" (e.g. "revenue_sum").

    Raises:
        ColumnNotFoundError: if any group_by, aggregated or sort_by column is absent.
        UnsupportedAggregationError: if a requested function is not supported.
    """
    group_by = group_by or []

    # group_by keys must be plain column names. A non-string entry (e.g. a CASE /
    # binning dict) is unhashable and would otherwise crash the membership check or
    # df.groupby with an opaque "unhashable type: 'dict'".
    non_str = [g for g in group_by if not isinstance(g, str)]
    if non_str:
        raise UnsupportedGroupByError(
            f"group_by entries must be column names (strings); got non-string "
            f"entr(ies) {non_str}. Derived groupings (e.g. CASE/binning expressions) "
            "are not supported by analyze_aggregate — group by an existing column."
        )

    # Validate columns first (fail-fast on caller mistakes).
    referenced = list(group_by) + list(aggregations.keys())
    missing = [c for c in referenced if c not in df.columns]
    if missing:
        raise ColumnNotFoundError(f"columns not found: {missing}")

    # Validate aggregation functions. Synonyms first: `avg`/`count_distinct` are the
    # QueryIR's names for `mean`/`nunique` and planner.md teaches them, so the same
    # plan can legitimately carry either spelling. (A23 Bug 2 -- see agg_vocab.)
    aggregations = {
        col: [normalise_agg(f) for f in funcs] for col, funcs in aggregations.items()
    }
    for col, funcs in aggregations.items():
        bad = [f for f in funcs if f not in SUPPORTED_AGGS]
        if bad:
            raise UnsupportedAggregationError(
                f"unsupported aggregation(s) {bad} for column '{col}'; "
                f"supported: {list(SUPPORTED_AGGS)}"
            )

    # Build named aggregations: {"revenue_sum": ("revenue", "sum"), ...}
    named = {
        f"{col}_{func}": (col, func)
        for col, funcs in aggregations.items()
        for func in funcs
    }

    # No grouping → aggregate the entire table into a single row.
    if not group_by:
        row = {name: _clean(df[col].agg(func)) for name, (col, func) in named.items()}
        return [row]

    # dropna=False keeps groups whose key is null so completeness is honest.
    grouped = df.groupby(group_by, dropna=False).agg(**named).reset_index()
    if sort_by is not None:
        # Test demo Q5: 261 units in key order; the answer read the first rows and
        # called once-only bad actors "recurring". A ranking needs the ranked column.
        key = _resolve_sort_key(sort_by, named, group_by)
        grouped = grouped.sort_values(key, ascending=ascending, kind="stable", na_position="last")
    if top_n is not None and top_n > 0:
        grouped = grouped.head(top_n)
    return [{k: _clean(v) for k, v in record.items()} for record in grouped.to_dict("records")]


def _resolve_sort_key(
    sort_by: str, named: dict[str, tuple[str, str]], group_by: list[str]
) -> str:
    """The output column `sort_by` names: an aggregate ("x_sum"), a group key, or a
    bare aggregated column when it has exactly one function ("x" -> "x_sum")."""
    if sort_by in named or sort_by in group_by:
        return sort_by
    by_source = [out for out, (col, _func) in named.items() if col == sort_by]
    if len(by_source) == 1:
        return by_source[0]
    raise ColumnNotFoundError(
        f"sort_by {sort_by!r} is not an output column; use one of "
        f"{list(named) + list(group_by)}"
    )