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

ctrl/seed/download.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# Downloads BIRD financial SQLite into ctrl/seed/data/financial.sqlite and
+# the golden eval set into api/evals/data/. Runs on the HOST (the api image
+# deliberately doesn't include the data). Override BIRD_URL if upstream
+# moves.
+#
+# Layout inside the outer dev.zip (as of the 2024-06-27 dump):
+#   dev_<date>/
+#     dev.sql, dev.json, dev_tables.json, dev_tied_append.json
+#     dev_databases.zip  → financial/financial.sqlite (and 10 other DBs)
+set -euo pipefail
+cd "$(dirname "$0")"
+
+DEST="data/financial.sqlite"
+EVAL_DIR="../../api/evals/data"
+mkdir -p data "$EVAL_DIR"
+
+if [ -f "$DEST" ]; then
+    echo "already have $DEST ($(du -h "$DEST" | cut -f1))"
+    exit 0
+fi
+
+URL="${BIRD_URL:-https://bird-bench.oss-cn-beijing.aliyuncs.com/dev.zip}"
+TMP="$(mktemp -d)"
+trap 'rm -rf "$TMP"' EXIT
+OUTER="$TMP/dev.zip"
+
+echo "fetching $URL (~330MB compressed; full dev set is ~3GB extracted)..."
+curl -fL --progress-bar -o "$OUTER" "$URL"
+
+echo "extracting outer archive..."
+unzip -q "$OUTER" -d "$TMP"
+ROOT=$(find "$TMP" -maxdepth 1 -type d -name 'dev_*' | head -1)
+test -n "$ROOT" || { echo "could not locate dev_* root inside outer zip" >&2; exit 1; }
+
+echo "extracting financial.sqlite from inner archive..."
+unzip -q -j "$ROOT/dev_databases.zip" '*/financial/financial.sqlite' -d data/
+
+echo "copying golden eval set to $EVAL_DIR/..."
+cp "$ROOT/dev.json" "$ROOT/dev.sql" "$ROOT/dev_tables.json" "$EVAL_DIR/"
+
+echo "ready:"
+echo "  $DEST ($(du -h "$DEST" | cut -f1))"
+echo "  $(ls -1 "$EVAL_DIR/" | wc -l) files in $EVAL_DIR/"

ctrl/seed/load_bird.py

