recon + sqlglot validator + drill_down package; guard ReAct dimension picks against candidate list

tests/test_analyses_registry.py

@@ -0,0 +1,23 @@
+"""Regression: Analysis registry exposes the expected Analyses."""
+
+from api.analyses.registry import REGISTRY, catalog, get
+
+
+def test_registry_has_expected_analyses():
+    assert set(REGISTRY) >= {"direct_answer", "compare_periods", "drill_down"}
+
+
+def test_registry_get_returns_instance_or_none():
+    a = get("direct_answer")
+    assert a is not None
+    assert a.name == "direct_answer"
+    assert get("does_not_exist") is None
+
+
+def test_catalog_shape_for_planner():
+    cat = catalog()
+    names = {entry["name"] for entry in cat}
+    assert names >= {"direct_answer", "compare_periods", "drill_down"}
+    for entry in cat:
+        assert isinstance(entry["description"], str) and entry["description"]
+        assert isinstance(entry["args_schema"], dict)

tests/test_events.py

@@ -2,6 +2,7 @@
 
 from api.analyses.types import Finding
 from api.runtime import events
+from api.runtime.context import current_analysis
 
 
 def test_run_start_shape():
@@ -22,10 +23,11 @@ def test_plan_ready_shape():
 
 
 def test_analysis_start_shape():
-    e = events.analysis_start("direct_answer", {"q": "?"}, "fits")
+    e = events.analysis_start("direct_answer", {"q": "?"}, "fits", step_id="direct_answer#0")
     assert e == {
         "type": "analysis_start",
         "analysis": "direct_answer",
+        "step_id": "direct_answer#0",
         "args": {"q": "?"},
         "why": "fits",
     }
@@ -33,19 +35,65 @@ def test_analysis_start_shape():
 
 def test_analysis_end_uses_finding():
     f = Finding(analysis="direct_answer", summary="ok", rows=[{"a": 1}], sql=["SELECT 1"])
-    e = events.analysis_end("direct_answer", f)
+    e = events.analysis_end("direct_answer", f, step_id="direct_answer#0")
     assert e["type"] == "analysis_end"
     assert e["analysis"] == "direct_answer"
+    assert e["step_id"] == "direct_answer#0"
     assert e["summary"] == "ok"
     assert e["error"] is None
     assert e["row_count"] == 1
 
 
-def test_tool_call_shapes():
+def test_analysis_fallback_shape():
+    e = events.analysis_fallback(
+        primary="drill_down", fallback="direct_answer",
+        reason="empty result", step_id="direct_answer#0.fallback",
+    )
+    assert e == {
+        "type": "analysis_fallback",
+        "primary": "drill_down",
+        "fallback": "direct_answer",
+        "reason": "empty result",
+        "step_id": "direct_answer#0.fallback",
+    }
+
+
+def test_tool_call_shapes_carry_analysis_field():
+    """Tool/LLM events must include the active Analysis name from the
+    contextvar — UI groups sub-events by it. None when outside an Analysis."""
+    # Outside any Analysis: analysis is None.
     s = events.tool_call_start("text_to_sql", input={"q": "?"})
-    assert s == {"type": "tool_call_start", "tool": "text_to_sql", "input": {"q": "?"}}
+    assert s["type"] == "tool_call_start"
+    assert s["tool"] == "text_to_sql"
+    assert s["analysis"] is None
+    assert s["input"] == {"q": "?"}
+
     e = events.tool_call_end("text_to_sql", output={"sql": "..."})
