app.py

"""
Minimal Gradio demo for LocalSQL.

Deploy this on a GPU-backed Space. The default sample database is generated at
startup so the repo does not need to track a binary SQLite file.
"""
from __future__ import annotations

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_ADAPTER, DEFAULT_BASE, load_model, predict

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


@lru_cache(maxsize=1)
def get_model():
    print("Loading LocalSQL model in bf16 for ZeroGPU...", flush=True)
    return load_model(DEFAULT_BASE, DEFAULT_ADAPTER, use_4bit=False)


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

    db_path = ensure_sample_db()
    try:
        model, tokenizer = get_model()
        result = predict(
            str(db_path),
            question.strip(),
            model,
            tokenizer,
            evidence=evidence.strip(),
            execute=run_sql,
        )
    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"]


with gr.Blocks(title="LocalSQL") as demo:
    gr.Markdown(
        "# LocalSQL\n"
        "Ask a sample SQLite database questions in plain English. "
        "This hosted demo uses sample data; run LocalSQL locally for private databases."
    )
    gr.Markdown(
        "On ZeroGPU, the GPU is allocated only while a query is running. "
        "The first query can take longer while the model downloads and loads."
    )

    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", value=True)
    submit = gr.Button("Ask LocalSQL", variant="primary")

    sql = gr.Code(label="Generated SQL", language="sql")
    rows = gr.Dataframe(label="Results", interactive=False)
    status = gr.Textbox(label="Status")
    schema = gr.Code(label="Schema", language="sql")

    submit.click(
        answer,
        inputs=[question, evidence, run_sql],
        outputs=[sql, rows, status, schema],
    )

    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],
    )


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