src/agents/planner/examples.py

"""Few-shot examples for the planner prompt.

Illustrative (question -> TaskList) pairs that teach the OUTPUT SHAPE:
stages, dependency edges, ordered tool-call chains, inline QueryIR,
"${t<id>}" placeholders, and the assumed data-flow convention — `retrieve_data`
pulls rows, then a composite `analyze_*` tool consumes them via a `data` placeholder
referencing the upstream result's column aliases (Pattern A; the tool team may
instead pick self-fetch by `source_id`, in which case these examples are reshaped
to match — see registry.py). They reference a hypothetical sales catalog
(`src_sales` / `t_orders`); these ids are part of the illustration and are not
validated against the user's real catalog. v1 is descriptive/diagnostic — no
modeling tasks.

**Exemplar balance (2026-09-17).** In-context learning over-predicts whatever a
prompt's examples show most and shows last (majority-label and recency bias —
Zhao et al., ICML 2021), so the mix here is deliberate, not incidental:

- **Question language is mixed EN/ID on purpose — do NOT translate these back to
  English.** Real traffic and the `eval/planner` dataset are majority Indonesian
  (17 ID / 14 EN); an all-English exemplar set teaches the planner that Indonesian
  questions are out-of-distribution.
- **Plan internals stay English regardless of question language** — `goal_restated`,
  `objective`, `success_criteria` are internal artifacts. The ONE field that follows
  the user's language is `infeasible_reason` (see `prompt.py`: "If (and only if) you
  return an infeasible plan, write `infeasible_reason` in {reply_language}").
  Example H demonstrates that pairing; keep it demonstrated.
- **Don't stack near-identical exemplars.** Three examples already share the
  `retrieve_data -> analyze_aggregate` shape; adding a fourth on the same shape and
  the same fictional domain makes that plan the model's default answer to everything.
  A new example should earn its place by teaching a shape or a rule no other one does.
- **Feasible/infeasible pairs are load-bearing.** G↔H and J↔K each pair a capability
  with its refusal so an explicit ask can't talk the planner past a missing column.
  Keep both halves, and keep each pair's phrasing parallel.

See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3 (Examples A and B).
"""

from __future__ import annotations

from .schemas import Task, TaskList, ToolCall

# --------------------------------------------------------------------------- #
# Example A — exploratory, no modeling.
# "Which product categories drove last quarter's revenue?"
# Shows: retrieve_data pulls rows -> analyze_aggregate sums revenue per
# category in one call (no manual per-category queries).
# --------------------------------------------------------------------------- #

_EXAMPLE_A = TaskList(
    plan_id="example_a",
    goal_restated="Identify which product categories contributed most to last quarter's revenue.",
    assumptions=["'last quarter' = 2026-01-01 to 2026-03-31."],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes category, revenue, and order date.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; the 3 needed columns are present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull last quarter's order-level category and revenue rows.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_category", "alias": "category"},
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                            ],
                            "filters": [
                                {
                                    "column_id": "c_order_date",
                                    "op": "between",
                                    "value": ["2026-01-01", "2026-03-31"],
                                    "value_type": "date",
                                }
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="quarter_rows",
            success_criteria="Produced last quarter's order rows with category and revenue.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Sum revenue per category for the quarter.",
            tool_calls=[
                ToolCall(
                    tool="analyze_aggregate",
                    args={
                        "data": "${t2}",
                        "aggregations": {"revenue": ["sum"]},
                        "group_by": ["category"],
                    },
                )
            ],
            expected_output="category_revenue",
            success_criteria="Produced total revenue per category, one row each.",
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example B — descriptive / trend.
# "How has monthly revenue trended by region this year, and what's unusual?"
# --------------------------------------------------------------------------- #

_EXAMPLE_B = TaskList(
    plan_id="example_b",
    goal_restated="Describe this year's monthly revenue trend and flag unusual months.",
    assumptions=["'this year' starts 2026-01-01."],
    open_questions=["'Unusual' is interpreted as months far from the typical monthly revenue."],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes order date, revenue, and region.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; the needed columns are present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull this year's order dates, revenue, and region.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {
                                    "kind": "column",
                                    "column_id": "c_order_date",
                                    "alias": "order_date",
                                },
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                                {"kind": "column", "column_id": "c_region", "alias": "region"},
                            ],
                            "filters": [
                                {
                                    "column_id": "c_order_date",
                                    "op": ">=",
                                    "value": "2026-01-01",
                                    "value_type": "date",
                                }
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="ytd_rows",
            success_criteria="Produced this year's order-level rows with date, revenue, region.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Bucket revenue into months and summarize the trend and movement.",
            tool_calls=[
                ToolCall(
                    tool="analyze_trend",
                    args={
                        "data": "${t2}",
                        "date_column": "order_date",
                        "value_column": "revenue",
                        "freq": "month",
                        "agg": "sum",
                    },
                )
            ],
            expected_output="monthly_trend",
            success_criteria=(
                "Produced a per-month revenue series with direction and change rate to "
                "flag months above/below the typical level."
            ),
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example C — mixed structured + unstructured.
# "Revenue dipped in Q1 — what happened?"
# Shows: a structured branch (query -> analyze_trend) runs alongside an
# INDEPENDENT retrieve_knowledge branch that pulls qualitative context. Note
# retrieve_knowledge takes a natural-language `query` (NOT a `${t<id>}` data
# placeholder — it is a source, not a consumer) and can run in parallel; the
# Assembler folds the document context into the explanation.
# --------------------------------------------------------------------------- #

_EXAMPLE_C = TaskList(
    plan_id="example_c",
    goal_restated="Explain Q1's revenue dip using both the numbers and the qualitative record.",
    assumptions=["'Q1' = 2026-01-01 to 2026-03-31."],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes order date and revenue.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; date and revenue columns present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull Q1 order dates and revenue.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {
                                    "kind": "column",
                                    "column_id": "c_order_date",
                                    "alias": "order_date",
                                },
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                            ],
                            "filters": [
                                {
                                    "column_id": "c_order_date",
                                    "op": "between",
                                    "value": ["2026-01-01", "2026-03-31"],
                                    "value_type": "date",
                                }
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="q1_rows",
            success_criteria="Produced Q1 order rows with date and revenue.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Summarize the Q1 monthly revenue trend to locate the dip.",
            tool_calls=[
                ToolCall(
                    tool="analyze_trend",
                    args={
                        "data": "${t2}",
                        "date_column": "order_date",
                        "value_column": "revenue",
                        "freq": "month",
                        "agg": "sum",
                    },
                )
            ],
            expected_output="q1_trend",
            success_criteria="Produced a per-month revenue series showing where revenue fell.",
            depends_on=["t2"],
            estimated_cost="low",
        ),
        Task(
            id="t4",
            stage="data_understanding",
            objective="Retrieve qualitative context on Q1 operational events behind a dip.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_knowledge",
                    args={
                        "query": "operational issues, outages, or notable events in Q1 2026",
                        "top_k": 5,
                    },
                )
            ],
            expected_output="q1_context_chunks",
            success_criteria="Produced relevant document chunks about Q1 operations.",
            depends_on=[],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example D — group-by aggregation (analyze_aggregate arg shape).
