verbose live UI + tool-level SSE events + Groq default + regression tests

api/tools/db.py

@@ -0,0 +1,39 @@
+"""SQLAlchemy engine for the warehouse.
+
+Per-connection defaults (search_path, statement_timeout) are baked into the
+engine via libpq's `options` connect_arg so we never SET them at the SQL
+layer. Read-only mode for query execution is requested via SQLAlchemy's
+`execution_options(postgresql_readonly=True)`.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+
+from sqlalchemy import create_engine
+from sqlalchemy.engine import Engine
+
+from api.config import get_settings
+
+SCHEMA = "financial"
+DEFAULT_STATEMENT_TIMEOUT_MS = 10_000
+
+
+@lru_cache(maxsize=1)
+def get_engine() -> Engine:
+    url = get_settings().database_url
+    if url.startswith("postgresql://"):
+        url = url.replace("postgresql://", "postgresql+psycopg://", 1)
+    return create_engine(
+        url,
+        pool_pre_ping=True,
+        future=True,
+        # libpq -c options run before any SQL the app issues — no need to
+        # SET search_path / statement_timeout at execution time.
+        connect_args={
+            "options": (
+                f"-c search_path={SCHEMA},public "
+                f"-c statement_timeout={DEFAULT_STATEMENT_TIMEOUT_MS}"
+            ),
+        },
+    )

api/tools/execute_sql.py

@@ -0,0 +1,55 @@
+"""Execute read-only SQL against the financial warehouse.
+
+Defence in depth:
+- Engine-level `search_path` and `statement_timeout` via libpq options in
+  `api.tools.db.get_engine` — no SQL-side SETs.
+- `execution_options(postgresql_readonly=True, postgresql_deferrable=True)`
+  — SQLAlchemy emits `SET TRANSACTION READ ONLY DEFERRABLE` internally; we
+  don't compose a SQL string for it.
+- Hard refusal of any statement that doesn't begin with SELECT or WITH.
+- Row cap on the returned result set.
+- No commit — read-only transactions roll back on context exit.
+
+The SQL passed in IS a string (it's the LLM's output); nothing on the Python
+side concatenates SQL fragments around it.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from sqlalchemy import text
+
+from api import langfuse_client as lf
+from api.tools.db import get_engine
+from api.tools.types import QueryResult
+
+logger = logging.getLogger("nvi.tools.execute_sql")
+
+MAX_ROWS = 1000
+
+
+def execute_sql(sql: str) -> QueryResult:
+    sql_clean = sql.strip().rstrip(";").strip()
+    head = sql_clean.split(None, 1)[0].upper() if sql_clean else ""
+    if head not in {"SELECT", "WITH"}:
+        raise ValueError(f"refusing non-SELECT statement (starts with {head!r})")
+
+    engine = get_engine()
+    with lf.span("execute_sql", input={"sql": sql_clean}) as span:
+        ro_engine = engine.execution_options(
+            postgresql_readonly=True,
+            postgresql_deferrable=True,
+        )
+        with ro_engine.connect() as conn:
+            res = conn.execute(text(sql_clean))
+            columns = list(res.keys())
+            fetched = res.fetchmany(MAX_ROWS + 1)
+
+        truncated = len(fetched) > MAX_ROWS
+        rows = [list(r) for r in fetched[:MAX_ROWS]]
+        result = QueryResult(
+            columns=columns, rows=rows, row_count=len(rows), truncated=truncated
+        )
+        span.update(output={"row_count": result.row_count, "truncated": truncated})
+        return result

api/tools/schema.py

@@ -0,0 +1,57 @@
+"""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)

api/tools/text_to_sql.py

@@ -0,0 +1,70 @@
+"""LLM-driven NL → SQL with sqlglot validation and one retry."""
+
+from __future__ import annotations
+
+import logging
+import re
+
+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.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()
+    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(),
+            retry_hint=retry_hint,
+        )
+
+    with lf.span(
+        "text_to_sql",
+        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)
+
+        result = T2SResult(sql=sql, used_tables=_extract_tables(sql))
+        span.update(output={"sql": sql, "used_tables": result.used_tables})
+        return result
+
+
+def _extract_sql(text: str) -> str:
+    m = re.search(r"```sql\s*(.*?)```", text, re.DOTALL | re.IGNORECASE)
+    if m:
+        return m.group(1).strip().rstrip(";").strip()
+    return text.strip().rstrip(";").strip()
+
+
+def _validate(sql: str) -> None:
+    parsed = sqlglot.parse_one(sql, dialect="postgres")
+    if parsed is None:
+        raise ValueError("sqlglot returned no parse tree")
+
+
+def _extract_tables(sql: str) -> list[str]:
+    try:
+        parsed = sqlglot.parse_one(sql, dialect="postgres")
+        return sorted({t.name for t in parsed.find_all(sqlglot.exp.Table)})
+    except Exception:
+        return []

api/tools/types.py

@@ -0,0 +1,128 @@
+"""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