recon-driven sql composer; pick → compose → execute; llm out of structural sql

api/analyses/_narrow.py

@@ -0,0 +1,95 @@
+"""Candidate narrowing for the Pick LLM call.
+
+The composer covers L1's schema-pollution problem — the LLM never sees raw
+SQL. But the LLM still needs to choose from *some* candidate set when
+picking a metric / dimension / filter column. Sending the full schema for
+every question is wasteful (cost, latency) and noisier (more candidates
+the model can pick wrong from).
+
+v1: **structural neighborhood**. Given the candidate metrics, compute the
+union of:
+  - columns of each candidate metric's `from_table`
+  - columns of every table reachable within `max_hops` via declared
+    relationships.
+
+This covers every column that could legitimately group/filter that metric
+without hand-curating a list. Cheap, deterministic, easy to reason about.
+
+v2 (future): embedding-based retrieval against question + metric/column
+descriptions, with usage-derived signal (pgvector). Out of scope here.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from api.recon.types import Metric, Recon
+
+
+@dataclass
+class Candidates:
+    metrics: list[Metric]
+    # column_refs: each entry is either a bare name (when unambiguous) or
+    # `table.column` (when the name lives on more than one reachable table).
+    column_refs: list[str]
+
+
+def neighborhood_tables(recon: Recon, seeds: list[str], max_hops: int = 2) -> set[str]:
+    """BFS from `seeds` over recon.relationships, returning every table
+    reachable within `max_hops`. Edges are undirected (FKs go both ways)."""
+    adj: dict[str, set[str]] = {t: set() for t in recon.tables}
+    for r in recon.relationships:
+        adj.setdefault(r.from_table, set()).add(r.to_table)
+        adj.setdefault(r.to_table, set()).add(r.from_table)
+    visited: set[str] = set()
+    frontier: set[str] = {s for s in seeds if s in recon.tables}
+    for _ in range(max_hops + 1):  # +1 to include the seed itself at hop 0
+        visited |= frontier
+        next_frontier: set[str] = set()
+        for t in frontier:
+            next_frontier |= adj.get(t, set())
+        frontier = next_frontier - visited
+        if not frontier:
+            break
+    return visited
+
+
+def narrow_candidates(
+    recon: Recon,
+    *,
+    allowed_metrics: list[str] | None = None,
+    max_hops: int = 2,
+) -> Candidates:
+    """Return the Pick-candidate set for the given metrics.
+
+    `allowed_metrics`: if provided, restricts the candidate metrics to this
+    list (intersected with recon.metrics). If None, every metric is in scope.
+    `max_hops`: how far to walk from each candidate metric's from_table.
+    """
+    metric_names = (
+        [m for m in allowed_metrics if m in recon.metrics]
+        if allowed_metrics is not None
+        else list(recon.metrics)
+    )
+    metrics = [recon.metrics[m] for m in metric_names]
+    seeds = [m.from_table for m in metrics]
+    tables = neighborhood_tables(recon, seeds, max_hops=max_hops)
+
+    # Build the column ref list. A column name appearing in more than one
+    # reachable table is exposed as `table.column` for each owner to keep
+    # the LLM's pick unambiguous.
+    appearances: dict[str, list[str]] = {}
+    for tname in tables:
+        for c in recon.tables[tname].columns:
+            appearances.setdefault(c.name, []).append(tname)
+
+    refs: list[str] = []
+    for col, owners in appearances.items():
+        if len(owners) == 1:
+            refs.append(col)
+        else:
+            for t in owners:
+                refs.append(f"{t}.{col}")
+    refs.sort()
+
+    return Candidates(metrics=metrics, column_refs=refs)