| """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 |
|
|
| |
| |
| 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 |
|
|
|
|
| |
| |
|
|
| 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. "pa_median" ' |
| "(<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 — only a MEDIAN ranking comes here: "which X has the highest / lowest \ |
| median" sets sort_by to `<column>_median` (ascending=true for "lowest") and top_n; \ |
| unsorted, the groups come back in key order and the top ones are not first. A \ |
| ranking by sum / average / count / min / max is a grouped retrieve_data with \ |
| order_by + limit — never pull raw rows to rank them here. |
| |
| 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 [] |
|
|
| |
| |
| |
| 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." |
| ) |
|
|
| |
| 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}") |
|
|
| |
| |
| |
| 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)}" |
| ) |
|
|
| |
| named = { |
| f"{col}_{func}": (col, func) |
| for col, funcs in aggregations.items() |
| for func in funcs |
| } |
|
|
| |
| if not group_by: |
| row = {name: _clean(df[col].agg(func)) for name, (col, func) in named.items()} |
| return [row] |
|
|
| |
| grouped = df.groupby(group_by, dropna=False).agg(**named).reset_index() |
| if sort_by is not None: |
| |
| |
| 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)}" |
| ) |
|
|
