app.py

"""
Gradio demo for LocalSQL.

Runs the merged CoT-SFT model (qwen2.5-coder-7b-bird-cot) on a GPU-backed Space.
Users can query a built-in sample database or upload their own .sqlite file. The
model reasons step by step, then emits the final SQL, which is executed
read-only against the chosen database.
"""
from __future__ import annotations

import html
import sqlite3
import traceback
import tempfile
from functools import lru_cache
from pathlib import Path

import gradio as gr
import pandas as pd

from src.text2sql import DEFAULT_MODEL, load_model, predict
from src.shared.schema_loader import get_schema_from_sqlite

try:
    import spaces
except ImportError:
    class _SpacesShim:
        @staticmethod
        def GPU(*args, **kwargs):
            def decorator(fn):
                return fn
            return decorator

    spaces = _SpacesShim()

SAMPLE_ROOT = Path(tempfile.gettempdir()) / "localsql_samples"
SAMPLE_DB = SAMPLE_ROOT / "music_store.sqlite"


def ensure_sample_db() -> Path:
    SAMPLE_ROOT.mkdir(parents=True, exist_ok=True)
    if SAMPLE_DB.exists():
        return SAMPLE_DB

    conn = sqlite3.connect(SAMPLE_DB)
    cur = conn.cursor()
    cur.executescript(
        """
        CREATE TABLE Artist (
            ArtistId INTEGER PRIMARY KEY,
            Name TEXT NOT NULL
        );
        CREATE TABLE Album (
            AlbumId INTEGER PRIMARY KEY,
            Title TEXT NOT NULL,
            ArtistId INTEGER NOT NULL,
            FOREIGN KEY (ArtistId) REFERENCES Artist(ArtistId)
        );
        CREATE TABLE Track (
            TrackId INTEGER PRIMARY KEY,
            Name TEXT NOT NULL,
            AlbumId INTEGER NOT NULL,
            Milliseconds INTEGER,
            UnitPrice REAL NOT NULL,
            FOREIGN KEY (AlbumId) REFERENCES Album(AlbumId)
        );
        CREATE TABLE Customer (
            CustomerId INTEGER PRIMARY KEY,
            FirstName TEXT NOT NULL,
            LastName TEXT NOT NULL,
            Country TEXT NOT NULL
        );
        CREATE TABLE Invoice (
            InvoiceId INTEGER PRIMARY KEY,
            CustomerId INTEGER NOT NULL,
            InvoiceDate TEXT NOT NULL,
            Total REAL NOT NULL,
            FOREIGN KEY (CustomerId) REFERENCES Customer(CustomerId)
        );

        INSERT INTO Artist VALUES
            (1, 'Iron Maiden'),
            (2, 'Led Zeppelin'),
            (3, 'Miles Davis'),
            (4, 'Nina Simone');
        INSERT INTO Album VALUES
            (1, 'Powerslave', 1),
            (2, 'Seventh Son of a Seventh Son', 1),
            (3, 'Led Zeppelin IV', 2),
            (4, 'Kind of Blue', 3),
            (5, 'I Put a Spell on You', 4);
        INSERT INTO Track VALUES
            (1, 'Aces High', 1, 271000, 0.99),
            (2, '2 Minutes to Midnight', 1, 360000, 0.99),
            (3, 'Stairway to Heaven', 3, 482000, 0.99),
            (4, 'So What', 4, 545000, 1.29),
            (5, 'Feeling Good', 5, 178000, 1.29);
        INSERT INTO Customer VALUES
            (1, 'Ada', 'Lovelace', 'United Kingdom'),
            (2, 'Grace', 'Hopper', 'United States'),
            (3, 'Katherine', 'Johnson', 'United States');
        INSERT INTO Invoice VALUES
            (1, 1, '2024-01-15', 12.87),
            (2, 2, '2024-02-20', 22.74),
            (3, 2, '2024-03-12', 8.91),
            (4, 3, '2024-03-30', 14.85);
        """
    )
    conn.commit()
    conn.close()
    return SAMPLE_DB


def is_sqlite_file(path: str) -> bool:
    """Cheap read-only check that a path is a usable SQLite database."""
    try:
        conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
        conn.execute("SELECT name FROM sqlite_master LIMIT 1")
        conn.close()
        return True
    except Exception:
        return False


def resolve_db(uploaded_db: str | None) -> tuple[str | None, str | None]:
    """Return (db_path, error). Uploaded file wins; else the sample database."""
    if uploaded_db:
        if is_sqlite_file(uploaded_db):
            return uploaded_db, None
        return None, "That file is not a valid SQLite database. Upload a .sqlite / .db file."
    return str(ensure_sample_db()), None


def show_schema(uploaded_db: str | None) -> str:
    """Preview the raw DDL of whichever database is currently selected."""
    db_path, err = resolve_db(uploaded_db)
    if err:
        return f"-- {err} --"
    try:
        return get_schema_from_sqlite(db_path)
    except Exception as exc:  # noqa: BLE001
        return f"-- could not read schema: {exc} --"


def _badge(text: str, color: str) -> str:
    return (
        f'<span style="background:{color}22;color:{color};border-radius:4px;'
        f'padding:0 5px;font-size:11px;margin-left:6px">{text}</span>'
    )


