Annie Voigt commited on
Commit
7944fee
·
1 Parent(s): 55b81a3

fix(planner): dataset_plan_analysis must say how to LOAD the dataset

Browse files

Verifying the previous commit on dev showed rule 7 worked — both runs now open
with dataset_plan_analysis — but the plan they got back had no loading step, so
the agent still had to guess an entry point and still burned steps on it.

_workflow_for_intent emitted a load step only for expression_source.type ==
"geo_series_matrix". Most of the registry is an h5ad over `url`, so those
datasets returned a plan starting at dataset_describe and jumping straight to
collapse/DE, with `required_inputs: ["Expression file ..."]` as if the user
supplied it. Loading now comes from _build_loading_plan — already the single
source of truth for how to open a dataset, and it knows the modality-first Path
P routing and the precomputed `collapsed_url`.

Two details that matter:
- Only the get-the-data-open prefix is spliced in. _build_loading_plan also
carries an analysis tail (validate_contrast -> DE -> scoring) which the
intent-specific steps already own; splicing both duplicated it (caught by two
existing planner tests).
- The plan's own decoupler_inspect_data step is dropped for a registered
dataset. That step is why the agent kept calling it despite efficiency rule 4
forbidding it — the plan was telling it to.

Measured on gse28735_pdac: the plan now names decoupler_load_url_counts with the
_collapsed.h5ad URL, so the annotate+collapse pair disappears entirely.

Suite 1325 passed / 60 skipped.

src/tools/dataset_tools/_base.py CHANGED
@@ -594,6 +594,23 @@ def _detect_intent(question: str) -> tuple[str, str, list[str]]:
594
  return best_intent, confidence, matched
595
 
