Add schema visualizer (table cards + PK/FK badges + relationships); guard prefetch to Space only

app.py

@@ -8,6 +8,7 @@ read-only against the chosen database.
 """
 from __future__ import annotations
 
+import html
 import sqlite3
 import traceback
 import tempfile
@@ -131,7 +132,7 @@ def resolve_db(uploaded_db: str | None) -> tuple[str | None, str | None]:
 
 
 def show_schema(uploaded_db: str | None) -> str:
-    """Preview the schema of whichever database is currently selected."""
+    """Preview the raw DDL of whichever database is currently selected."""
     db_path, err = resolve_db(uploaded_db)
     if err:
         return f"-- {err} --"
@@ -141,6 +142,87 @@ def show_schema(uploaded_db: str | None) -> str:
         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 prefetch_weights() -> None:
     """Download weights to disk at startup (CPU, no GPU timer).
 
@@ -218,6 +300,9 @@ 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.Row():
         with gr.Column(scale=1):
             uploaded_db = gr.File(
@@ -253,7 +338,7 @@ with gr.Blocks(title="LocalSQL") as demo:
             status = gr.Textbox(label="Status")
             with gr.Accordion("Model reasoning (chain of thought)", open=False):
                 reasoning = gr.Markdown()
-            with gr.Accordion("Schema", open=False):
+            with gr.Accordion("Schema (raw SQL)", open=False):
                 schema = gr.Code(language="sql")
 
     submit.click(
@@ -261,14 +346,18 @@ with gr.Blocks(title="LocalSQL") as demo:
         inputs=[question, evidence, run_sql, uploaded_db],
         outputs=[sql, rows, status, schema, reasoning],
     )
-    # Preview the schema of whichever database is selected.
-    uploaded_db.change(show_schema, inputs=[uploaded_db], outputs=[schema])
-    demo.load(show_schema, inputs=[uploaded_db], outputs=[schema])
+    # Refresh the schema map + raw DDL for whichever database is selected.
+    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])
 
 
-# Run at import so the Space downloads weights during startup, not on first query.
 ensure_sample_db()
-prefetch_weights()
+# 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()