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

api/recon/__init__.py

@@ -0,0 +1,50 @@
+"""Recon package — load the dataset's knowledge graph.
+
+Typical use:
+    from api.recon import load_recon
+    recon = load_recon()
+    tables = recon.owning_tables("A2")          # ["district"]
+    path = recon.join_path("loan", "district")  # ["loan", "account", "district"]
+
+The Recon is read from `api/datasets/<active_dataset>/recon.json`. If the
+file is missing on first call, it's built on demand (with a log line) so
+local dev doesn't require remembering `make recon`. For deploys, run
+`python -m api.recon.build` (or `make recon`) to pre-bake.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from functools import lru_cache
+from pathlib import Path
+
+from api.config import get_settings
+from api.datasets import DATASETS_DIR
+from api.recon.types import (
+    Column, Metric, Recon, Relationship, Table,
+)
+
+logger = logging.getLogger("nvi.recon")
+
+__all__ = [
+    "load_recon", "recon_path",
+    "Recon", "Column", "Table", "Metric", "Relationship",
+]
+
+
+def recon_path(dataset: str | None = None) -> Path:
+    dataset = dataset or get_settings().dataset
+    return DATASETS_DIR / dataset / "recon.json"
+
+
+@lru_cache(maxsize=1)
+def load_recon() -> Recon:
+    dataset = get_settings().dataset
+    path = recon_path(dataset)
+    if not path.exists():
+        logger.info("recon.json missing for %s; building on demand", dataset)
+        from api.recon.build import build_recon, write_recon
+        write_recon(dataset, build_recon(dataset))
+    data = json.loads(path.read_text())
+    return Recon.from_dict(data)

api/recon/build.py

@@ -0,0 +1,192 @@
+"""Recon build — two-stage:
+
+1. **DDL extraction** (auto, from Postgres): reads every table's columns +
+   types + nullability via SQLAlchemy Inspector and writes
+   `api/datasets/<name>/extracted_schema.json`. This is the structural
+   source of truth. Re-run whenever the warehouse changes.
+
+2. **Augmentation merge** (human YAML): reads `schema_docs.yaml` (sparse —
+   table descriptions, column descriptions, relationships) and `metrics.yaml`,
+   merges them onto the extracted schema, and writes
+   `api/datasets/<name>/recon.json` — the artefact the runtime reads.
+
+Run via `make recon` or `python -m api.recon.build`. Auto-triggered by
+load_recon() on first call if recon.json is missing.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import sys
+from pathlib import Path
+from typing import Any
+
+import yaml
+from sqlalchemy import inspect
+
+from api.datasets import DATASETS_DIR
+from api.recon.types import Column, Metric, Recon, Relationship, Table
+from api.tools.db import get_engine
+
+logger = logging.getLogger("nvi.recon.build")
+
+
+# ── Stage 1: DDL extraction ──
+
+def extract_ddl(dataset: str) -> dict[str, Any]:
+    """Introspect Postgres for `<dataset>` schema → structural dict.
+
+    Shape:
+      {
+        "schema": "<name>",
+        "tables": {
+          "<table>": {
+            "columns": [
+              {"name": ..., "sql_type": ..., "nullable": ...},
+              ...
+            ]
+          },
+          ...
+        }
+      }
+    """
+    insp = inspect(get_engine())
+    tables: dict[str, dict[str, Any]] = {}
+    for name in insp.get_table_names(schema=dataset):
+        cols = []
+        for c in insp.get_columns(name, schema=dataset):
+            cols.append({
+                "name": c["name"],
+                "sql_type": str(c["type"]).upper(),
+                "nullable": bool(c["nullable"]),
+            })
+        tables[name] = {"columns": cols}
+    return {"schema": dataset, "tables": tables}
+
+
+def write_extracted_schema(dataset: str, extracted: dict[str, Any]) -> Path:
+    out = DATASETS_DIR / dataset / "extracted_schema.json"
+    out.write_text(json.dumps(extracted, indent=2) + "\n", encoding="utf-8")
+    return out
+
+
+# ── Stage 2: augmentation merge ──
+
+def _read_yaml(path: Path) -> dict[str, Any]:
+    if not path.exists():
+        return {}
+    return yaml.safe_load(path.read_text()) or {}
+
+
+def _parse_column_descs(raw: dict[str, Any]) -> dict[str, str]:
+    """Accept either the flat form `{table.col: desc}` or the nested form
+    `{table: {col: desc}}`. Returns flat form."""
+    out: dict[str, str] = {}
+    for k, v in raw.items():
+        if isinstance(v, dict):
+            for col, desc in v.items():
+                out[f"{k}.{col}"] = desc
+        else:
+            out[k] = v
+    return out
+
+
+def merge_into_recon(dataset: str, extracted: dict[str, Any]) -> Recon:
+    """Combine extracted DDL with human-authored augmentation YAML."""
+    schema_name = extracted.get("schema", dataset)
+
+    aug = _read_yaml(DATASETS_DIR / dataset / "schema_docs.yaml")
+    metrics_yaml = _read_yaml(DATASETS_DIR / dataset / "metrics.yaml")
+
+    table_descs: dict[str, str] = aug.get("tables", {}) or {}
+    col_descs = _parse_column_descs(aug.get("columns", {}) or {})
+
+    tables: dict[str, Table] = {}
+    for tname, t_data in extracted["tables"].items():
+        cols = [
+            Column(
+                name=c["name"],
+                sql_type=c["sql_type"],
+                nullable=c["nullable"],
+                description=col_descs.get(f"{tname}.{c['name']}"),
+            )
+            for c in t_data["columns"]
+        ]
+        tables[tname] = Table(
+            name=tname,
+            description=table_descs.get(tname),
+            columns=cols,
+        )
+
+    metrics: dict[str, Metric] = {
+        name: Metric.from_spec(name, spec)
+        for name, spec in (metrics_yaml.get("metrics", {}) or {}).items()
+    }
+
+    relationships = [Relationship.parse(r) for r in (aug.get("relationships", []) or [])]
+
+    column_to_tables: dict[str, list[str]] = {}
+    for t in tables.values():
+        for c in t.columns:
+            column_to_tables.setdefault(c.name, []).append(t.name)
+    for k in column_to_tables:
+        column_to_tables[k] = sorted(column_to_tables[k])
+
+    return Recon(
+        schema=schema_name,
+        tables=tables,
+        metrics=metrics,
+        column_to_tables=column_to_tables,
+        relationships=relationships,
+    )
+
+
+def write_recon(dataset: str, recon: Recon) -> Path:
+    out = DATASETS_DIR / dataset / "recon.json"
+    out.write_text(json.dumps(recon.to_dict(), indent=2) + "\n", encoding="utf-8")
+    return out
+
+
+# ── Public entry: full build ──
+
+def build_recon(dataset: str) -> Recon:
+    """End-to-end: extract DDL, write the extracted artefact, merge with
+    YAML, write the recon artefact. Returns the in-memory Recon."""
+    extracted = extract_ddl(dataset)
+    write_extracted_schema(dataset, extracted)
+    recon = merge_into_recon(dataset, extracted)
+    write_recon(dataset, recon)
+    return recon
+
+
+def main() -> None:
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
+    ap = argparse.ArgumentParser(description="Build the recon artefacts for one or all datasets.")
+    ap.add_argument("--dataset", help="Dataset name (default: all subdirs of api/datasets/).")
+    args = ap.parse_args()
+
+    if args.dataset:
+        targets = [args.dataset]
+    else:
+        targets = [
+            p.name for p in DATASETS_DIR.iterdir()
+            if p.is_dir() and not p.name.startswith("_") and not p.name.startswith(".")
+        ]
+
+    if not targets:
+        print("no datasets to build", file=sys.stderr)
+        sys.exit(1)
+
+    for name in targets:
+        recon = build_recon(name)
+        print(
+            f"recon[{name}]: {len(recon.tables)} tables, "
+            f"{len(recon.metrics)} metrics, {len(recon.relationships)} relationships "
+            f"→ api/datasets/{name}/{{extracted_schema,recon}}.json"
+        )
+
+
+if __name__ == "__main__":
+    main()