def schema_map_html(uploaded_db: str | None) -> str:
    """Render the selected database as a grid of table cards (columns + PK/FK
    badges) plus a relationships line, so users can see what they can ask about.
    """
    db_path, err = resolve_db(uploaded_db)
    if err:
        return f"<p>{html.escape(err)}</p>"
    try:
        con = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
        con.row_factory = sqlite3.Row
        tables = [
            r[0] for r in con.execute(
                "SELECT name FROM sqlite_master WHERE type='table' "
                "AND name NOT LIKE 'sqlite_%' ORDER BY name"
            )
        ]
        if not tables:
            con.close()
            return "<p>No tables found in this database.</p>"

        cards, rels = [], []
        for t in tables:
            cols = con.execute(f'PRAGMA table_info("{t}")').fetchall()
            fks = con.execute(f'PRAGMA foreign_key_list("{t}")').fetchall()
            try:
                n = con.execute(f'SELECT COUNT(*) FROM "{t}"').fetchone()[0]
            except Exception:  # noqa: BLE001
                n = "?"
            fk_from = {f["from"]: (f["table"], f["to"]) for f in fks}

            col_rows = []
            for c in cols:
                name = html.escape(str(c["name"]))
                ctype = html.escape(str(c["type"] or ""))
                if c["pk"]:
                    badge = _badge("PK", "#3b82f6")
                elif c["name"] in fk_from:
                    badge = _badge("FK&rarr;" + html.escape(fk_from[c["name"]][0]), "#10b981")
                else:
                    badge = ""
                col_rows.append(
                    '<div style="display:flex;justify-content:space-between;gap:10px;padding:1px 0">'
                    f'<span>{name}{badge}</span>'
                    f'<span style="opacity:.5;font-size:11px">{ctype}</span></div>'
                )
            for f in fks:
                rels.append(
                    f'{html.escape(t)}.{html.escape(str(f["from"]))} &rarr; '
                    f'{html.escape(str(f["table"]))}.{html.escape(str(f["to"]))}'
                )

            cards.append(
                '<div style="border:1px solid rgba(128,128,128,.35);border-radius:10px;'
                'padding:10px 12px;min-width:190px;flex:0 0 auto">'
                f'<div style="font-weight:600;border-bottom:1px solid rgba(128,128,128,.25);'
                f'padding-bottom:5px;margin-bottom:5px">{html.escape(t)} '
                f'<span style="opacity:.5;font-weight:400;font-size:12px">{n} rows</span></div>'
                + "".join(col_rows) + "</div>"
            )
        con.close()

        grid = (
            '<div style="display:flex;flex-wrap:wrap;gap:12px;align-items:flex-start">'
            + "".join(cards) + "</div>"
        )
        rel_html = (
            '<div style="margin-top:10px;opacity:.8;font-size:13px"><b>Relationships:</b> '
            + "; ".join(rels) + "</div>"
        ) if rels else ""
        return grid + rel_html
    except Exception as exc:  # noqa: BLE001
        return f"<p>Could not read database: {html.escape(str(exc))}</p>"


def _list_tables(db_path: str) -> list[str]:
    con = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
    tables = [
        r[0] for r in con.execute(
            "SELECT name FROM sqlite_master WHERE type='table' "
            "AND name NOT LIKE 'sqlite_%' ORDER BY name"
        )
    ]
    con.close()
    return tables


def browse_table(uploaded_db: str | None, table_name: str | None) -> pd.DataFrame:
    """Return the first rows of the selected table (read-only preview)."""
    if not table_name:
        return pd.DataFrame()
    db_path, err = resolve_db(uploaded_db)
    if err:
        return pd.DataFrame()
    try:
        con = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
        # table_name comes from our own dropdown (populated from sqlite_master).
        cur = con.execute(f'SELECT * FROM "{table_name}" LIMIT 10')
        cols = [d[0] for d in cur.description]
        data = cur.fetchall()
        con.close()
        return pd.DataFrame(data, columns=cols)
    except Exception:  # noqa: BLE001
        return pd.DataFrame()


def refresh_browser(uploaded_db: str | None):
    """Repopulate the table dropdown and preview the first table's rows."""
    db_path, err = resolve_db(uploaded_db)
    if err:
        return gr.update(choices=[], value=None), pd.DataFrame()
    tables = _list_tables(db_path)
    first = tables[0] if tables else None
    return gr.update(choices=tables, value=first), browse_table(uploaded_db, first)


def prefetch_weights() -> None:
    """Download weights to disk at startup (CPU, no GPU timer).

    ZeroGPU caps a GPU function at 120s on the free tier, so the 15 GB download
    must not happen inside the GPU window — only the load-to-GPU + inference do.
    """
    try:
        from huggingface_hub import snapshot_download

        print(f"Pre-fetching {DEFAULT_MODEL} weights to disk...", flush=True)
        snapshot_download(DEFAULT_MODEL)
        print("Weights pre-fetched.", flush=True)
    except Exception as exc:  # noqa: BLE001
        print(f"Weight prefetch skipped ({exc}); will download on first query.", flush=True)


