'
)
def _schema_preview_html(table: str, schema: list[dict]) -> str:
"""Render the column names + types of the loaded table."""
if not schema:
return ""
cols = "".join(
f''
f'{c["name"]}'
f'{c["type"]}'
f''
for c in schema
)
return (
'
'
f'
{table} · {len(schema)} columns
'
f'
{cols}
'
'
'
)
def _download_links_html(sql: str, results: list[dict], svg: str) -> str:
"""Inline data-URL download links for CSV (results) and SVG (chart)."""
import base64
import csv
import io as _io
parts = []
# CSV download
if results:
buf = _io.StringIO()
writer = csv.DictWriter(buf, fieldnames=list(results[0].keys()))
writer.writeheader()
for r in results:
writer.writerow({k: ("" if v is None else v) for k, v in r.items()})
csv_b64 = base64.b64encode(buf.getvalue().encode("utf-8")).decode("ascii")
parts.append(
f''
f'↓ CSV ({len(results):,} rows)'
)
# SVG download (standalone version: explicit dims, white bg, XML prolog)
if svg and "'
f'↓ SVG'
)
if not parts:
return ""
return f'
{"".join(parts)}
'
def _suggestions_html(qs: list[str]) -> str:
if not qs:
return ""
chips = "".join(
f'{q}'
for q in qs
)
return f'
{chips}
'
def _data_table_html(rows: list[dict], max_rows: int = 10) -> str:
if not rows:
return '
No rows.
'
df = pd.DataFrame(rows[:max_rows])
cols = df.columns.tolist()
head = "".join(f"
{c}
" for c in cols)
body = "".join(
"
" + "".join(
f"
{('' if r.get(c) is None else r.get(c, ''))}
" for c in cols
) + "
"
for r in rows[:max_rows]
)
note = (
f'
Showing {min(max_rows, len(rows))} of {len(rows):,} rows
'
if len(rows) > max_rows else ""
)
return f'
{head}
{body}
{note}'
PIPELINE_STAGES = [
("sql", "Generating SQL"),
("execute", "Running query"),
("chart", "Designing chart"),
("render", "Rendering visualization"),
]
def _turn_html_progress(question: str, current_stage: str = "sql") -> str:
"""Render a turn that's still in progress with multi-step pipeline view."""
items = []
seen_active = False
for key, label in PIPELINE_STAGES:
if key == current_stage:
cls = "active"
seen_active = True
elif seen_active:
cls = "pending"
else:
cls = "done"
items.append(
f'
'
f''
f'{label}'
f'
'
)
return (
'
'
f'
{question}
'
f'
{"".join(items)}
'
'
'
)
SQL_KEYWORDS = {
"SELECT", "FROM", "WHERE", "GROUP", "BY", "ORDER", "HAVING", "LIMIT",
"OFFSET", "JOIN", "LEFT", "RIGHT", "INNER", "OUTER", "FULL", "ON",
"AS", "AND", "OR", "NOT", "IN", "BETWEEN", "LIKE", "IS", "NULL",
"DISTINCT", "WITH", "UNION", "ALL", "CASE", "WHEN", "THEN", "ELSE",
"END", "ASC", "DESC", "INSERT", "INTO", "VALUES", "UPDATE", "SET",
"DELETE", "CREATE", "TABLE", "DROP", "ALTER",
}
SQL_FUNCTIONS = {
"COUNT", "SUM", "AVG", "MIN", "MAX", "DATE", "EXTRACT", "CAST",
"COALESCE", "ROUND", "ABS", "LENGTH", "LOWER", "UPPER", "SUBSTR",
"TRIM", "REPLACE", "CONCAT", "STRFTIME", "DATEDIFF", "NOW",
}
def _highlight_sql(sql: str) -> str:
"""Tokenize and add span highlighting to a SQL query."""
import re as _re, html as _html
out = []
i = 0
n = len(sql)
while i < n:
c = sql[i]
# Comments: -- to end of line
if c == "-" and i + 1 < n and sql[i + 1] == "-":
j = sql.find("\n", i)
if j == -1: j = n
out.append(f'{_html.escape(sql[i:j])}')
i = j
# Strings
elif c in ("'", '"'):
quote = c
j = i + 1
while j < n and sql[j] != quote:
if sql[j] == "\\" and j + 1 < n:
j += 2
else:
j += 1
j = min(j + 1, n)
out.append(f'{_html.escape(sql[i:j])}')
i = j
# Numbers
elif c.isdigit():
j = i
while j < n and (sql[j].isdigit() or sql[j] == "."):
j += 1
out.append(f'{sql[i:j]}')
i = j
# Identifiers / keywords / functions
elif c.isalpha() or c == "_":
j = i
while j < n and (sql[j].isalnum() or sql[j] == "_"):
j += 1
word = sql[i:j]
upper = word.upper()
if upper in SQL_KEYWORDS:
out.append(f'{word}')
elif upper in SQL_FUNCTIONS:
out.append(f'{word}')
else:
out.append(_html.escape(word))
i = j
else:
out.append(_html.escape(c))
i += 1
return "".join(out)
def _chart_badge(spec: dict, results_count: int) -> str:
"""Small uppercase badge showing chart type and row count."""
if not spec:
return ""
chart_type = (spec.get("chart_type") or "?").upper()
if results_count > 30:
return f'
')
if result.get("svg"):
badge = _chart_badge(result.get("chart_spec"), len(result.get("results") or []))
parts.append(
f'
{badge}{result["svg"]}
'
)
# Narration: 1-2 sentence finding from the analyst persona
if result.get("narration"):
parts.append(f'
{result["narration"]}
')
# Inline download links for CSV + SVG
parts.append(_download_links_html(
result.get("sql") or "",
result.get("results") or [],
result.get("svg") or "",
))
if result.get("sql"):
sql_safe = _highlight_sql(result["sql"])
sql_raw = result["sql"].replace("`", "\\`").replace("\\", "\\\\")
copy_js = (
f"navigator.clipboard.writeText(`{sql_raw}`);"
"this.textContent='Copied';setTimeout(()=>{this.textContent='Copy'},1200);"
)
parts.append(
'SQL query'
'
'
f''
f'
{sql_safe}
'
'
'
''
)
if result.get("results"):
parts.append(
'Data'
f'{_data_table_html(result["results"])}'
''
)
return f'
{"".join(parts)}
'
def _conversation_html(history: list[dict], in_progress: tuple[str, str] | None = None) -> str:
"""Conversation HTML reacts to: data loaded? history? in-progress?
in_progress is a tuple of (question, current_stage_key)."""
has_data = bool(get_agent().list_tables())
has_turns = bool(history) or in_progress is not None
if not has_data and not has_turns:
return _empty_state_html()
if has_data and not has_turns:
return _ready_state_html()
out = "".join(_turn_html_complete(t) for t in history)
if in_progress:
out += _turn_html_progress(in_progress[0], in_progress[1])
return out
def _empty_state_html() -> str:
"""No data loaded — polished placeholder for the right panel."""
icon = (
''
)
return (
'
'
f'
{icon}
'
'
No data loaded
'
'
Upload a CSV, JSON, Parquet or Excel file '
'on the left, or click Demo to try sample data.
'
'
'
)
def _ready_state_html() -> str:
"""Shown when data is loaded but no queries asked yet.
Suggestions are derived from the actual loaded table's columns."""
agent = get_agent()
tables = agent.list_tables()
suggestions: list[str] = []
if tables:
schema = agent.executor.get_table_schema(tables[0])
suggestions = _suggest_questions(tables[0], schema)
return (
'
'
'
Ready
'
'
Ask a question above, or try one of these:
'
f'{_suggestions_html(suggestions)}'
'
'
)
# ============================================================ EVENT HANDLERS
def _build_schema_html(table: str) -> str:
agent = get_agent()
schema = agent.executor.get_table_schema(table)
return _schema_preview_html(table, schema)
def _placeholder_for_table(table: str) -> str:
"""Return a smart placeholder based on the loaded table's columns."""
try:
agent = get_agent()
if not agent.list_tables():
return "Ask anything about your data…"
schema = agent.executor.get_table_schema(table)
suggestions = _suggest_questions(table, schema)
if suggestions:
return f"Try: {suggestions[0]}"
except Exception:
pass
return "Ask anything about your data…"
def on_upload(file):
"""Returns: (chip, schema_html, conversation, history, question_update)."""
if file is None:
return "", "", _conversation_html([]), [], gr.update(placeholder="Ask anything about your data…")
agent = get_agent()
agent.reset()
try:
path = Path(file.name if hasattr(file, "name") else file)
table = agent.load_data(path)
rows = agent.executor.con.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0]
cols = len(agent.executor.get_table_schema(table))
chip = _file_chip_html(path.name, rows, cols)
schema = _build_schema_html(table)
placeholder = _placeholder_for_table(table)
return chip, schema, _conversation_html([]), [], gr.update(placeholder=placeholder)
except Exception as e:
logger.exception("upload failed")
return "", "", f'
', [], gr.update()
@spaces.GPU(duration=60)
def _gpu_process(question: str) -> dict:
"""Inference only — models already on cuda from module-level loading."""
agent = get_agent()
return agent.process(question)
def on_ask(question: str, history: list) -> Generator[Tuple[str, str, list], None, None]:
"""
Generator: yields conversation HTML at each pipeline step so the user
sees real-time progress instead of waiting silently.
"""
# Single-shot mode: each Ask is a fresh query, replaces any previous result
question = (question or "").strip()
if not question:
yield _conversation_html([]), "", []
return
if not get_agent().list_tables():
result = {
"question": question,
"error": "Upload a file first or load the demo dataset.",
}
yield _conversation_html([result]), "", [result]
return
# Stream pipeline stages so user sees progress (4 visual steps).
# The actual GPU call is one shot (yielding mid-call would break the
# @spaces.GPU window) but we still surface the stages around it.
yield _conversation_html([], in_progress=(question, "sql")), "", []
try:
result = _gpu_process(question)
except Exception as e:
logger.exception("ask failed")
result = {"question": question, "error": str(e)}
yield _conversation_html([result]), "", [result]
def on_reset():
"""Clear ONLY the dataset (file, chip, schema, agent state).
Also wipes the displayed result since it's now stale without data.
Order: upload, chip_html, schema_html, conversation, history_state
"""
get_agent().reset()
return None, "", "", _conversation_html([]), []
# ====================================================================== APP
def build_app() -> gr.Blocks:
with gr.Blocks(
theme=gr.themes.Base(),
css=THEME_CSS,
title="SQL Agent",
analytics_enabled=False,
) as demo:
# Header
gr.HTML(
'
'
'
'
'
SQL Agent
'
'
Ask anything about your data.
'
'
'
'
'
)
# Two-column layout: LEFT = data/question controls, RIGHT = results
with gr.Row(elem_classes=["split-layout"]):
# ---------- LEFT panel ----------
with gr.Column(elem_classes=["split-left"], scale=0):
gr.HTML('