96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
"""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)
|