@lru_cache(maxsize=1)
def get_model():
    print(f"Loading {DEFAULT_MODEL} in bf16 for ZeroGPU...", flush=True)
    return load_model(DEFAULT_MODEL, adapter=None, use_4bit=False)


@spaces.GPU(duration=120)
def answer(question: str, evidence: str, run_sql: bool, uploaded_db: str | None):
    if not question.strip():
        return "", pd.DataFrame(), "Ask a question first.", "", ""

    db_path, err = resolve_db(uploaded_db)
    if err:
        return "", pd.DataFrame(), err, "", ""

    try:
        model, tokenizer = get_model()
        result = predict(
            db_path,
            question.strip(),
            model,
            tokenizer,
            evidence=evidence.strip(),
            execute=run_sql,
            cot=True,
            max_new_tokens=512,
        )
    except Exception as exc:
        traceback.print_exc()
        message = f"{type(exc).__name__}: {exc!s}" if str(exc) else repr(exc)
        return "", pd.DataFrame(), f"Model error: {message}", "", ""

    frame = (
        pd.DataFrame(result["rows"], columns=result["columns"])
        if result["rows"]
        else pd.DataFrame()
    )
    status = result["error"] or f"{result['row_count']} rows"
    return result["sql"], frame, status, result["schema"], result["reasoning"]


INTRO = """\
# LocalSQL

**Ask a database questions in plain English.** A 7B model reasons over your
schema, writes the SQL, and runs it read-only.

Model: [`qwen2.5-coder-7b-bird-cot`](https://huggingface.co/jk200201/qwen2.5-coder-7b-bird-cot) &nbsp;·&nbsp; BIRD dev: **52.1%** greedy, **58.5%** self-consistency (K=8) &nbsp;·&nbsp; matches DeepSeek V4-Pro (1.6T) at ~0.4% the size.
"""

NOTE = """\
Try the built-in sample store below, or upload your own `.sqlite` / `.db` file.
Uploaded files are used only for the current query and are not stored. The first
query takes a minute or two while the model loads onto the GPU, then it is fast.
"""


with gr.Blocks(title="LocalSQL") as demo:
    gr.Markdown(INTRO)
    gr.Markdown(NOTE)

    with gr.Accordion("What's in this database? (tables, columns, keys)", open=True):
        schema_map = gr.HTML()

    with gr.Accordion("Browse data (peek at rows)", open=False):
        # allow_custom_value: choices are populated dynamically by refresh_browser,
        # so the backend must not reject a value it has not registered yet.
        table_dropdown = gr.Dropdown(
            label="Table", choices=[], interactive=True, allow_custom_value=True
        )
        browse_df = gr.Dataframe(label="First 10 rows", interactive=False)

    with gr.Row():
        with gr.Column(scale=1):
            uploaded_db = gr.File(
                label="Database (optional) — upload a .sqlite / .db file",
                file_types=[".sqlite", ".db", ".sqlite3", ".db3"],
                type="filepath",
            )
            question = gr.Textbox(
                label="Question",
                value="Which artists have the most albums?",
                lines=2,
            )
            evidence = gr.Textbox(
                label="Optional domain hint",
                placeholder="Example: revenue = unit price * quantity",
                lines=2,
            )
            run_sql = gr.Checkbox(label="Execute generated SQL (read-only)", value=True)
            submit = gr.Button("Ask LocalSQL", variant="primary")

            gr.Examples(
                examples=[
                    ["Which artists have the most albums?", "", True],
                    ["List customers by total invoice amount.", "", True],
                    ["What is the longest track?", "", True],
                ],
                inputs=[question, evidence, run_sql],
            )

        with gr.Column(scale=1):
            sql = gr.Code(label="Generated SQL", language="sql")
            rows = gr.Dataframe(label="Results", interactive=False)
            status = gr.Textbox(label="Status")
            with gr.Accordion("Model reasoning (chain of thought)", open=False):
                reasoning = gr.Markdown()
            with gr.Accordion("Schema (raw SQL)", open=False):
                schema = gr.Code(language="sql")

    submit.click(
        answer,
        inputs=[question, evidence, run_sql, uploaded_db],
        outputs=[sql, rows, status, schema, reasoning],
    )
    # Preview rows when the user picks a table.
    table_dropdown.change(
        browse_table, inputs=[uploaded_db, table_dropdown], outputs=[browse_df]
    )
    # Refresh the schema map, raw DDL, and table browser for the selected database.
    for trigger in (uploaded_db.change, demo.load):
        trigger(schema_map_html, inputs=[uploaded_db], outputs=[schema_map])
        trigger(show_schema, inputs=[uploaded_db], outputs=[schema])
        trigger(refresh_browser, inputs=[uploaded_db], outputs=[table_dropdown, browse_df])


ensure_sample_db()
# On HF Spaces (SPACE_ID is set) prefetch weights at startup so the 15 GB pull is
# outside the GPU window. Skipped locally so importing the module stays cheap.
import os  # noqa: E402
if os.getenv("SPACE_ID"):
    prefetch_weights()

if __name__ == "__main__":
    demo.launch()