129 lines
3.4 KiB
Python
129 lines
3.4 KiB
Python
"""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.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
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]
|
|
rows: list[list[Any]]
|
|
row_count: int
|
|
truncated: bool
|
|
|
|
def as_dicts(self) -> list[dict[str, Any]]:
|
|
return [dict(zip(self.columns, r)) for r in self.rows]
|
|
|
|
|
|
# ── Text-to-SQL output ──
|
|
|
|
@dataclass
|
|
class T2SResult:
|
|
sql: str
|
|
used_tables: list[str]
|
|
explanation: str | None = None
|