sofhiaazzhr commited on
Commit
958c76f
Β·
1 Parent(s): 5a010a6

[NOTICKET] fix(planner): reject analyze_* fed from metadata/documents tools

Browse files

The planner sometimes wired an analyze_* tool's 'data' to check_data (catalog metadata) instead of retrieve_data (rows); the plan passed structural validation but analyze_* then failed to find its columns at runtime (silent ColumnNotFoundError). Add PlannerValidator Check 9: an analyze_* 'data' placeholder must not reference a catalog.introspection (check_data/check_knowledge) or retrieval.documents (retrieve_knowledge) task β€” a category denylist, so retrieve_data and table-producing analyze_* chains stay valid. Reinforce with prompt guidance (planner.md: never feed analyze_* from check_data) + a new analyze_descriptive few-shot (Example F). Also log a compact plan dump on 'analysis planned' for diagnosability.

src/agents/planner/examples.py CHANGED
@@ -448,12 +448,89 @@ _EXAMPLE_E = TaskList(
448
  )
449
 
450
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
451
  EXAMPLES: list[tuple[str, TaskList]] = [
452
  ("Which product categories drove last quarter's revenue?", _EXAMPLE_A),
453
  ("How has monthly revenue trended by region this year, and what's unusual?", _EXAMPLE_B),
454
  ("Revenue dipped in Q1 β€” what happened?", _EXAMPLE_C),
455
  ("What is the average and total order value per region?", _EXAMPLE_D),
456
  ("Total revenue for the East and West regions, counting orders of at least 100.", _EXAMPLE_E),
 
457
  ]
458
 
459
 
 
448
  )
449
 
450
 
451
+ # --------------------------------------------------------------------------- #
452
+ # Example F β€” descriptive statistics (analyze_descriptive).
453
+ # "Give me the summary statistics for order revenue and quantity."
454
+ # Shows: retrieve_data pulls the numeric columns -> analyze_descriptive summarizes
455
+ # them. The `data` comes from the retrieve_data task (t2), NEVER from the check_data
456
+ # inspection step (t1) β€” check_data returns column metadata, not data rows.
457
+ # --------------------------------------------------------------------------- #
458
+
459
+ _EXAMPLE_F = TaskList(
460
+ plan_id="example_f",
461
+ goal_restated="Summarize the distribution of order revenue and quantity.",
462
+ assumptions=[],
463
+ open_questions=[],
464
+ tasks=[
465
+ Task(
466
+ id="t1",
467
+ stage="data_understanding",
468
+ objective="Confirm the sales source exposes order revenue and quantity.",
469
+ tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
470
+ expected_output="source_shape",
471
+ success_criteria="Produced the orders table schema; the needed columns are present.",
472
+ depends_on=[],
473
+ estimated_cost="low",
474
+ ),
475
+ Task(
476
+ id="t2",
477
+ stage="data_preparation",
478
+ objective="Pull the order revenue and quantity rows.",
479
+ tool_calls=[
480
+ ToolCall(
481
+ tool="retrieve_data",
482
+ args={
483
+ "ir": {
484
+ "source_id": "src_sales",
485
+ "table_id": "t_orders",
486
+ "select": [
487
+ {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
488
+ {"kind": "column", "column_id": "c_quantity", "alias": "quantity"},
489
+ ],
490
+ "limit": 10000,
491
+ }
492
+ },
493
+ )
494
+ ],
495
+ expected_output="order_rows",
496
+ success_criteria="Produced order-level rows with revenue and quantity.",
497
+ depends_on=["t1"],
498
+ estimated_cost="medium",
499
+ ),
500
+ Task(
501
+ id="t3",
502
+ stage="evaluation",
503
+ objective="Summarize the distribution of revenue and quantity.",
504
+ tool_calls=[
505
+ ToolCall(
506
+ tool="analyze_descriptive",
507
+ args={
508
+ # `data` references t2 (retrieve_data rows), NOT t1 (check_data).
509
+ "data": "${t2}",
510
+ # column refs are the retrieve_data output aliases.
511
+ "column_ids": ["revenue", "quantity"],
512
+ },
513
+ )
514
+ ],
515
+ expected_output="summary_stats",
516
+ success_criteria=(
517
+ "Produced mean/median/std/quartiles for revenue and quantity, above/below "
518
+ "the typical range."
519
+ ),
520
+ depends_on=["t2"],
521
+ estimated_cost="low",
522
+ ),
523
+ ],
524
+ )
525
+
526
+
527
  EXAMPLES: list[tuple[str, TaskList]] = [
528
  ("Which product categories drove last quarter's revenue?", _EXAMPLE_A),
529
  ("How has monthly revenue trended by region this year, and what's unusual?", _EXAMPLE_B),
530
  ("Revenue dipped in Q1 β€” what happened?", _EXAMPLE_C),
531
  ("What is the average and total order value per region?", _EXAMPLE_D),
532
  ("Total revenue for the East and West regions, counting orders of at least 100.", _EXAMPLE_E),
533
+ ("Give me the summary statistics for order revenue and quantity.", _EXAMPLE_F),
534
  ]
535
 
536
 
src/agents/planner/service.py CHANGED
@@ -142,6 +142,25 @@ class PlannerService:
142
  plan_id=task_list.plan_id,
143
  n_tasks=len(task_list.tasks),
144
  retry=attempt > 1,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
145
  )