# "Berapa rata-rata dan total nilai order per region?"  [ID — see the exemplar
# balance note in the module docstring; the plan itself stays English.]
# Shows the EXACT analyze_aggregate args: `aggregations` is an OBJECT mapping each
# column to a LIST of functions ({"revenue": ["mean", "sum"]}), and `group_by` is a
# SEPARATE array — NOT a nested list of metric specs. Supported funcs: sum, mean,
# count, min, max, median, nunique.
# --------------------------------------------------------------------------- #

_EXAMPLE_D = TaskList(
    plan_id="example_d",
    goal_restated="Report the average and total order value for each region.",
    assumptions=[],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes region and revenue.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; region and revenue present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull order-level region and revenue.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_region", "alias": "region"},
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="region_rows",
            success_criteria="Produced order rows with region and revenue.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Aggregate mean and total revenue per region.",
            tool_calls=[
                ToolCall(
                    tool="analyze_aggregate",
                    args={
                        "data": "${t2}",
                        "aggregations": {"revenue": ["mean", "sum"]},
                        "group_by": ["region"],
                    },
                )
            ],
            expected_output="region_aggregates",
            success_criteria="Produced one row per region with mean and total revenue.",
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example E — non-date filters (value_type is the ELEMENT type, never a container).
# "Total revenue for the East and West regions, counting orders of at least 100."
# Shows: an `in` filter over a list of strings uses value_type "string" (NOT
# "list"); a numeric comparison uses "decimal"; and every filter carries a
# value_type copied from the column's catalog [data_type].
# --------------------------------------------------------------------------- #

_EXAMPLE_E = TaskList(
    plan_id="example_e",
    goal_restated="Total revenue for the East and West regions, orders of at least 100.",
    assumptions=["'at least 100' filters order revenue >= 100."],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes region and revenue.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; region and revenue present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull East/West order rows of at least 100 with region and revenue.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_region", "alias": "region"},
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                            ],
                            "filters": [
                                {
                                    "column_id": "c_region",
                                    "op": "in",
                                    "value": ["East", "West"],
                                    "value_type": "string",
                                },
                                {
                                    "column_id": "c_revenue",
                                    "op": ">=",
                                    "value": 100,
                                    "value_type": "decimal",
                                },
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="filtered_rows",
            success_criteria="Produced East/West order rows above the revenue threshold.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Sum revenue per region.",
            tool_calls=[
                ToolCall(
                    tool="analyze_aggregate",
                    args={
                        "data": "${t2}",
                        "aggregations": {"revenue": ["sum"]},
                        "group_by": ["region"],
                    },
                )
            ],
            expected_output="region_revenue",
            success_criteria="Produced total revenue per region, one row each.",
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example F — descriptive statistics (analyze_descriptive).
# "Give me the summary statistics for order revenue and quantity."
# Shows: retrieve_data pulls the numeric columns -> analyze_descriptive summarizes
# them. The `data` comes from the retrieve_data task (t2), NEVER from the check_data
# inspection step (t1) — check_data returns column metadata, not data rows.
# --------------------------------------------------------------------------- #

_EXAMPLE_F = TaskList(
    plan_id="example_f",
    goal_restated="Summarize the distribution of order revenue and quantity.",
    assumptions=[],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes order revenue and quantity.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; the needed columns are present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Pull the order revenue and quantity rows.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
                                {"kind": "column", "column_id": "c_quantity", "alias": "quantity"},
                            ],
                            "limit": 10000,
                        }
                    },
                )
            ],
            expected_output="order_rows",
            success_criteria="Produced order-level rows with revenue and quantity.",
            depends_on=["t1"],
            estimated_cost="medium",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Summarize the distribution of revenue and quantity.",
            tool_calls=[
                ToolCall(
                    tool="analyze_descriptive",
                    args={
                        # `data` references t2 (retrieve_data rows), NOT t1 (check_data).
                        "data": "${t2}",
                        # column refs are the retrieve_data output aliases.
                        "column_ids": ["revenue", "quantity"],
                    },
                )
            ],
            expected_output="summary_stats",
            success_criteria=(
                "Produced mean/median/std/quartiles for revenue and quantity, above/below "
                "the typical range."
            ),
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example G — top-N ranking.
# "Top 3 product categories by total revenue."
# Shows: top-N is ONE retrieve_data query — group by the entity, aggregate the
# measure with an alias, order by that alias, limit N. NEVER a bare
# order-by-measure + limit (that ranks raw rows, so the same entity can appear
# twice — observed in production: "top 3 models" returned one model twice).
# --------------------------------------------------------------------------- #