api/recon/types.py

@@ -0,0 +1,250 @@
+"""Recon types — the dataset's knowledge graph as seen by the rest of the api.
+
+The dataset-model types (Column, Table, Metric, SchemaContext) used to live
+in `api/tools/types.py`. They're authored here now because they describe the
+dataset, not the tools that consume it.
+
+`Recon` is a SchemaContext extended with derived indexes that let Analyses
+ask things like "which table owns column A2?" or "how do I join loan to
+district?" without re-deriving from raw YAML each time.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+# ── Schema model ──
+
+@dataclass
+class Column:
+    name: str
+    sql_type: str
+    nullable: bool
+    description: str | None = None
+
+    @classmethod
+    def from_inspector(cls, info: dict[str, Any], description: str | None = None) -> "Column":
+        """Build from a SQLAlchemy Inspector column dict."""
+        return cls(
+            name=info["name"],
+            sql_type=str(info["type"]).upper(),
+            nullable=bool(info["nullable"]),
+            description=description,
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "sql_type": self.sql_type,
+            "nullable": self.nullable,
+            "description": self.description,
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Column":
+        return cls(name=d["name"], sql_type=d["sql_type"],
+                   nullable=d["nullable"], description=d.get("description"))
+
+
+@dataclass
+class Table:
+    name: str
+    description: str | None
+    columns: list[Column]
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "description": self.description,
+            "columns": [c.to_dict() for c in self.columns],
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Table":
+        return cls(
+            name=d["name"],
+            description=d.get("description"),
+            columns=[Column.from_dict(c) for c in d.get("columns", [])],
+        )
+
+
+@dataclass
+class Metric:
+    name: str
+    description: str
+    sql: str
+    from_table: str
+    filter: str | None
+    unit: str | None
+
+    @classmethod
+    def from_spec(cls, name: str, spec: dict[str, Any]) -> "Metric":
+        return cls(
+            name=name,
+            description=spec.get("description", ""),
+            sql=spec["sql"],
+            from_table=spec["from_table"],
+            filter=spec.get("filter"),
+            unit=spec.get("unit"),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "description": self.description,
+            "sql": self.sql,
+            "from_table": self.from_table,
+            "filter": self.filter,
+            "unit": self.unit,
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Metric":
+        return cls(
+            name=d["name"], description=d.get("description", ""),
+            sql=d["sql"], from_table=d["from_table"],
+            filter=d.get("filter"), unit=d.get("unit"),
+        )
+
+
+# ── Relationships ──
+
+@dataclass
+class Relationship:
+    from_table: str
+    from_column: str
+    to_table: str
+    to_column: str
+
+    @classmethod
+    def parse(cls, raw: dict[str, str]) -> "Relationship":
+        """Accept either `{from: 'a.b', to: 'c.d'}` or `{from_table, from_column, ...}`."""
+        if "from_table" in raw:
+            return cls(raw["from_table"], raw["from_column"], raw["to_table"], raw["to_column"])
+        ft, fc = raw["from"].split(".", 1)
+        tt, tc = raw["to"].split(".", 1)
+        return cls(ft, fc, tt, tc)
+
+    def to_dict(self) -> dict[str, str]:
+        return {
+            "from_table": self.from_table, "from_column": self.from_column,
+            "to_table": self.to_table, "to_column": self.to_column,
+        }
+
+
+# ── Recon — top-level dataset knowledge bundle ──
+
+@dataclass
+class Recon:
+    schema: str
+    tables: dict[str, Table]
+    metrics: dict[str, Metric]
+    column_to_tables: dict[str, list[str]] = field(default_factory=dict)
+    relationships: list[Relationship] = field(default_factory=list)
+
+    # ── Read-side helpers used by Analyses + text_to_sql ──
+
+    def table_names(self) -> list[str]:
+        return sorted(self.tables)
+
+    def metric_names(self) -> list[str]:
+        return sorted(self.metrics)
+
+    def owning_tables(self, column: str) -> list[str]:
+        """Tables that have a column with this name (case-sensitive)."""
+        return self.column_to_tables.get(column, [])
+
+    def join_path(self, src: str, dst: str) -> list[str] | None:
+        """Shortest sequence of tables from src to dst via declared relationships.
+        Returns None if no path exists. The path includes both endpoints."""
+        if src == dst:
+            return [src]
+        if src not in self.tables or dst not in self.tables:
+            return None
+        # Build undirected adjacency once.
+        adj: dict[str, set[str]] = {t: set() for t in self.tables}
+        for r in self.relationships:
+            adj.setdefault(r.from_table, set()).add(r.to_table)
+            adj.setdefault(r.to_table, set()).add(r.from_table)
+        # BFS.
+        from collections import deque
+        prev: dict[str, str | None] = {src: None}
+        q: deque[str] = deque([src])
+        while q:
+            cur = q.popleft()
+            if cur == dst:
+                # reconstruct
+                path: list[str] = []
+                node: str | None = cur
+                while node is not None:
+                    path.append(node)
+                    node = prev[node]
+                return list(reversed(path))
+            for nb in adj.get(cur, ()):
+                if nb not in prev:
+                    prev[nb] = cur
+                    q.append(nb)
+        return None
+
+    # ── Prompt-rendering helpers ──
+
+    def render_tables(self, names: list[str] | None = None) -> str:
+        """CREATE TABLE-like rendering for prompt context."""
+        sel = [self.tables[n] for n in (names or self.table_names()) if n in self.tables]
+        out: list[str] = []
+        for t in sel:
+            header = f"-- {t.description}\n" if t.description else ""
+            cols: list[str] = []
+            for c in t.columns:
+                line = f'  "{c.name}" {c.sql_type}'
+                if not c.nullable:
+                    line += " NOT NULL"
+                if c.description:
+                    line += f"  -- {c.description}"
+                cols.append(line)
+            out.append(header + f'CREATE TABLE "{t.name}" (\n' + ",\n".join(cols) + "\n);")
+        return "\n\n".join(out)
+
+    def render_metrics(self) -> str:
+        if not self.metrics:
+            return "(no metrics defined)"
+        lines: list[str] = []
+        for m in self.metrics.values():
+            lines.append(
+                f"- {m.name} ({m.unit or 'unitless'}): {m.description}\n"
+                f"    sql: {m.sql}\n"
+                f"    from: {m.from_table}"
+                + (f"\n    filter: {m.filter}" if m.filter else "")
+            )
+        return "\n".join(lines)
+
+    def render_relationships(self) -> str:
+        if not self.relationships:
+            return "(no relationships declared)"
+        return "\n".join(
+            f"  {r.from_table}.{r.from_column} → {r.to_table}.{r.to_column}"
+            for r in self.relationships
+        )
+
+    # ── Cache serialisation ──
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "schema": self.schema,
+            "tables": {n: t.to_dict() for n, t in self.tables.items()},
+            "metrics": {n: m.to_dict() for n, m in self.metrics.items()},
+            "column_to_tables": self.column_to_tables,
+            "relationships": [r.to_dict() for r in self.relationships],
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Recon":
+        return cls(
+            schema=d["schema"],
+            tables={n: Table.from_dict(t) for n, t in d.get("tables", {}).items()},
+            metrics={n: Metric.from_dict(m) for n, m in d.get("metrics", {}).items()},
+            column_to_tables=d.get("column_to_tables", {}),
+            relationships=[Relationship(**r) for r in d.get("relationships", [])],
+        )

api/recon/validate.py

@@ -0,0 +1,76 @@
+"""SQL validation against the recon schema.
+
+`validate_sql` parses the SQL, runs sqlglot's `qualify` optimizer pass with
+the dataset's column→type schema, and raises with a clear message when the
+LLM references a column that doesn't exist on the table it's bound to.
+
+This is the deterministic complement to the LLM prompt: prompt hints help
+the LLM get it right; the validator GUARANTEES we don't ship a schema-
+wrong query to Postgres. If validation fails, the caller can either bubble
+the error up (no silent retry) or do one *guided* re-prompt that includes
+the validator's message — which is a correction with concrete facts, not
+a blind retry.
+"""
+
+from __future__ import annotations
+
+import sqlglot
+from sqlglot.errors import OptimizeError
+from sqlglot.optimizer.qualify import qualify
+
+from api.recon import load_recon
+from api.recon.types import Recon
+
+
+class ReconValidationError(ValueError):
+    """The SQL references a table/column combination that doesn't exist
+    in the recon. The original sqlglot error is in `__cause__`."""
+
+
+def _build_sqlglot_schema(recon: Recon) -> dict[str, dict[str, str]]:
+    """Convert recon → sqlglot's expected schema shape: {table: {col: type}}."""
+    return {
+        name: {col.name: col.sql_type for col in t.columns}
+        for name, t in recon.tables.items()
+    }
+
+
+def validate_sql(sql: str, recon: Recon | None = None) -> None:
+    """Raise `ReconValidationError` if any column reference in `sql` doesn't
+    exist on the table it's bound to. Returns None on success."""
+    recon = recon or load_recon()
+    schema = _build_sqlglot_schema(recon)
+    try:
+        parsed = sqlglot.parse_one(sql, dialect="postgres")
+        qualify(parsed, schema=schema, dialect="postgres")
+    except OptimizeError as e:
+        # Try to enrich the message with hints from the recon.
+        msg = str(e)
+        hint = _column_hint(msg, recon)
+        full = f"{msg}{hint}" if hint else msg
+        raise ReconValidationError(full) from e
+
+
+_COL_PATTERNS = [
+    r"Column '([^']+)'",                       # "Column 'X' could not be resolved."
+    r"Unknown column:\s*\"?([^\"\s,]+)\"?",  # "Unknown column: X"
+    r"column \"([^\"]+)\"",                   # 'column "X" does not exist'
+]
+
+
+def _column_hint(msg: str, recon: Recon) -> str:
+    """If the error names a specific column, append a hint about which
+    table(s) actually own it. Tries the wordings sqlglot and psycopg use."""
+    import re
+    col_name: str | None = None
+    for pat in _COL_PATTERNS:
+        m = re.search(pat, msg)
+        if m:
+            col_name = m.group(1)
+            break
+    if col_name is None:
+        return ""
+    owners = recon.owning_tables(col_name)
+    if owners:
+        return f"  (Column '{col_name}' actually lives in: {', '.join(owners)}.)"
+    return f"  (No table in the {recon.schema!r} schema has a column named '{col_name}'.)"