Rifqi Hafizuddin Claude Opus 4.8 commited on
Commit
f07e524
·
1 Parent(s): 8cb3714

[KM-626][AI] Planner few-shot: add analyze_aggregate (Example D) to fix arg shape

Browse files

Live test surfaced an agent-side bug: for group-by questions the Planner emitted
analyze_aggregate args as a nested list of metric specs, but the tool expects
`aggregations` as an OBJECT {column: [funcs]} with `group_by` as a SEPARATE array.
The tool then did aggregations.keys() on a list -> AttributeError. There was no
few-shot for analyze_aggregate, so the LLM guessed the wrong structure.

Add Example D ("average and total order value per region") demonstrating the exact
shape: aggregations={"revenue": ["mean","sum"]}, group_by=["region"].

Verified live: the same question that failed now plans the correct shape
({"base_price": ["mean"]}, group_by=["category"]) and analyze_aggregate succeeds
end-to-end. Planner suite green (32 passed).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Files changed (1) hide show
  1. src/agents/planner/examples.py +77 -0
src/agents/planner/examples.py CHANGED
@@ -301,10 +301,87 @@ _EXAMPLE_C = TaskList(
301
  )
302
 
303
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
304
  EXAMPLES: list[tuple[str, TaskList]] = [
305
  ("Which product categories drove last quarter's revenue?", _EXAMPLE_A),
306
  ("How has monthly revenue trended by region this year, and what's unusual?", _EXAMPLE_B),
307
  ("Revenue dipped in Q1 — what happened?", _EXAMPLE_C),
 
308
  ]
309
 
310
 
 
301
  )
302
 
303
 
304
+ # --------------------------------------------------------------------------- #
305
+ # Example D — group-by aggregation (analyze_aggregate arg shape).
306
+ # "What is the average and total order value per region?"
307
+ # Shows the EXACT analyze_aggregate args: `aggregations` is an OBJECT mapping each
308
+ # column to a LIST of functions ({"revenue": ["mean", "sum"]}), and `group_by` is a
309
+ # SEPARATE array — NOT a nested list of metric specs. Supported funcs: sum, mean,
310
+ # count, min, max, median, nunique.
311
+ # --------------------------------------------------------------------------- #
312
+
313
+ _EXAMPLE_D = TaskList(
314
+ plan_id="example_d",
315
+ goal_restated="Report the average and total order value for each region.",
316
+ assumptions=[],
317
+ open_questions=[],
318
+ tasks=[
319
+ Task(
320
+ id="t1",
321
+ stage="data_understanding",
322
+ objective="Confirm the sales source exposes region and revenue.",
323
+ tool_calls=[ToolCall(tool="describe_source", args={"source_id": "src_sales"})],
324
+ expected_output="source_shape",
325
+ success_criteria="Produced the orders table schema; region and revenue present.",
326
+ depends_on=[],
327
+ parallelizable_with=[],
328
+ estimated_cost="low",
329
+ ),
330
+ Task(
331
+ id="t2",
332
+ stage="data_preparation",
333
+ objective="Pull order-level region and revenue.",
334
+ tool_calls=[
335
+ ToolCall(
336
+ tool="query_structured",
337
+ args={
338
+ "ir": {
339
+ "source_id": "src_sales",
340
+ "table_id": "t_orders",
341
+ "select": [
342
+ {"kind": "column", "column_id": "c_region", "alias": "region"},
343
+ {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
344
+ ],
345
+ "limit": 10000,
346
+ }
347
+ },
348
+ )
349
+ ],
350
+ expected_output="region_rows",
351
+ success_criteria="Produced order rows with region and revenue.",
352
+ depends_on=["t1"],
353
+ parallelizable_with=[],
354
+ estimated_cost="medium",
355
+ ),
356
+ Task(
357
+ id="t3",
358
+ stage="evaluation",
359
+ objective="Aggregate mean and total revenue per region.",
360
+ tool_calls=[
361
+ ToolCall(
362
+ tool="analyze_aggregate",
363
+ args={
364
+ "data": "${t2}",
365
+ "aggregations": {"revenue": ["mean", "sum"]},
366
+ "group_by": ["region"],
367
+ },
368
+ )
369
+ ],
370
+ expected_output="region_aggregates",
371
+ success_criteria="Produced one row per region with mean and total revenue.",
372
+ depends_on=["t2"],
373
+ parallelizable_with=[],
374
+ estimated_cost="low",
375
+ ),
376
+ ],
377
+ )
378
+
379
+
380
  EXAMPLES: list[tuple[str, TaskList]] = [
381
  ("Which product categories drove last quarter's revenue?", _EXAMPLE_A),
382
  ("How has monthly revenue trended by region this year, and what's unusual?", _EXAMPLE_B),
383
  ("Revenue dipped in Q1 — what happened?", _EXAMPLE_C),
384
+ ("What is the average and total order value per region?", _EXAMPLE_D),
385
  ]
386
 
387