recon + sqlglot validator + drill_down package; guard ReAct dimension picks against candidate list
This commit is contained in:
@@ -1,57 +0,0 @@
|
||||
"""Warehouse schema retrieval.
|
||||
|
||||
Physical layer via SQLAlchemy's Inspector against the active dataset's
|
||||
Postgres schema; semantic layer (table/column descriptions, metric catalog)
|
||||
from YAML files under `api/datasets/<name>/`.
|
||||
|
||||
Cached on first call — the schema is stable for a process's lifetime.
|
||||
Restart the api to pick up a dataset switch or YAML edits.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import inspect
|
||||
|
||||
from api.config import get_settings
|
||||
from api.datasets import read_yaml
|
||||
from api.tools.db import get_engine
|
||||
from api.tools.types import Column, Metric, SchemaContext, Table
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def load_schema() -> SchemaContext:
|
||||
dataset = get_settings().dataset
|
||||
docs = read_yaml(dataset, "schema_docs.yaml")
|
||||
metrics_yaml = read_yaml(dataset, "metrics.yaml")
|
||||
|
||||
# Postgres schema name comes from the dataset id; the schema_docs file
|
||||
# can override it, but normally they match.
|
||||
schema_name: str = docs.get("schema", dataset)
|
||||
table_descs: dict[str, dict[str, Any]] = docs.get("tables", {}) or {}
|
||||
|
||||
insp = inspect(get_engine())
|
||||
tables: dict[str, Table] = {
|
||||
name: Table(
|
||||
name=name,
|
||||
description=(desc := table_descs.get(name, {})).get("description"),
|
||||
columns=[
|
||||
Column.from_inspector(c, (desc.get("columns") or {}).get(c["name"]))
|
||||
for c in insp.get_columns(name, schema=schema_name)
|
||||
],
|
||||
)
|
||||
for name in insp.get_table_names(schema=schema_name)
|
||||
}
|
||||
|
||||
metrics: dict[str, Metric] = {
|
||||
name: Metric.from_spec(name, spec)
|
||||
for name, spec in (metrics_yaml.get("metrics", {}) or {}).items()
|
||||
}
|
||||
|
||||
return SchemaContext(schema=schema_name, tables=tables, metrics=metrics)
|
||||
|
||||
|
||||
def get_metric(name: str) -> Metric | None:
|
||||
return load_schema().metrics.get(name)
|
||||
@@ -10,23 +10,24 @@ import sqlglot
|
||||
from api import langfuse_client as lf
|
||||
from api.llm import chat
|
||||
from api.prompts import load, render
|
||||
from api.tools.schema import load_schema
|
||||
from api.recon import load_recon
|
||||
from api.recon.validate import ReconValidationError, validate_sql
|
||||
from api.tools.types import T2SResult
|
||||
|
||||
logger = logging.getLogger("nvi.tools.text_to_sql")
|
||||
|
||||
|
||||
def text_to_sql(question: str, *, hint_tables: list[str] | None = None) -> T2SResult:
|
||||
schema = load_schema()
|
||||
tables = hint_tables or schema.table_names()
|
||||
recon = load_recon()
|
||||
tables = hint_tables or recon.table_names()
|
||||
system = load("text_to_sql.system")
|
||||
|
||||
def _user(retry_hint: str = "") -> str:
|
||||
return render(
|
||||
"text_to_sql.user",
|
||||
question=question,
|
||||
schema_block=schema.render_tables(tables),
|
||||
metrics_block=schema.render_metrics(),
|
||||
schema_block=recon.render_tables(tables),
|
||||
metrics_block=recon.render_metrics(),
|
||||
retry_hint=retry_hint,
|
||||
)
|
||||
|
||||
@@ -35,14 +36,9 @@ def text_to_sql(question: str, *, hint_tables: list[str] | None = None) -> T2SRe
|
||||
as_type="generation",
|
||||
input={"question": question, "hint_tables": hint_tables},
|
||||
) as span:
|
||||
sql = _extract_sql(chat(system=system, user=_user()))
|
||||
try:
|
||||
_validate(sql)
|
||||
except Exception as e:
|
||||
logger.info("t2s first attempt invalid (%s); retrying once", e)
|
||||
hint = f"\n\nPrevious attempt failed parsing: {e}. Fix and return only the SQL."
|
||||
sql = _extract_sql(chat(system=system, user=_user(hint)))
|
||||
_validate(sql)
|
||||
raw = _extract_sql(chat(system=system, user=_user()))
|
||||
sql = _normalize(raw) # raises on parse error — fail fast
|
||||
validate_sql(sql, recon) # raises ReconValidationError — fail fast
|
||||
|
||||
result = T2SResult(sql=sql, used_tables=_extract_tables(sql))
|
||||
span.update(output={"sql": sql, "used_tables": result.used_tables})
|
||||
@@ -56,10 +52,20 @@ def _extract_sql(text: str) -> str:
|
||||
return text.strip().rstrip(";").strip()
|
||||
|
||||
|
||||
def _validate(sql: str) -> None:
|
||||
def _normalize(sql: str) -> str:
|
||||
"""Parse the LLM's SQL, then re-render with every identifier quoted.
|
||||
|
||||
Postgres folds unquoted identifiers to lowercase, so `d.A13` becomes
|
||||
`d.a13` and breaks against case-preserving columns like district."A13".
|
||||
Re-rendering with `identify=True` puts double quotes around every
|
||||
identifier — case is preserved, and Postgres treats a quoted lowercase
|
||||
identifier the same as the unquoted form, so this is safe in both
|
||||
directions.
|
||||
"""
|
||||
parsed = sqlglot.parse_one(sql, dialect="postgres")
|
||||
if parsed is None:
|
||||
raise ValueError("sqlglot returned no parse tree")
|
||||
return parsed.sql(dialect="postgres", identify=True)
|
||||
|
||||
|
||||
def _extract_tables(sql: str) -> list[str]:
|
||||
|
||||
@@ -1,11 +1,7 @@
|
||||
"""Public data shapes for the Tools layer.
|
||||
|
||||
Behavior modules in api/tools/ import from here. Keeping types in one file
|
||||
makes the layer's surface visible at a glance.
|
||||
|
||||
Secondary constructors (`from_inspector`, `from_spec`, …) live as
|
||||
classmethods so call sites read as one line instead of multi-line kwarg
|
||||
blocks.
|
||||
Tool-output types only. The dataset model (Column, Table, Metric, Recon)
|
||||
lives in `api/recon/types.py` — Tools consume it but don't own it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -14,100 +10,6 @@ from dataclasses import dataclass
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Table:
|
||||
name: str
|
||||
description: str | None
|
||||
columns: list[Column]
|
||||
|
||||
|
||||
@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":
|
||||
"""Build from a metrics.yaml entry."""
|
||||
return cls(
|
||||
name=name,
|
||||
description=spec.get("description", ""),
|
||||
sql=spec["sql"],
|
||||
from_table=spec["from_table"],
|
||||
filter=spec.get("filter"),
|
||||
unit=spec.get("unit"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SchemaContext:
|
||||
schema: str
|
||||
tables: dict[str, Table]
|
||||
metrics: dict[str, Metric]
|
||||
|
||||
def table_names(self) -> list[str]:
|
||||
return sorted(self.tables)
|
||||
|
||||
def metric_names(self) -> list[str]:
|
||||
return sorted(self.metrics)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
# ── Query execution ──
|
||||
|
||||
@dataclass
|
||||
class QueryResult:
|
||||
columns: list[str]
|
||||
@@ -119,8 +21,6 @@ class QueryResult:
|
||||
return [dict(zip(self.columns, r)) for r in self.rows]
|
||||
|
||||
|
||||
# ── Text-to-SQL output ──
|
||||
|
||||
@dataclass
|
||||
class T2SResult:
|
||||
sql: str
|
||||
|
||||
Reference in New Issue
Block a user