146
  return task_list
147
 
 
142
  plan_id=task_list.plan_id,
143
  n_tasks=len(task_list.tasks),
144
  retry=attempt > 1,
145
+ # Compact plan dump: task id, its tools, deps, and each tool's arg keys
146
+ # (+ any analyze_* column refs) β€” enough to see how retrieve_data feeds
147
+ # the analyze_* tools without dumping full inline IRs.
148
+ plan=[
149
+ {
150
+ "id": t.id,
151
+ "depends_on": t.depends_on,
152
+ "tools": [
153
+ {
154
+ "tool": c.tool,
155
+ "args": sorted(c.args.keys()),
156
+ "cols": c.args.get("column_ids") or c.args.get("column"),
157
+ "data": c.args.get("data"),
158
+ }
159
+ for c in t.tool_calls
160
+ ],
161
+ }
162
+ for t in task_list.tasks
163
+ ],
164
  )
165
  return task_list
166
 
src/agents/planner/validator.py CHANGED
@@ -29,6 +29,15 @@ logger = get_logger("ir_repair")
29
  # Heuristic: a checkable success_criteria mentions a measurable signal.
30
  _CHECKABLE_TOKENS = ("rate", "count", "match", "produced", "above", "below", "equal")
31
 
 
 
 
 
 
 
 
 
 
32
  # DFS colors for cycle detection.
33
  _WHITE, _GREY, _BLACK = 0, 1, 2
34
 
@@ -109,6 +118,11 @@ class PlannerValidator:
109
  if call.tool == "retrieve_data":
110
  self._validate_inline_ir(task.id, call.args, catalog)
111
 
 
 
 
 
 
112
  # Check 7 β€” success_criteria is checkable.
113
  if not _is_checkable(task.success_criteria):
114
  raise PlannerValidationError(
@@ -157,6 +171,37 @@ class PlannerValidator:
157
  f"task {task_id}: retrieve_data IR failed catalog validation: {e}"
158
  ) from e
159
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
160
  @staticmethod
161
  def _validate_dag(tasks_by_id: dict, id_set: set[str]) -> None:
162
  for task in tasks_by_id.values():
 
29
  # Heuristic: a checkable success_criteria mentions a measurable signal.
30
  _CHECKABLE_TOKENS = ("rate", "count", "match", "produced", "above", "below", "equal")
31
 
32
+ # Tool categories whose output is NOT analyzable data rows: `catalog.introspection`
33
+ # (check_data/check_knowledge β†’ catalog metadata) and `retrieval.documents`
34
+ # (retrieve_knowledge β†’ prose chunks). An analyze_* `data` handoff must not come from
35
+ # one of these β€” their output is kind="table"/documents so the structural checks pass,
36
+ # but feeding metadata rows to analyze_* makes it fail to find the requested columns.
37
+ # A denylist (not an allowlist) so any data/table-producing tool β€” retrieve_data AND a
38
+ # table-producing analyze_* chained into another analyze_* β€” stays valid.
39
+ _NON_DATA_SOURCE_CATEGORIES = frozenset({"catalog.introspection", "retrieval.documents"})
40
+
41
  # DFS colors for cycle detection.