_EXAMPLE_G = TaskList(
    plan_id="example_g",
    goal_restated="Rank product categories by total revenue and return the top 3.",
    assumptions=[],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes category and revenue.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria=(
                "Produced the orders table schema; category and revenue columns "
                "are present."
            ),
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Aggregate revenue per category, rank descending, keep the top 3.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_category", "alias": "category"},
                                {
                                    "kind": "agg",
                                    "fn": "sum",
                                    "column_id": "c_revenue",
                                    "alias": "total_revenue",
                                },
                            ],
                            "group_by": ["c_category"],
                            "order_by": [{"column_id": "total_revenue", "dir": "desc"}],
                            "limit": 3,
                        }
                    },
                )
            ],
            expected_output="top3_categories",
            success_criteria=(
                "Produced at most 3 rows, one distinct category each, ranked by "
                "total revenue."
            ),
            depends_on=["t1"],
            estimated_cost="low",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example H — infeasible question (see planner.md "When the catalog cannot
# answer"). "Berapa tingkat churn pelanggan kita?" against a sales catalog with no
# subscription/churn data: no task list is forced onto unrelated columns;
# instead `infeasible_reason` states the gap + the nearest available data.
# ALSO the one example of the reply-language rule: the question is Indonesian, so
# `infeasible_reason` is Indonesian — while every other field stays English,
# because only this field is user-facing prose (see `prompt.py`). Keep it that way.
# --------------------------------------------------------------------------- #

_EXAMPLE_H = TaskList(
    plan_id="example_h",
    goal_restated="Measure the customer churn rate.",
    assumptions=[],
    open_questions=[],
    tasks=[],
    infeasible_reason=(
        "Sumber data yang terhubung tidak memuat data churn atau status langganan — "
        "tabel orders hanya menyimpan kategori, revenue, kuantitas, dan tanggal di "
        "level order. Analisis terdekat yang tersedia: perilaku pembelian ulang, "
        "atau revenue per pelanggan dari waktu ke waktu."
    ),
)


# --------------------------------------------------------------------------- #
# Example I — combine two measures per entity (KM-703).
# "Which category has both the highest revenue and the highest average order
# quantity?" Shows: each measure is computed in its OWN grouped retrieve_data
# task (a "${t<id>}" placeholder resolves to a task's LAST output, so the two
# retrievals must be separate tasks), then analyze_merge aligns them on the
# shared entity alias. The merged table answers "both A and B" questions that
# a single query cannot express.
# --------------------------------------------------------------------------- #

_EXAMPLE_I = TaskList(
    plan_id="example_i",
    goal_restated=(
        "Identify the product category with both the highest total revenue and the "
        "highest average order quantity."
    ),
    assumptions=[],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes category, revenue, and quantity.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria=(
                "Produced the orders table schema; category, revenue, and quantity "
                "columns are present."
            ),
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Total revenue per category.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_category", "alias": "category"},
                                {
                                    "kind": "agg",
                                    "fn": "sum",
                                    "column_id": "c_revenue",
                                    "alias": "total_revenue",
                                },
                            ],
                            "group_by": ["c_category"],
                        }
                    },
                )
            ],
            expected_output="revenue_per_category",
            success_criteria="Produced one total-revenue row per category.",
            depends_on=["t1"],
            estimated_cost="low",
        ),
        Task(
            id="t3",
            stage="data_preparation",
            objective="Average order quantity per category.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_category", "alias": "category"},
                                {
                                    "kind": "agg",
                                    "fn": "avg",
                                    "column_id": "c_quantity",
                                    "alias": "avg_quantity",
                                },
                            ],
                            "group_by": ["c_category"],
                        }
                    },
                )
            ],
            expected_output="quantity_per_category",
            success_criteria="Produced one average-quantity row per category.",
            depends_on=["t1"],
            estimated_cost="low",
        ),
        Task(
            id="t4",
            stage="evaluation",
            objective="Align both measures per category to find the category leading on both.",
            tool_calls=[
                ToolCall(
                    tool="analyze_merge",
                    args={
                        "data": "${t2}",
                        "data_right": "${t3}",
                        "on": ["category"],
                    },
                )
            ],
            expected_output="combined_measures",
            success_criteria=(
                "Produced one row per category carrying both total_revenue and "
                "avg_quantity."
            ),
            depends_on=["t2", "t3"],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example J — visualization tail (SPINE_V2_PLAN §4.3, recipe "viz tail").
# "Show me a bar chart of total revenue per product category."
# Shows: render_chart is planned ONLY on an explicit ask ("bar chart"), and is
# ALWAYS a tail step — its `data` consumes a TABLE-producing upstream (here the
# grouped top-N-style retrieve from Example G, without the limit), never stats/
# series/metadata. `x`/`y` reference that table's column aliases; the chart
# carries the already-aggregated rows (one bar per category), never raw order
# rows. The `assumptions` line carries the feasibility check on purpose: a chart
# ask never licenses aliasing a stand-in column (see Example K).
# --------------------------------------------------------------------------- #

_EXAMPLE_J = TaskList(
    plan_id="example_j",
    goal_restated="Chart total revenue per product category as a bar chart.",
    assumptions=[
        "The chart dimension exists in the catalog (c_category) — a chart ask is "
        "planned only when the asked-for dimension and measure have real catalog "
        "columns; otherwise it is infeasible, chart or no chart."
    ],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes category and revenue.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; category and revenue present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Aggregate total revenue per product category.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_category", "alias": "category"},
                                {
                                    "kind": "agg",
                                    "fn": "sum",
                                    "column_id": "c_revenue",
                                    "alias": "total_revenue",
                                },
                            ],
                            "group_by": ["c_category"],
                        }
                    },
                )
            ],
            expected_output="revenue_per_category",
            success_criteria="Produced one total-revenue row per category.",
            depends_on=["t1"],
            estimated_cost="low",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Render the per-category revenue table as a bar chart.",
            tool_calls=[
                ToolCall(
                    tool="render_chart",
                    args={
                        # `data` references the TABLE task (t2) — a chart is always a
                        # tail on an aggregated table, never on raw rows or stats.
                        "data": "${t2}",
                        "chart_type": "bar",
                        "x": "category",
                        "y": "total_revenue",
                        "title": "Total revenue by product category",
                    },
                )
            ],
            expected_output="revenue_bar_chart",
            success_criteria="Produced a bar-chart spec with one bar per category.",
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example K — a chart ask that is INFEASIBLE (SPINE_V2_PLAN §4.3 + planner.md
# "Charts ... never relaxes feasibility"). "Plot the customer churn rate by
# month as a line chart" against the sales catalog: churn does not exist, and
# the explicit chart request does NOT license mapping some other column into
# the asked-for measure — the verdict is the same infeasible_reason Example H
# would give, chart or no chart.
# --------------------------------------------------------------------------- #

_EXAMPLE_K = TaskList(
    plan_id="example_k",
    goal_restated="Chart the monthly customer churn rate as a line chart.",
    assumptions=[],
    open_questions=[],
    tasks=[],
    infeasible_reason=(
        "The connected source has no churn or subscription-status data — the "
        "orders table only carries order-level category, revenue, quantity, and "
        "dates — so there is nothing to chart. Nearest chartable analyses: "
        "monthly revenue trend, or order counts per category."
    ),
)


# --------------------------------------------------------------------------- #
# Example L — scalar count with a filter (no grouping, no analyze step).
# "Ada berapa order yang revenue-nya nol?"  [ID — R2b "berapa banyak …" is the
# recipe users most often reach for in Indonesian; the plan stays English.]
# Shows: a "how many rows match X" question is answered by a SINGLE retrieve_data
# IR whose select is a COUNT(*) aggregate ({"kind": "agg", "fn": "count"}, with
# column_id omitted) plus the filter — it returns the exact number in one row.
# Do NOT select the raw column and count the returned rows: that caps at `limit`
# and leaves the tally to be eyeballed. count(*) omits column_id (the validator
# lets only 'count' do so); no group_by is needed for a scalar count.
# --------------------------------------------------------------------------- #

_EXAMPLE_L = TaskList(
    plan_id="example_l",
    goal_restated="Count how many orders have revenue equal to 0.",
    assumptions=[],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes order revenue.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; revenue is present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="evaluation",
            objective="Count the orders whose revenue equals 0.",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "agg", "fn": "count", "alias": "order_count"}
                            ],
                            "filters": [
                                {
                                    "column_id": "c_revenue",
                                    "op": "=",
                                    "value": 0,
                                    "value_type": "decimal",
                                }
                            ],
                        }
                    },
                )
            ],
            expected_output="zero_revenue_count",
            success_criteria="Produced one row holding the count of zero-revenue orders.",
            depends_on=["t1"],
            estimated_cost="low",
        ),
    ],
)


