recon + sqlglot validator + drill_down package; guard ReAct dimension picks against candidate list
This commit is contained in:
250
api/recon/types.py
Normal file
250
api/recon/types.py
Normal file
@@ -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", [])],
|
||||
)
|
||||
Reference in New Issue
Block a user