-    assert e == {"type": "tool_call_end", "tool": "text_to_sql", "output": {"sql": "..."}, "error": None}
+    assert e["analysis"] is None
+    assert e["output"] == {"sql": "..."}
+    assert e["error"] is None
+
+    # Inside an Analysis context: analysis field is populated.
+    token = current_analysis.set("direct_answer#0")
+    try:
+        s2 = events.tool_call_start("execute_sql", input={"sql": "SELECT 1"})
+        e2 = events.tool_call_end("execute_sql", output={"row_count": 1})
+        assert s2["analysis"] == "direct_answer#0"
+        assert e2["analysis"] == "direct_answer#0"
+    finally:
+        current_analysis.reset(token)
+
+
+def test_llm_call_carries_analysis_field():
+    token = current_analysis.set("compare_periods#0")
+    try:
+        e = events.llm_call("compare_periods.interpret", system_len=100, user_len=200)
+        assert e["type"] == "llm_call"
+        assert e["analysis"] == "compare_periods#0"
+        assert e["system_chars"] == 100
+        assert e["user_chars"] == 200
+    finally:
+        current_analysis.reset(token)
 
 
 def test_synth_start_shape():

tests/test_imports.py

@@ -17,8 +17,10 @@ MODULES = [
     "api.llm",
     "api.datasets",
     "api.prompts",
+    "api.recon",
+    "api.recon.types",
+    "api.recon.build",
     "api.tools.db",
-    "api.tools.schema",
     "api.tools.text_to_sql",
     "api.tools.execute_sql",
     "api.tools.types",

tests/test_plan_types.py

@@ -29,3 +29,28 @@ def test_plan_roundtrip():
     assert plan.rationale == "single step suffices"
     assert len(plan.steps) == 1
     assert plan.steps[0].analysis == "direct_answer"
+
+
+def test_planstep_fallback_roundtrip():
+    raw = {
+        "analysis": "drill_down",
+        "args": {"metric": "amount"},
+        "why": "speculative",
+        "fallback": {
+            "analysis": "direct_answer",
+            "args": {"question": "fallback q"},
+            "why": "safe option",
+        },
+    }
+    step = PlanStep.from_dict(raw)
+    assert step.fallback is not None
+    assert step.fallback.analysis == "direct_answer"
+    assert step.fallback.args == {"question": "fallback q"}
+    assert step.fallback.why == "safe option"
+    assert step.to_dict() == raw
+
+
+def test_planstep_no_fallback_omits_key():
+    step = PlanStep.from_dict({"analysis": "direct_answer"})
+    assert step.fallback is None
+    assert "fallback" not in step.to_dict()

tests/test_recon.py

@@ -0,0 +1,106 @@
+"""Regression tests for the recon package — type round-trips, join_path,
+and column_to_tables index."""
+
+from api.recon.types import Recon, Relationship, Table, Column, Metric
+
+
+def _sample_recon() -> Recon:
+    cols_account = [
+        Column("account_id", "BIGINT", False),
+        Column("district_id", "BIGINT", True),
+    ]
+    cols_district = [
+        Column("district_id", "BIGINT", False),
+        Column("A2", "TEXT", True),
+        Column("A11", "BIGINT", True),
+    ]
+    cols_loan = [
+        Column("loan_id", "BIGINT", False),
+        Column("account_id", "BIGINT", True),
+        Column("amount", "BIGINT", True),
+        Column("status", "TEXT", True),
+    ]
+    rels = [
+        Relationship("account", "district_id", "district", "district_id"),
+        Relationship("loan", "account_id", "account", "account_id"),
+    ]
+    c2t = {
+        "account_id": sorted(["account", "loan"]),
+        "district_id": sorted(["account", "district"]),
+        "A2": ["district"],
+        "A11": ["district"],
+        "loan_id": ["loan"],
+        "amount": ["loan"],
+        "status": ["loan"],
+    }
+    return Recon(
+        schema="financial",
+        tables={
+            "account": Table("account", "Accounts", cols_account),
+            "district": Table("district", "Districts", cols_district),
+            "loan": Table("loan", "Loans", cols_loan),
+        },
+        metrics={
+            "loan_volume": Metric("loan_volume", "Total CZK lent",
+                                  "SUM(amount)", "loan", None, "CZK"),
+        },
+        column_to_tables=c2t,
+        relationships=rels,
+    )
+
+
+def test_owning_tables():
+    r = _sample_recon()
+    assert r.owning_tables("A2") == ["district"]
+    assert r.owning_tables("amount") == ["loan"]
+    assert r.owning_tables("account_id") == ["account", "loan"]
+    assert r.owning_tables("nonsense") == []
+
+
+def test_join_path_direct():
+    r = _sample_recon()
+    assert r.join_path("loan", "account") == ["loan", "account"]
+    assert r.join_path("account", "district") == ["account", "district"]
+
+
+def test_join_path_multi_hop():
+    r = _sample_recon()
+    assert r.join_path("loan", "district") == ["loan", "account", "district"]
+
+
+def test_join_path_same():
+    r = _sample_recon()
+    assert r.join_path("loan", "loan") == ["loan"]
+
+
+def test_join_path_unknown_returns_none():
+    r = _sample_recon()
+    assert r.join_path("loan", "does_not_exist") is None
+
+
+def test_recon_to_dict_from_dict_roundtrip():
+    r = _sample_recon()
+    d = r.to_dict()
+    r2 = Recon.from_dict(d)
+    assert r2.schema == r.schema
+    assert sorted(r2.tables) == sorted(r.tables)
+    assert sorted(r2.metrics) == sorted(r.metrics)
+    assert r2.column_to_tables == r.column_to_tables
+    assert len(r2.relationships) == len(r.relationships)
+    assert r2.join_path("loan", "district") == ["loan", "account", "district"]
+
+
+def test_relationship_parse_dotted():
+    r = Relationship.parse({"from": "loan.account_id", "to": "account.account_id"})
+    assert r.from_table == "loan"
+    assert r.from_column == "account_id"
+    assert r.to_table == "account"
+    assert r.to_column == "account_id"
+
+
+def test_relationship_parse_explicit():
+    r = Relationship.parse({
+        "from_table": "loan", "from_column": "account_id",
+        "to_table": "account", "to_column": "account_id",
+    })
+    assert r.from_table == "loan"

tests/test_recon_validate.py

@@ -0,0 +1,94 @@
+"""Regression: recon SQL validator catches LLM-generated schema mismatches.
+
+Bug captured: drill_down generated `SELECT district_id FROM loan` —
+syntactically valid, parses, but `loan` doesn't have a `district_id` column.
+We want to catch this BEFORE Postgres does, with a clear message naming
+where the column actually lives.
+"""
+
+import pytest
+
+from api.recon.types import Column, Metric, Recon, Relationship, Table
+from api.recon.validate import ReconValidationError, validate_sql
+
+
+def _recon() -> Recon:
+    return Recon(
+        schema="financial",
+        tables={
+            "loan": Table("loan", None, [
+                Column("loan_id", "BIGINT", False),
+                Column("account_id", "BIGINT", True),
+                Column("amount", "BIGINT", True),
+                Column("status", "TEXT", True),
+            ]),
+            "account": Table("account", None, [
+                Column("account_id", "BIGINT", False),
+                Column("district_id", "BIGINT", True),
+            ]),
+            "district": Table("district", None, [
+                Column("district_id", "BIGINT", False),
+                Column("A2", "TEXT", True),
+                Column("A13", "BIGINT", True),
+            ]),
+        },
+        metrics={},
+        column_to_tables={
+            "loan_id": ["loan"],
+            "account_id": ["account", "loan"],
+            "amount": ["loan"],
+            "status": ["loan"],
+            "district_id": ["account", "district"],
+            "A2": ["district"],
+            "A13": ["district"],
+        },
+        relationships=[
+            Relationship("account", "district_id", "district", "district_id"),
+            Relationship("loan", "account_id", "account", "account_id"),
+        ],
+    )
+
+
+def test_valid_sql_passes():
+    sql = "SELECT account_id, amount FROM loan WHERE status = 'B'"
+    validate_sql(sql, _recon())
+
+
+def test_valid_join_passes():
+    sql = """
+      SELECT d.A2, AVG(l.amount)
+      FROM loan l
+      JOIN account a ON l.account_id = a.account_id
+      JOIN district d ON a.district_id = d.district_id
+      GROUP BY d.A2
+    """
+    validate_sql(sql, _recon())
+
+
+def test_unknown_column_on_table_raises():
+    # The exact bug we hit in production: district_id picked from loan
+    # without joining through account.
+    sql = """
+      WITH x AS (SELECT district_id, AVG(amount) FROM loan GROUP BY district_id)
+      SELECT * FROM x
+    """
+    with pytest.raises(ReconValidationError) as exc:
+        validate_sql(sql, _recon())
+    msg = str(exc.value)
+    assert "district_id" in msg
+    # The hint should name where district_id actually lives.
+    assert "account" in msg and "district" in msg
+
+
+def test_misspelled_column_raises():
+    sql = "SELECT amnt FROM loan"
+    with pytest.raises(ReconValidationError) as exc:
+        validate_sql(sql, _recon())
+    assert "amnt" in str(exc.value)
+
+
+def test_qualified_column_on_wrong_table_raises():
+    # district.amount doesn't exist (amount is on loan).
+    sql = "SELECT d.amount FROM district d"
+    with pytest.raises(ReconValidationError):
+        validate_sql(sql, _recon())

tests/test_yaml_datasets.py

@@ -3,6 +3,11 @@
 Bug captured: schema_docs.yaml had values that began with a quote character
 (e.g. `gender: 'M' or 'F'.`), which YAML treats as a quoted scalar — anything
 after the closing quote is a parse error.
+
+The YAML structure is now flat — tables.<name> is a description string,
+columns are keyed `table.column`, and relationships are an explicit list.
+The DDL structure (column types, nullability) lives in extracted_schema.json
+generated by `make recon`.
 """
 
 from pathlib import Path
@@ -20,6 +25,9 @@ def test_schema_docs_parses(dataset):
     assert docs is not None
     assert "tables" in docs
     assert isinstance(docs["tables"], dict) and docs["tables"]
+    # New shape: each table value is a description string, not a dict.
+    for name, val in docs["tables"].items():
+        assert isinstance(val, str), f"tables.{name} should be a string description, got {type(val)}"
 
 
 @pytest.mark.parametrize("dataset", DATASETS)
@@ -31,16 +39,24 @@ def test_metrics_parses(dataset):
 
 
 def test_financial_columns_with_quotes_survive_parse():
-    """Specific regression: values that contain quoted codes (M, F, OWNER,
-    PRIJEM, etc.) must parse without truncation."""
+    """Values containing quoted codes (M, F, OWNER, PRIJEM, ...) must parse
+    without truncation. Columns are flat `table.col` keys now."""
     docs = yaml.safe_load(
         (DATASETS_DIR / "financial" / "schema_docs.yaml").read_text()
     )
-    cols = docs["tables"]["client"]["columns"]
-    assert '"M"' in cols["gender"] and '"F"' in cols["gender"]
+    cols = docs["columns"]
+    assert '"M"' in cols["client.gender"] and '"F"' in cols["client.gender"]
+    assert '"OWNER"' in cols["disp.type"] and '"DISPONENT"' in cols["disp.type"]
+    assert '"PRIJEM"' in cols["trans.type"] and '"VYDAJ"' in cols["trans.type"]
 
-    disp = docs["tables"]["disp"]["columns"]
-    assert '"OWNER"' in disp["type"] and '"DISPONENT"' in disp["type"]
 
-    trans = docs["tables"]["trans"]["columns"]
-    assert '"PRIJEM"' in trans["type"] and '"VYDAJ"' in trans["type"]
+def test_financial_has_relationships():
+    docs = yaml.safe_load(
+        (DATASETS_DIR / "financial" / "schema_docs.yaml").read_text()
+    )
+    rels = docs.get("relationships") or []
+    assert len(rels) >= 5  # we have 8; assert ≥5 to allow some pruning
+    # all entries have from/to
+    for r in rels:
+        assert "from" in r and "to" in r
+        assert "." in r["from"] and "." in r["to"]