# --------------------------------------------------------------------------- #
# Example M — line/TREND chart over time (viz tail on a time series).
# "Tampilkan grafik garis rata-rata revenue dari waktu ke waktu."  [ID — also keeps
# this from reading as a near-duplicate of Example J's English bar-chart ask.]
# Shows: a trend line chart aggregates PER TIME PERIOD first — group_by the date
# column + aggregate the measure in the retrieve_data IR (one row per date), THEN
# tail render_chart. Never chart raw per-record rows over time (many rows per date
# = an unreadable vertical smear). Mirrors Example J (bar) but grouping by the date
# instead of a category. (analyze_trend -> render_chart is an equally valid shape.)
# --------------------------------------------------------------------------- #

_EXAMPLE_M = TaskList(
    plan_id="example_m",
    goal_restated="Chart average order revenue per day as a line chart over time.",
    assumptions=["The date and revenue columns exist in the catalog (c_order_date, c_revenue)."],
    open_questions=[],
    tasks=[
        Task(
            id="t1",
            stage="data_understanding",
            objective="Confirm the sales source exposes order date and revenue.",
            tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
            expected_output="source_shape",
            success_criteria="Produced the orders table schema; order date and revenue present.",
            depends_on=[],
            estimated_cost="low",
        ),
        Task(
            id="t2",
            stage="data_preparation",
            objective="Aggregate average revenue per day (one row per date).",
            tool_calls=[
                ToolCall(
                    tool="retrieve_data",
                    args={
                        "ir": {
                            "source_id": "src_sales",
                            "table_id": "t_orders",
                            "select": [
                                {"kind": "column", "column_id": "c_order_date", "alias": "order_date"},
                                {"kind": "agg", "fn": "avg", "column_id": "c_revenue", "alias": "avg_revenue"},
                            ],
                            "group_by": ["c_order_date"],
                            "order_by": [{"column_id": "c_order_date", "dir": "asc"}],
                        }
                    },
                )
            ],
            expected_output="daily_revenue",
            success_criteria="Produced one average-revenue row per date, ordered by date.",
            depends_on=["t1"],
            estimated_cost="low",
        ),
        Task(
            id="t3",
            stage="evaluation",
            objective="Render the daily average-revenue series as a line chart.",
            tool_calls=[
                ToolCall(
                    tool="render_chart",
                    args={
                        # `data` is the AGGREGATED daily table (t2) — one point per date,
                        # never the raw per-order rows.
                        "data": "${t2}",
                        "chart_type": "line",
                        "x": "order_date",
                        "y": "avg_revenue",
                        "title": "Average revenue over time",
                    },
                )
            ],
            expected_output="revenue_line_chart",
            success_criteria="Produced a line-chart spec with one point per date.",
            depends_on=["t2"],
            estimated_cost="low",
        ),
    ],
)


