src/agents/planner/examples.py

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

Two 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.

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).
# "What is the average and total order value per region?"
# 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",
        ),
    ],
)


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),
    ("What is the average and total order value per region?", _EXAMPLE_D),
]


def render_examples() -> str:
    """Render the few-shots as text for the planner prompt."""
    blocks: list[str] = []
    for i, (question, plan) in enumerate(EXAMPLES, 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)