@@ -0,0 +1,147 @@
+"""Load a BIRD SQLite dataset into Postgres.
+
+Runs inside the api container (`make seed`). Reads which dataset to load
+from `settings.dataset` (env var `NVI_DATASET`, default `financial`) and
+expects the source SQLite at `ctrl/seed/data/<dataset>.sqlite` — produced
+on the host by `ctrl/seed/download.sh` and reaching the container via
+Tilt live_update sync (kind) or the docker-compose bind mount.
+
+DDL is built from SQLAlchemy `MetaData` + `Table` + `Column` objects, not
+string SQL. Bulk insert uses psycopg's COPY protocol via the engine's raw
+connection — there's no SQLAlchemy equivalent, and INSERT for ~1M rows
+would be orders of magnitude slower. Identifiers in the COPY statement
+are rendered through SA's dialect-aware identifier preparer.
+"""
+
+from __future__ import annotations
+
+import sqlite3
+import sys
+from pathlib import Path
+from typing import Any
+
+from sqlalchemy import (
+    BigInteger, Boolean, Column, Date, DateTime, Float, Integer, LargeBinary,
+    MetaData, Numeric, SmallInteger, Table, Text, text,
+)
+from sqlalchemy.engine import Engine
+
+from api.config import get_settings
+from api.tools.db import get_engine
+
+# SQLite is type-affinity, not strict — pick a SQLAlchemy type per affinity.
+# Anything unrecognised falls through to Text, which is always safe.
+TYPE_MAP: dict[str, Any] = {
+    "INT": BigInteger,
+    "INTEGER": BigInteger,
+    "TINYINT": SmallInteger,
+    "SMALLINT": SmallInteger,
+    "MEDIUMINT": Integer,
+    "BIGINT": BigInteger,
+    "REAL": Float,
+    "DOUBLE": Float,
+    "FLOAT": Float,
+    "NUMERIC": Numeric,
+    "DECIMAL": Numeric,
+    "BOOLEAN": Boolean,
+    "DATE": Date,
+    "DATETIME": DateTime,
+    "TIMESTAMP": DateTime,
+    "TEXT": Text,
+    "CHAR": Text,
+    "VARCHAR": Text,
+    "CLOB": Text,
+    "BLOB": LargeBinary,
+}
+
+
+def _sa_type(sqlite_type: str):
+    base = sqlite_type.upper().split("(")[0].strip()
+    return TYPE_MAP.get(base, Text)
+
+
+def _list_tables(src: sqlite3.Connection) -> list[str]:
+    return [
+        r[0] for r in src.execute(
+            "SELECT name FROM sqlite_master "
+            "WHERE type='table' AND name NOT LIKE 'sqlite_%' "
+            "ORDER BY name"
+        ).fetchall()
+    ]
+
+
+def _discover_table(src: sqlite3.Connection, metadata: MetaData,
+                    name: str, schema: str) -> Table:
+    cols = src.execute(f'PRAGMA table_info("{name}")').fetchall()
+    # PRAGMA returns (cid, name, type, notnull, dflt_value, pk).
+    sa_cols = [
+        Column(
+            c[1],
+            _sa_type(c[2]),
+            nullable=not c[3],
+            primary_key=bool(c[5]),
+        )
+        for c in cols
+    ]
+    return Table(name, metadata, *sa_cols, schema=schema)
+
+
+def _copy_table(engine: Engine, src: sqlite3.Connection, sa_table: Table) -> int:
+    """Bulk-load via psycopg COPY. Identifiers go through SA's dialect preparer
+    so we never hand-quote names."""
+    prep = engine.dialect.identifier_preparer
+    table_ident = prep.format_table(sa_table)                       # "<schema>"."<table>"
+    col_list = ", ".join(prep.quote(c.name) for c in sa_table.columns)
+    src_ident = '"' + sa_table.name.replace('"', '""') + '"'        # SQLite side
+
+    count = 0
+    raw = engine.raw_connection()
+    try:
+        cur = raw.cursor()
+        with cur.copy(f"COPY {table_ident} ({col_list}) FROM STDIN") as copy:
+            for row in src.execute(f"SELECT * FROM {src_ident}"):
+                copy.write_row(tuple(row))
+                count += 1
+                if count % 100_000 == 0:
+                    print(f"  {sa_table.name}: {count:,} rows...")
+        raw.commit()
+    finally:
+        raw.close()
+    return count
+
+
+def main() -> None:
+    dataset = get_settings().dataset
+    sqlite_path = Path(f"seed/data/{dataset}.sqlite")
+
+    if not sqlite_path.exists():
+        sys.exit(
+            f"{dataset} SQLite not found at {sqlite_path}.\n"
+            f"Run `bash ctrl/seed/download.sh` on the host first."
+        )
+
+    engine = get_engine()
+    src = sqlite3.connect(sqlite_path)
+    src.row_factory = sqlite3.Row
+
+    prep = engine.dialect.identifier_preparer
+    schema_ident = prep.quote_schema(dataset)
+    with engine.begin() as conn:
+        conn.execute(text(f"DROP SCHEMA IF EXISTS {schema_ident} CASCADE"))
+        conn.execute(text(f"CREATE SCHEMA {schema_ident}"))
+
+    metadata = MetaData(schema=dataset)
+    tables = [_discover_table(src, metadata, t, dataset) for t in _list_tables(src)]
+    metadata.create_all(engine)
+    print(f"created {len(tables)} tables in schema {dataset!r}")
+
+    for sa_table in tables:
+        count = _copy_table(engine, src, sa_table)
+        print(f"  loaded {sa_table.name}: {count:,} rows")
+
+    src.close()
+    print("done.")
+
+
+if __name__ == "__main__":
+    main()