# =========================================================================== #
# Examples N–W — one per tool that had no exemplar (A18, 2026-09-17).
#
# These deliberately do NOT open with `check_data`. All 13 earlier examples do,
# making it 13/13 — a prior strong enough that the planner emits a catalog task
# it rarely needs (planner.md: "smallest plan that answers the question"; the
# recipe chains are `retrieve_data -> analyze_*`, with no `check_data`). Ten more
# would deepen the imbalance the exemplar-balance note above already flags.
#
# Stages follow the rule rewritten in planner.md: `modeling` ONLY for a task that
# actually FITS a model (forecast, importance, cluster, anomaly). A statistic
# (correlation), an arithmetic breakdown (pareto), a rule-based split
# (segmentation) and a hypothesis test (ab_test) are `evaluation`; profiling is
# `data_understanding`.
#
# None carries a `render_chart` tail: the ML tools ship their own chart, so a tail
# would demonstrate exactly the double-chart rule 5 forbids.
# =========================================================================== #


def _ir(select, filters=None, group_by=None, limit=10000):
    """Compact inline QueryIR for the examples below (same shape as A–M)."""
    ir = {"source_id": "src_sales", "table_id": "t_orders", "select": select, "limit": limit}
    if filters:
        ir["filters"] = filters
    if group_by:
        ir["group_by"] = group_by
    return {"ir": ir}


def _col(cid, alias):
    return {"kind": "column", "column_id": cid, "alias": alias}


def _agg(fn, alias, cid=None):
    out = {"kind": "agg", "fn": fn, "alias": alias}
    if cid:
        out["column_id"] = cid
    return out


def _pull(objective, select, expected, criteria, **kw):
    return Task(
        id="t1",
        stage="data_preparation",
        objective=objective,
        tool_calls=[ToolCall(tool="retrieve_data", args=_ir(select, **kw))],
        expected_output=expected,
        success_criteria=criteria,
        depends_on=[],
        estimated_cost="medium",
    )