42
  _WHITE, _GREY, _BLACK = 0, 1, 2
43
 
 
118
  if call.tool == "retrieve_data":
119
  self._validate_inline_ir(task.id, call.args, catalog)
120
 
121
+ # Check 9 β€” a `data` handoff (Pattern A) must reference a task that
122
+ # produces analyzable data rows, not one producing catalog metadata
123
+ # (check_data/check_knowledge) or documents (retrieve_knowledge).
124
+ self._validate_data_source(task.id, call, tasks_by_id, registry)
125
+
126
  # Check 7 β€” success_criteria is checkable.
127
  if not _is_checkable(task.success_criteria):
128
  raise PlannerValidationError(
 
171
  f"task {task_id}: retrieve_data IR failed catalog validation: {e}"
172
  ) from e
173
 
174
+ @staticmethod
175
+ def _validate_data_source(
176
+ task_id: str, call, tasks_by_id: dict, registry: ToolRegistry
177
+ ) -> None:
178
+ """A `data` placeholder must reference a data-producing task, not a
179
+ metadata (check_data/check_knowledge) or documents (retrieve_knowledge) one.
180
+
181
+ Those pass the structural checks (check_* also returns kind="table"), but
182
+ their rows are catalog schema, so a downstream analyze_* fails to find the
183
+ requested columns. Resolving points at the referenced task's representative
184
+ output β€” its last tool call (matches TaskRunner's `outputs[-1]`).
185
+ """
186
+ data_arg = call.args.get("data")
187
+ if not isinstance(data_arg, str):
188
+ return
189
+ match = PLACEHOLDER_RE.fullmatch(data_arg.strip())
190
+ if not match:
191
+ return
192
+ ref_task = tasks_by_id.get(match.group(1))
193
+ if ref_task is None or not ref_task.tool_calls:
194
+ return # a dangling placeholder is reported by the DAG check
195
+ ref_tool = ref_task.tool_calls[-1].tool
196
+ ref_spec = registry.get(ref_tool)
197
+ if ref_spec is not None and ref_spec.category in _NON_DATA_SOURCE_CATEGORIES:
198
+ raise PlannerValidationError(
199
+ f"task {task_id}: tool {call.tool!r} takes its 'data' from task "
200
+ f"{match.group(1)} ({ref_tool!r}, category {ref_spec.category!r}), "
201
+ "which produces metadata/documents β€” not analyzable data rows. Feed "
202
+ "analyze_* from a data-producing tool (e.g. retrieve_data)."
203
+ )
204
+
205
  @staticmethod
206
  def _validate_dag(tasks_by_id: dict, id_set: set[str]) -> None:
207
  for task in tasks_by_id.values():
src/config/prompts/planner.md CHANGED
@@ -39,6 +39,12 @@ only a `TaskList` object that conforms to the provided schema.
39
  profiling β€” run `retrieve_data` to fetch the rows, then pass its output to
40
  the matching composite `analyze_*` tool via a `"${t<id>}"` `data` argument
41
  (referencing the upstream result's column aliases).
 
 
 
 
 
 
42
  - **Measure by a dimension in another table (joins).** When the number you are
43
  aggregating and the grouping dimension live in DIFFERENT tables of the same
44
  database source, add a `joins` entry to the `retrieve_data` IR. **Join ONLY on a
 
39
  profiling β€” run `retrieve_data` to fetch the rows, then pass its output to
40
  the matching composite `analyze_*` tool via a `"${t<id>}"` `data` argument
41
  (referencing the upstream result's column aliases).
42
+ An `analyze_*` tool's `data` MUST reference a `retrieve_data` task (or a
43
+ table-producing `analyze_*`) β€” **NEVER `check_data` or `check_knowledge`**.
44
+ Those return catalog **metadata** (a listing of columns/types), not data rows,
45
+ so an `analyze_*` fed from them finds no columns to analyze and fails.
46
+ `check_data` is only for inspecting *what exists*; always `retrieve_data` to
47
+ pull the rows before analyzing them.
48
  - **Measure by a dimension in another table (joins).** When the number you are
49
  aggregating and the grouping dimension live in DIFFERENT tables of the same
50
  database source, add a `joins` entry to the `retrieve_data` IR. **Join ONLY on a