[KM-567] planner support for analyze_merge: data_right guard + two-retrieve merge few-shot

src/agents/planner/examples.py

@@ -612,6 +612,124 @@ _EXAMPLE_H = TaskList(
 )
 
 
+# --------------------------------------------------------------------------- #
+# 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",
+        ),
+    ],
+)
+
+
 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),
@@ -621,6 +739,11 @@ EXAMPLES: list[tuple[str, TaskList]] = [
     ("Give me the summary statistics for order revenue and quantity.", _EXAMPLE_F),
     ("Which 3 product categories have the best revenue performance?", _EXAMPLE_G),
     ("What is our customer churn rate?", _EXAMPLE_H),
+    (
+        "Which product category has both the highest revenue and the highest average "
+        "order quantity?",
+        _EXAMPLE_I,
+    ),
 ]
 
 

src/agents/planner/validator.py

@@ -220,24 +220,28 @@ class PlannerValidator:
         requested columns. Resolving points at the referenced task's representative
         output — its last tool call (matches TaskRunner's `outputs[-1]`).
         """
-        data_arg = call.args.get("data")
-        if not isinstance(data_arg, str):
-            return
-        match = PLACEHOLDER_RE.fullmatch(data_arg.strip())
-        if not match:
-            return
-        ref_task = tasks_by_id.get(match.group(1))
-        if ref_task is None or not ref_task.tool_calls:
-            return  # a dangling placeholder is reported by the DAG check
-        ref_tool = ref_task.tool_calls[-1].tool
-        ref_spec = registry.get(ref_tool)
-        if ref_spec is not None and ref_spec.category in _NON_DATA_SOURCE_CATEGORIES:
-            raise PlannerValidationError(
-                f"task {task_id}: tool {call.tool!r} takes its 'data' from task "
-                f"{match.group(1)} ({ref_tool!r}, category {ref_spec.category!r}), "
-                "which produces metadata/documents — not analyzable data rows. Feed "
-                "analyze_* from a data-producing tool (e.g. retrieve_data)."
-            )
+        # `data_right` is analyze_merge's second table input (KM-703) — same
+        # Pattern A handoff, so it gets the same guard.
+        for arg_name in ("data", "data_right"):
+            data_arg = call.args.get(arg_name)
+            if not isinstance(data_arg, str):
+                continue
+            match = PLACEHOLDER_RE.fullmatch(data_arg.strip())
+            if not match:
+                continue
+            ref_task = tasks_by_id.get(match.group(1))
+            if ref_task is None or not ref_task.tool_calls:
+                continue  # a dangling placeholder is reported by the DAG check
+            ref_tool = ref_task.tool_calls[-1].tool
+            ref_spec = registry.get(ref_tool)
+            if ref_spec is not None and ref_spec.category in _NON_DATA_SOURCE_CATEGORIES:
+                raise PlannerValidationError(
+                    f"task {task_id}: tool {call.tool!r} takes its {arg_name!r} from "
+                    f"task {match.group(1)} ({ref_tool!r}, category "
+                    f"{ref_spec.category!r}), which produces metadata/documents — not "
+                    "analyzable data rows. Feed analyze_* from a data-producing tool "
+                    "(e.g. retrieve_data)."
+                )
 
     @staticmethod
     def _validate_dag(tasks_by_id: dict, id_set: set[str]) -> None:

src/config/prompts/planner.md

@@ -67,6 +67,14 @@ only a `TaskList` object that conforms to the provided schema.
   `orders.total_amount`); if they genuinely aren't linked, say the data isn't
   connected rather than guessing. Prefer an existing measure column over
   recomputing. Joins are database-only — not available for tabular/file sources.
+- **Two measures per entity ("which X has both the worst A and the biggest B").**
+  Compute each measure in its OWN grouped `retrieve_data` task (one aggregate per
+  entity each), then align them with `analyze_merge`:
+  `{"data": "${tA}", "data_right": "${tB}", "on": ["<entity alias>"]}`.
+  The two retrievals MUST be separate tasks — a `"${t<id>}"` placeholder resolves
+  to a task's LAST output, so two retrievals inside one task lose the first
+  table. The merged table (one row per entity, both measures) answers the
+  question, or feeds a further `analyze_*` step.
 - **Mixing structured + unstructured.** If qualitative context helps, add a
   `retrieve_knowledge` task against an unstructured source listed in the catalog.
 - **CRISP-DM stages.** Tag each task with the stage it serves: