R3 polish: drop 'matched to your data shape' overclaim on no-first_treat path

Codex R2 narrowed P1: the no-`first_treat` branch unconditionally emitted
`DifferenceInDifferences().fit(...)` and labeled it "matched to your data
shape". That overclaim doesn't hold for continuous-dose or heterogeneous-
adoption designs without first_treat — DiD validates and rejects non-
binary treatment / time at fit-time (`diff_diff/estimators.py:307-312`).

Fix in agent_workflow.py:
- With `first_treat`: keep CallawaySantAnna example, relabel as
  "your data has first_treat -> staggered structure; CS is canonical".
- Without `first_treat`: keep DiD example as the simple 2x2 case but
  reframe the label to explicitly condition on that shape, name the
  alternative candidates (ContinuousDiD, HeterogeneousAdoptionDiD),
  and reference DiD's fit-time validation so the agent knows when to
  switch. No claim of universal match.

New regression test `test_no_first_treat_step3_does_not_overclaim_match`:
- asserts "matched to your data shape" is NOT in the no-first_treat
  script (negative)
- asserts "2x2" and "substitute" appear (positive — substitution hint)
- asserts ContinuousDiD and HeterogeneousAdoptionDiD remain enumerated
  in Step 2's routing patterns

Total tests: 17 → 18 in test_agent_workflow.py.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: igerber

diff_diff/agent_workflow.py

@@ -145,23 +145,43 @@ def agent_workflow(
     )
     guide_call = 'diff_diff.get_llm_guide("autonomous")'
 
-    # Step 3 example: branch on first_treat presence. Staggered estimators
-    # (CallawaySantAnna et al.) take `first_treat=` and do NOT accept
-    # `treatment=` on fit(); simple 2x2 `DifferenceInDifferences.fit()`
-    # takes `treatment=` and does not accept `first_treat=`. Emitting the
-    # wrong shape produces a TypeError when the agent copy-pastes.
+    # Step 3 example: branch on first_treat presence.
+    # - With first_treat: a staggered structure is strongly implied;
+    #   show CallawaySantAnna with the matching `first_treat=` signature
+    #   (NB: staggered estimators reject `treatment=` on fit()).
+    # - Without first_treat: the orchestrator does not inspect df, so it
+    #   CANNOT infer whether the panel is 2x2 binary vs continuous-dose
+    #   vs heterogeneous-adoption. Show a DifferenceInDifferences call
+    #   as the "simple 2x2" example and label it explicitly conditional
+    #   on that shape. The agent is expected to substitute the right
+    #   estimator class from Step 2's routing patterns when the data
+    #   is continuous / heterogeneous / non-binary-time
+    #   (DifferenceInDifferences.fit() rejects non-binary treatment/
+    #   time per diff_diff/estimators.py).
     if first_treat is not None:
         fit_example_class = "CallawaySantAnna"
         fit_example_kwargs = _join_kwargs(
             outcome=outcome, unit=unit, time=time, first_treat=first_treat
         )
         fit_example_call = f"diff_diff.CallawaySantAnna().fit(df, {fit_example_kwargs})"
+        step3_label = (
+            f"Step 3 - Fit. Your data has `first_treat` -> staggered structure; "
+            f"{fit_example_class} is the canonical starting point:"
+        )
     else:
         fit_example_class = "DifferenceInDifferences"
         fit_example_kwargs = _join_kwargs(
             outcome=outcome, unit=unit, time=time, treatment=treatment
         )
         fit_example_call = f"diff_diff.DifferenceInDifferences().fit(df, {fit_example_kwargs})"
+        step3_label = (
+            "Step 3 - Fit. Pick a candidate from Step 2's patterns based on your "
+            "treatment/time shape. The example below shows the simple 2x2 case "
+            "(binary treatment + binary time); substitute ContinuousDiD / "
+            "HeterogeneousAdoptionDiD / etc. when your design is not 2x2 "
+            "(DifferenceInDifferences.fit() validates and rejects non-binary "
+            "treatment or time):"
+        )
 
     validation_calls = [
         "diff_diff.practitioner_next_steps(result)",
@@ -200,7 +220,7 @@ def agent_workflow(
 
 {diagnostics_block}
 
-Step 3 - Fit (matched to your data shape; {fit_example_class} shown):
+{step3_label}
     result = {fit_example_call}
 
 Step 4 - Validate:

tests/test_agent_workflow.py

@@ -142,6 +142,35 @@ def test_first_treat_switches_step3_estimator(df):
     assert "treatment=" not in step3_lines[0]
 
 
+def test_no_first_treat_step3_does_not_overclaim_match(df):
+    """Without `first_treat`, the orchestrator cannot infer panel shape.
+
+    The Step 3 label must NOT claim the emitted DiD example is "matched to
+    your data shape" — for continuous-dose or heterogeneous-adoption
+    designs without first_treat, DifferenceInDifferences would reject at
+    fit time. The label must instead frame the example as conditional on
+    the simple 2x2 case and tell the agent to substitute the matching
+    candidate from Step 2 for other shapes.
+    """
+    out = diff_diff.agent_workflow(
+        df,
+        unit="firm_id",
+        time="year",
+        treatment="treated",
+        outcome="logwage",
+        verbose=False,
+    )
+    script = out["script"]
+    # Negative: should not claim universal match.
+    assert "matched to your data shape" not in script
+    # Positive: must qualify with the 2x2 conditionality + substitution hint.
+    assert "2x2" in script
+    assert "substitute" in script.lower()
+    # Other workflow patterns must remain enumerated so the agent can substitute.
+    for name in ("ContinuousDiD", "HeterogeneousAdoptionDiD"):
+        assert name in script, f"{name} routing pattern missing from Step 2 hints"
+
+
 def test_does_not_inspect_df():
     # Pure orchestrator: a structurally-empty DataFrame must still produce
     # the templated script (no df inspection happens).