Agentic-Service-Data-Eyond-Catalog

Running

Agentic-Service-Data-Eyond-Catalog / src /agents /planner /examples.py

Rifqi Hafizuddin

Claude Opus 4.8

[NOTICKET][AI] Planner: drop parallelizable_with from Task

272a186 5 months ago

15.7 kB

"""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 — `query_structured`
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: query_structured pulls rows -> analyze_contribution computes each
# category's share of the total in one call (no manual per-category + total
# 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="describe_source", 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="query_structured",
                    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="Rank each category's revenue share of the quarter total.",
            tool_calls=[
                ToolCall(
                    tool="analyze_contribution",
                    args={
                        "data": "${t2}",
                        "dimension": "category",
                        "value_column": "revenue",
                        "agg": "sum",
                    },
                )
            ],
            expected_output="category_contribution",
            success_criteria="Produced each category's revenue share, ranked high to low.",
            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="describe_source", 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="query_structured",
                    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_documents branch that pulls qualitative context. Note
# retrieve_documents 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="describe_source", 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="query_structured",
                    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_documents",
                    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="describe_source", 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="query_structured",
                    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)

	"""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 — `query_structured`
	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: query_structured pulls rows -> analyze_contribution computes each
	# category's share of the total in one call (no manual per-category + total
	# 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="describe_source", 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="query_structured",
	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="Rank each category's revenue share of the quarter total.",
	tool_calls=[
	ToolCall(
	tool="analyze_contribution",
	args={
	"data": "${t2}",
	"dimension": "category",
	"value_column": "revenue",
	"agg": "sum",
	},
	)
	],
	expected_output="category_contribution",
	success_criteria="Produced each category's revenue share, ranked high to low.",
	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="describe_source", 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="query_structured",
	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_documents branch that pulls qualitative context. Note
	# retrieve_documents 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="describe_source", 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="query_structured",
	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_documents",
	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="describe_source", 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="query_structured",
	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)