File size: 5,825 Bytes
81e5fe7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """analyze_segment — bucket rows into segments (KM-608).
An analytical "family" tool: in ONE call it bins a numeric column into
segments and reports how rows distribute across them (count, and optionally an
aggregate of another column per segment). Two binning modes: explicit cut
"edges" (e.g. age 0-18-35-60) or equal-frequency "quantile" buckets (quartiles,
deciles). Answers questions like "split customers into age brackets" or "bucket
orders into value tiers".
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 math
import pandas as pd
from src.tools.analytics.descriptive import ColumnNotFoundError
# Binning strategies.
SUPPORTED_METHODS = ("edges", "quantile")
# How to aggregate the value column within each segment.
SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median")
class InvalidMethodError(ValueError):
"""The requested binning method is not supported (maps to INVALID_METHOD)."""
class NonNumericColumnError(ValueError):
"""The column to segment on is not numeric (maps to NON_NUMERIC_COLUMN)."""
class UnsupportedAggregationError(ValueError):
"""The requested aggregation is not supported (maps to UNSUPPORTED_AGG)."""
def _clean(value: object) -> object:
"""Convert numpy scalars to plain Python; NaN -> None for JSON-clean output."""
if value is None:
return None
if hasattr(value, "item"):
value = value.item()
if isinstance(value, float) and math.isnan(value):
return None
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.
DESCRIPTION = """\
Summary: Bins a NUMERIC column into segments and counts how rows distribute \
across them (optionally aggregating another column per segment). Two modes: \
explicit cut edges (e.g. age 0-18-35-60) or equal-frequency quantile buckets \
(quartiles, deciles).
USE WHEN the question asks to bucket/bracket a continuous number into ranges. \
Trigger words: "segment" (segmen), "bucket/bracket" (kelompokkan ke rentang), \
"age groups/tiers" (kelompok umur/tingkatan), "quartiles/deciles", "bins".
DON'T USE WHEN:
- the category already exists (no binning needed) -> analyze_contribution
- it aggregates by an existing key -> analyze_aggregate
- it compares two named groups -> analyze_comparison
Example questions:
- "split customers into age brackets 0-18, 18-35, 35-60"
- "bucket orders into value tiers"
- "divide users into spending quartiles"
- "how many customers fall in each income band?"
"""
def analyze_segment(
df: pd.DataFrame,
column: str,
bins: list[float] | int,
method: str = "edges",
labels: list[str] | None = None,
value_column: str | None = None,
agg: str = "sum",
) -> dict[str, object]:
"""Segment rows by binning a numeric column.
Args:
df: already-materialized data (in the real system the wrapper fetches
this from a source_id).
column: numeric column to bin on.
bins: for method "edges", the list of cut boundaries (e.g.
[0, 18, 35, 60]); for method "quantile", the number of equal-
frequency buckets (e.g. 4 for quartiles).
method: "edges" (explicit boundaries) or "quantile" (equal frequency).
labels: optional segment names; for "edges" there must be
len(bins) - 1 of them.
value_column: if given, also aggregate this column per segment.
agg: how to aggregate value_column — one of SUPPORTED_AGGS.
Returns:
dict with:
column, method — echo of the chosen settings
agg — present only when value_column is given
segments — [{"segment", "count", ("value")}], in bin order
Raises:
ColumnNotFoundError: if column or value_column is absent.
NonNumericColumnError: if column is not numeric.
InvalidMethodError: if method is unknown.
UnsupportedAggregationError: if agg is not supported.
"""
referenced = [column] + ([value_column] if value_column else [])
missing = [c for c in referenced if c not in df.columns]
if missing:
raise ColumnNotFoundError(f"columns not found: {missing}")
if not pd.api.types.is_numeric_dtype(df[column]):
raise NonNumericColumnError(f"column '{column}' is not numeric")
if method not in SUPPORTED_METHODS:
raise InvalidMethodError(
f"unknown method '{method}'; supported: {list(SUPPORTED_METHODS)}"
)
if value_column is not None and agg not in SUPPORTED_AGGS:
raise UnsupportedAggregationError(
f"unsupported aggregation '{agg}'; supported: {list(SUPPORTED_AGGS)}"
)
if method == "edges":
cats = pd.cut(df[column], bins=bins, labels=labels, include_lowest=True)
else: # quantile
cats = pd.qcut(df[column], q=bins, labels=labels, duplicates="drop")
grouped = df.groupby(cats, observed=False)
counts = grouped.size()
segments = []
for seg in counts.index:
row = {"segment": str(seg), "count": int(counts[seg])}
if value_column is not None:
row["value"] = _clean(grouped[value_column].agg(agg).get(seg))
segments.append(row)
out: dict[str, object] = {"column": column, "method": method, "segments": segments}
if value_column is not None:
out["agg"] = agg
return out
|