596
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
597
  def _workflow_for_intent(
598
  intent: str,
599
  dataset_id: str,
@@ -632,26 +649,50 @@ def _workflow_for_intent(
632
  ]
633
  step_n += 1
634
 
635
- # Optional loading steps from manifest
 
 
 
 
 
 
 
636
  if manifest is not None and hasattr(manifest, "expression_source"):
637
- if manifest.expression_source.get("type") == "geo_series_matrix":
638
- steps.append(
639
- {
640
- "step": step_n,
641
- "tool": "decoupler_load_geo_series_matrix",
642
- "purpose": "Download and parse GEO series matrix to h5ad.",
643
- "status": "available",
644
- "args_hint": {
645
- "url_or_path": manifest.expression_source.get("url"),
646
- "condition_column": (
647
- manifest.group_columns[0] if manifest.group_columns else None
648
- ),
649
- },
650
- }
651
- )
652
- step_n += 1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
653
 
654
- if collapse_required:
 
 
655
  steps.append(
656
  {
657
  "step": step_n,
 
594
  return best_intent, confidence, matched
595
 
596
 
597
+ # The prefix of _build_loading_plan that is about OPENING the dataset. Anything
598
+ # else it emits (validate_contrast, DE, scoring) is the analysis tail, which the
599
+ # intent-specific steps in _workflow_for_intent own instead.
600
+ _LOADING_PLAN_TOOLS = frozenset(
601
+ {
602
+ "decoupler_load_url_counts",
603
+ "decoupler_load_geo_series_matrix",
604
+ "decoupler_load_gdc_star_counts",
605
+ "decoupler_load_and_visualize_data",
606
+ "decoupler_annotate_probes_with_gpl",
607
+ "decoupler_collapse_probes_to_genes",
608
+ "decoupler_join_clinical_metadata",
609
+ "dataset_filter_to_curated_samples",
610
+ }
611
+ )
612
+
613
+
614
  def _workflow_for_intent(
615
  intent: str,
616
  dataset_id: str,
 
649
  ]
650
  step_n += 1
651
 
652
+ # Loading steps come from _build_loading_plan — the single source of truth
653
+ # for HOW to open a dataset. This used to emit a load step only for
654
+ # `geo_series_matrix`, so every h5ad-over-`url` dataset (most of the
655
+ # registry) got a plan with NO loading step at all and the agent had to
656
+ # guess an entry point, which cost 3-5 steps per run and often failed
657
+ # outright. The loading plan also knows about `collapsed_url`, which skips
658
+ # the annotate+collapse pair below.
659
+ loading_steps: list[dict] = []
660
  if manifest is not None and hasattr(manifest, "expression_source"):
661
+ try:
662
+ loading_steps = _build_loading_plan(manifest)
663
+ except Exception: # a malformed manifest must not sink the whole plan
664
+ loading_steps = []
665
+
666
+ collapse_precomputed = False
667
+ for raw in loading_steps:
668
+ tool = raw.get("tool")
669
+ # Take only the GET-THE-DATA-OPEN prefix. _build_loading_plan also
670
+ # carries an analysis tail (validate_contrast → DE → scoring); the
671
+ # intent-specific steps below own that, and splicing both in would
672
+ # duplicate it.
673
+ if tool not in _LOADING_PLAN_TOOLS:
674
+ continue
675
+ # The loading plan's own inspect step is redundant for a REGISTERED
676
+ # dataset — data_level and analysis_path are manifest facts (efficiency
677
+ # rule 4). Keeping it here is what kept the agent calling it.
678
+ if tool == "decoupler_inspect_data":
679
+ continue
680
+ if tool == "decoupler_collapse_probes_to_genes":
681
+ collapse_precomputed = True
682
+ steps.append(
683
+ {
684
+ "step": step_n,
685
+ "tool": tool,
686
+ "purpose": raw.get("note", "Load the dataset as the manifest specifies."),
687
+ "status": "available",
688
+ "args_hint": raw.get("key_args", {}),
689
+ }
690
+ )
691
+ step_n += 1
692
 
693
+ # Only add a collapse step if the loading plan did not already cover it
694
+ # (a precomputed `collapsed_url` makes it unnecessary).
695
+ if collapse_required and not collapse_precomputed and not loading_steps:
696
  steps.append(
697
  {
698
  "step": step_n,
tests/test_planner_tool.py CHANGED
@@ -327,3 +327,66 @@ class TestManifestDrivenArgs:
327
  result = dataset_plan_analysis("nonexistent_dataset_xyz", "compare groups")
328
  assert result["detected_intent"] == "compare_groups"
329
  assert len(result["recommended_tools"]) > 0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
327
  result = dataset_plan_analysis("nonexistent_dataset_xyz", "compare groups")
328
  assert result["detected_intent"] == "compare_groups"
329
  assert len(result["recommended_tools"]) > 0
330
+
331
+
332
+ class TestPlanIncludesLoading:
333
+ """The plan must say HOW to open the dataset, not just what to run.
334
+
335
+ Measured on dev 2026-08-04: `_workflow_for_intent` emitted a load step only
336
+ for `geo_series_matrix`, so every h5ad-over-`url` dataset (most of the
337
+ registry) returned a plan with no loading step. The agent then guessed an
338
+ entry point — 3-5 wasted steps per run, and for the DE query it never found
339
+ a working one. Loading now comes from `_build_loading_plan`.
340
+ """
341
+
342
+ def test_h5ad_url_dataset_gets_a_load_step(self):
343
+ plan = dataset_plan_analysis(
344
+ dataset_id="gse28735_pdac",
345
+ user_question="Compare tumor vs normal for differential expression",
346
+ )
347
+ tools = [s["tool"] for s in plan["recommended_tools"]]
348
+ assert "decoupler_load_url_counts" in tools, tools
349
+ # ...and it must come before the analysis it feeds.
350
+ assert tools.index("decoupler_load_url_counts") < tools.index(
351
+ "decoupler_differential_expression"
352
+ )
353
+
354
+ def test_load_step_uses_the_precollapsed_url(self):
355
+ """The uncollapsed URL costs two extra steps (annotate + collapse)."""
356
+ plan = dataset_plan_analysis(
357
+ dataset_id="gse28735_pdac", user_question="differential expression tumor vs normal"
358
+ )
359
+ load = next(
360
+ s for s in plan["recommended_tools"] if s["tool"] == "decoupler_load_url_counts"
361
+ )
362
+ assert "collapsed" in load["args_hint"]["url_or_path"]
363
+ tools = [s["tool"] for s in plan["recommended_tools"]]
364
+ assert "decoupler_annotate_probes_with_gpl" not in tools
365
+ assert "decoupler_collapse_probes_to_genes" not in tools
366
+
367
+ def test_registered_dataset_plan_omits_inspect_data(self):
368
+ """Efficiency rule 4: data_level/analysis_path are manifest facts."""
369
+ for did in ("gse28735_pdac", "paca_au_rnaseq"):
370
+ plan = dataset_plan_analysis(dataset_id=did, user_question="compare groups")
371
+ tools = [s["tool"] for s in plan["recommended_tools"]]
372
+ assert "decoupler_inspect_data" not in tools, (did, tools)
373
+
374
+ def test_single_cell_plan_uses_the_sc_loader(self):
375
+ """Path P must not be handed a bulk flat-file loader (ADR-0006)."""
376
+ plan = dataset_plan_analysis(
377
+ dataset_id="gse155698_steele", user_question="compare tumor vs normal"
378
+ )
379
+ tools = [s["tool"] for s in plan["recommended_tools"]]
380
+ assert "decoupler_load_and_visualize_data" in tools, tools
381
+ assert "decoupler_load_url_counts" not in tools
382
+
383
+ def test_steps_are_numbered_consecutively(self):
384
+ plan = dataset_plan_analysis(dataset_id="gse28735_pdac", user_question="compare groups")
385
+ nums = [s["step"] for s in plan["recommended_tools"]]
386
+ assert nums == list(range(1, len(nums) + 1)), nums
387
+
388
+ def test_unknown_dataset_still_returns_a_plan(self):
389
+ """A missing manifest must degrade to a warning, not an exception."""
390
+ plan = dataset_plan_analysis(dataset_id="not_a_dataset", user_question="compare groups")
391
+ assert plan["recommended_tools"]
392
+ assert plan["warnings"]