def _step(stage, objective, tool, args, expected, criteria, cost="low"):
    return Task(
        id="t2",
        stage=stage,
        objective=objective,
        tool_calls=[ToolCall(tool=tool, args={"data": "${t1}", **args})],
        expected_output=expected,
        success_criteria=criteria,
        depends_on=["t1"],
        estimated_cost=cost,
    )


# --------------------------------------------------------------------------- #
# Example N — forecast (timeseries). "Proyeksikan revenue 6 bulan ke depan."
# Pairs against Example B (trend): the SAME retrieve shape — raw date + value
# rows, the tool buckets them itself via `freq`/`agg` — but the question projects
# forward instead of describing the past. Past vs project is the entire
# distinction between `analyze_trend` and `analyze_forecast`.
# --------------------------------------------------------------------------- #

_EXAMPLE_N = TaskList(
    plan_id="example_n",
    goal_restated="Project monthly revenue six months beyond the last actual month.",
    assumptions=[],
    open_questions=[],
    tasks=[
        _pull(
            "Pull order dates and revenue across the available history.",
            [_col("c_order_date", "order_date"), _col("c_revenue", "revenue")],
            "history_rows",
            "Produced dated revenue rows covering the available history.",
        ),
        _step(
            "modeling",
            "Fit a forecast on monthly revenue and project six months ahead.",
            "analyze_forecast",
            {
                "date_column": "order_date",
                "value_column": "revenue",
                "horizon": 6,
                "freq": "month",
                "agg": "sum",
            },
            "revenue_forecast",
            "Produced six future monthly points plus the backtest score.",
            cost="medium",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example O — correlation (relationship). "Is order quantity related to revenue?"
# Pairs against Example P, which uses an identical retrieve: correlation answers
# how TWO NAMED columns move together, so `column_ids` lists exactly that pair.
# --------------------------------------------------------------------------- #

_EXAMPLE_O = TaskList(
    plan_id="example_o",
    goal_restated="Measure how order quantity and revenue move together.",
    assumptions=[],
    open_questions=[],
    tasks=[
        _pull(
            "Pull order-level quantity and revenue.",
            [_col("c_quantity", "quantity"), _col("c_revenue", "revenue")],
            "pair_rows",
            "Produced order rows carrying both quantity and revenue.",
        ),
        _step(
            "evaluation",
            "Correlate quantity against revenue.",
            "analyze_correlation",
            {"column_ids": ["quantity", "revenue"]},
            "quantity_revenue_correlation",
            "Produced a correlation coefficient for the pair.",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example P — importance (relationship). "Apa yang paling mempengaruhi revenue?"
# Deliberately the SAME retrieve as Example O, so the contrast is the question
# alone: correlation takes TWO NAMED columns; importance takes ONE TARGET and
# ranks every other numeric column against it. "How do X and Y relate" -> O;
# "what most affects X" -> P. A `modeling` task — it fits a model.
# --------------------------------------------------------------------------- #

_EXAMPLE_P = TaskList(
    plan_id="example_p",
    goal_restated="Rank which order attributes most explain revenue.",
    assumptions=[],
    open_questions=[],
    tasks=[
        _pull(
            "Pull revenue alongside its candidate explanatory columns.",
            [_col("c_quantity", "quantity"), _col("c_revenue", "revenue")],
            "feature_rows",
            "Produced rows carrying revenue and its candidate explainers.",
        ),
        _step(
            "modeling",
            "Rank the columns that explain revenue.",
            "analyze_importance",
            {"target_column": "revenue"},
            "revenue_drivers",
            "Produced ranked features with a holdout score.",
            cost="medium",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example Q — pareto (relationship). "80% revenue kita dari kategori apa?"
# Contrast with Example P: pareto is ARITHMETIC share of a total — which few
# categories add up to most of it — where importance is statistical explanation.
# A "which few account for most" question is always pareto, never importance.
# --------------------------------------------------------------------------- #

_EXAMPLE_Q = TaskList(
    plan_id="example_q",
    goal_restated="Identify the few categories that make up most of total revenue.",
    assumptions=[],
    open_questions=[],
    tasks=[
        _pull(
            "Pull category and revenue per order.",
            [_col("c_category", "category"), _col("c_revenue", "revenue")],
            "category_rows",
            "Produced order rows with category and revenue.",
        ),
        _step(
            "evaluation",
            "Break total revenue down by category, largest first.",
            "analyze_pareto",
            {"dimension": "category", "value_column": "revenue", "agg": "sum"},
            "revenue_pareto",
            "Produced categories with share and cumulative share of the total.",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example R — segmentation (grouping). "Split customers into spend tiers
# 0-100, 100-500, 500+."
# Pairs against Example S. The user STATED the rule, so the split is
# deterministic: explicit `bins` edges and `labels`. Note the retrieve groups to
# one row per customer first — the tiers are per customer, not per order.
# --------------------------------------------------------------------------- #

_EXAMPLE_R = TaskList(
    plan_id="example_r",
    goal_restated="Split customers into the three stated spend tiers.",
    assumptions=["Tier boundaries are the ones the user gave: 0-100, 100-500, 500+."],
    open_questions=[],
    tasks=[
        _pull(
            "Total revenue per customer.",
            [
                _col("c_customer_id", "customer_id"),
                _agg("sum", "total_revenue", "c_revenue"),
            ],
            "customer_totals",
            "Produced one row per customer with their total revenue.",
            group_by=["c_customer_id"],
        ),
        _step(
            "evaluation",
            "Bin customers into the stated spend tiers.",
            "analyze_segmentation",
            {
                "column": "total_revenue",
                "bins": [0, 100, 500, 1000000],
                "method": "edges",
                "labels": ["0-100", "100-500", "500+"],
            },
            "spend_tiers",
            "Produced a customer count per stated tier.",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example S — cluster (grouping). "Ada berapa kelompok pelanggan secara alami?"
# Pairs against Example R. The user gave NO rule, so the groups are discovered:
# no bins, no labels — just the columns describing each customer, and the tool
# picks k itself. Rule given -> R; no rule -> S. A `modeling` task.
# --------------------------------------------------------------------------- #

_EXAMPLE_S = TaskList(
    plan_id="example_s",
    goal_restated="Discover how many natural customer groups the data contains.",
    assumptions=[],
    open_questions=[],
    tasks=[
        _pull(
            "Build a per-customer profile of spend and volume.",
            [
                _col("c_customer_id", "customer_id"),
                _agg("sum", "total_revenue", "c_revenue"),
                _agg("sum", "total_quantity", "c_quantity"),
            ],
            "customer_profiles",
            "Produced one row per customer with spend and volume.",
            group_by=["c_customer_id"],
        ),
        _step(
            "modeling",
            "Discover the natural customer groups.",
            "analyze_cluster",
            {"column_ids": ["total_revenue", "total_quantity"]},
            "customer_clusters",
            "Produced a chosen k with its silhouette and one profile per cluster.",
            cost="medium",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example T — A/B test (evaluate). "Revenue East vs West, bedanya signifikan?"
# The question names TWO specific groups and asks whether the gap is real, so
# the plan retrieves the dimension plus the measure and hands both group values
# to the tool. "Is the difference significant" is always this tool — a plain
# total per region would be `analyze_aggregate` instead.
# --------------------------------------------------------------------------- #

_EXAMPLE_T = TaskList(
    plan_id="example_t",
    goal_restated="Test whether East and West differ significantly on revenue.",
    assumptions=[],
    open_questions=[],
    tasks=[
        _pull(
            "Pull region and revenue for the two regions being compared.",
            [_col("c_region", "region"), _col("c_revenue", "revenue")],
            "region_rows",
            "Produced order rows carrying region and revenue.",
            filters=[
                {
                    "column_id": "c_region",
                    "op": "in",
                    "value": ["East", "West"],
                    "value_type": "string",
                }
            ],
        ),
        _step(
            "evaluation",
            "Test East against West on revenue.",
            "analyze_ab_test",
            {
                "dimension": "region",
                "value_column": "revenue",
                "group_a": "East",
                "group_b": "West",
                "agg": "sum",
            },
            "east_west_test",
            "Produced both group values, the gap, and a significance verdict.",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example U — anomaly (evaluate). "Any unusual orders this month?"
# Unusual ROWS, not a comparison between groups: the plan pulls the numeric
# columns that define "normal" and the tool flags the outliers among them. A
# `modeling` task — it fits a model.
# --------------------------------------------------------------------------- #

_EXAMPLE_U = TaskList(
    plan_id="example_u",
    goal_restated="Find this month's orders that look unusual.",
    assumptions=["'this month' = 2026-03-01 to 2026-03-31."],
    open_questions=[],
    tasks=[
        _pull(
            "Pull this month's order revenue and quantity.",
            [_col("c_revenue", "revenue"), _col("c_quantity", "quantity")],
            "month_rows",
            "Produced this month's orders with revenue and quantity.",
            filters=[
                {
                    "column_id": "c_order_date",
                    "op": "between",
                    "value": ["2026-03-01", "2026-03-31"],
                    "value_type": "date",
                }
            ],
        ),
        _step(
            "modeling",
            "Flag the unusual orders.",
            "analyze_anomaly",
            {"column_ids": ["revenue", "quantity"]},
            "unusual_orders",
            "Produced the flagged rows with their anomaly scores.",
            cost="medium",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example V — data quality (explore). "Data ordernya bersih nggak? Ada yang kosong?"
# A question ABOUT the data rather than an answer computed from it. Still needs a
# retrieve — the tool profiles rows it is handed, it never self-fetches — and the
# profiling task is `data_understanding`, not `evaluation`.
# --------------------------------------------------------------------------- #

_EXAMPLE_V = TaskList(
    plan_id="example_v",
    goal_restated="Assess how complete and clean the order data is.",
    assumptions=[],
    open_questions=[],
    tasks=[
        _pull(
            "Pull the order columns to be profiled.",
            [
                _col("c_order_date", "order_date"),
                _col("c_category", "category"),
                _col("c_revenue", "revenue"),
                _col("c_quantity", "quantity"),
            ],
            "profile_rows",
            "Produced order rows across the columns to be profiled.",
        ),
        _step(
            "data_understanding",
            "Profile the columns for nulls, types and outliers.",
            "analyze_data_quality",
            {},
            "order_data_quality",
            "Produced per-column completeness and any flagged issues.",
        ),
    ],
)

# --------------------------------------------------------------------------- #
# Example W — infeasible METHOD (recipe R15). "Which customer will churn next?"
# Distinct from Example H, where the measure simply does not exist. Here the data
# is fine; the METHOD does not apply. This product forecasts a measure over time
# and ranks what explains one — it does not predict which individual entity a
# future event will happen to. An explicit ask for a prediction never relaxes
# that (planner.md rule 5), exactly as a chart ask never relaxes feasibility.
# --------------------------------------------------------------------------- #

_EXAMPLE_W = TaskList(
    plan_id="example_w",
    goal_restated="Predict which individual customer will churn next.",
    assumptions=[],
    open_questions=[],
    tasks=[],
    infeasible_reason=(
        "Naming the next customer to churn means predicting a future event for one "
        "entity, which none of the available tools do — they project a measure over "
        "time and rank what explains one. The orders table also carries no churn or "
        "subscription-status column to learn from. Nearest available analyses: revenue "
        "per customer over time, or which attributes most explain order value."
    ),
)


EXAMPLES: list[tuple[str, TaskList]] = [
    ("Which product categories drove last quarter's revenue?", _EXAMPLE_A),
    ("How has monthly revenue trended by region this year, and what's unusual?", _EXAMPLE_B),
    ("Revenue dipped in Q1 — what happened?", _EXAMPLE_C),
    ("Berapa rata-rata dan total nilai order per region?", _EXAMPLE_D),
    ("Total revenue for the East and West regions, counting orders of at least 100.", _EXAMPLE_E),
    ("Give me the summary statistics for order revenue and quantity.", _EXAMPLE_F),
    ("Which 3 product categories have the best revenue performance?", _EXAMPLE_G),
    ("Berapa tingkat churn pelanggan kita?", _EXAMPLE_H),
    (
        "Which product category has both the highest revenue and the highest average "
        "order quantity?",
        _EXAMPLE_I,
    ),
    ("Show me a bar chart of total revenue per product category.", _EXAMPLE_J),
    ("Plot the customer churn rate by month as a line chart.", _EXAMPLE_K),
    ("Ada berapa order yang revenue-nya nol?", _EXAMPLE_L),
    ("Tampilkan grafik garis rata-rata revenue dari waktu ke waktu.", _EXAMPLE_M),
    ("Proyeksikan revenue 6 bulan ke depan.", _EXAMPLE_N),
    ("Is order quantity related to revenue?", _EXAMPLE_O),
    ("Apa yang paling mempengaruhi revenue?", _EXAMPLE_P),
    ("80% revenue kita dari kategori apa?", _EXAMPLE_Q),
    ("Split customers into spend tiers 0-100, 100-500, and 500+.", _EXAMPLE_R),
    ("Ada berapa kelompok pelanggan secara alami?", _EXAMPLE_S),
    ("Revenue East vs West, bedanya signifikan?", _EXAMPLE_T),
    ("Any unusual orders this month?", _EXAMPLE_U),
    ("Data ordernya bersih nggak? Ada yang kosong?", _EXAMPLE_V),
    ("Which customer will churn next?", _EXAMPLE_W),
]


def _tools_used(plan: TaskList) -> set[str]:
    """Every tool name a plan calls (empty for an infeasible plan)."""
    return {call.tool for task in plan.tasks for call in task.tool_calls}


def render_examples(visible_tools: set[str] | None = None) -> str:
    """Render the few-shots as text for the planner prompt.

    `visible_tools` is the name set the planner can actually call this turn. Group
    scoping (A17) shows the planner ~8 of 18 tools, so an example demonstrating a tool
    outside that slice teaches it to reach for something the validator will reject —
    and A19 established that these examples dominate when they and the rules disagree.
    Examples whose tools are not all visible are therefore dropped.

    Dropping is safe because it never starves the prompt: the narrowest group still
    keeps 7 of 13, and the two infeasible examples survive EVERY scope for free, since
    an infeasible plan calls no tools. Those are the ones that matter most — they teach
    the refusal discipline, and the feasible/infeasible pairs stay intact.

    `None` renders everything (the unscoped/fail-open path, and back-compat).
    """
    shown = [
        (q, p)
        for q, p in EXAMPLES
        if visible_tools is None or _tools_used(p) <= visible_tools
    ]
    blocks: list[str] = []
    for i, (question, plan) in enumerate(shown, start=1):
        blocks.append(
            f"## Example {i}\n\n"
            f"Question:\n{question}\n\n"
            f"TaskList:\n{plan.model_dump_json(indent=2)}"
        )
    return "\n\n".join(blocks)