Upload 3 files

app.py

@@ -1,77 +1,133 @@
-import os, sqlite3, textwrap, requests, json, time
+import os
+import json
+import time
+import sqlite3
+import textwrap
+import requests
 import gradio as gr
 
-DB_PATH = "company.db"
+# -----------------------------
+# Config
+# -----------------------------
+DB_PATH     = "company.db"
 SCHEMA_FILE = "schema.sql"
-MODEL_ID = "defog/sqlcoder-7b-2"
-HF_TOKEN = os.getenv("HF_TOKEN")
-API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
-HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
+MODEL_ID    = "defog/sqlcoder-7b-2"           # swap to another HF model if desired
+HF_TOKEN    = os.getenv("HF_TOKEN")           # set in Space > Settings > Secrets
+API_URL     = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
+HEADERS     = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
 
+# -----------------------------
+# Helper: Build DB once
+# -----------------------------
 def create_db_if_needed():
-    if not os.path.exists(DB_PATH):
-        with open(SCHEMA_FILE) as f:
-            schema = f.read()
-        with sqlite3.connect(DB_PATH) as conn:
-            conn.executescript(schema)
+    if os.path.exists(DB_PATH):
+        return
+    if not os.path.isfile(SCHEMA_FILE):
+        raise FileNotFoundError("schema.sql not found in Space.")
+    with open(SCHEMA_FILE) as f:
+        schema_sql = f.read()
+    with sqlite3.connect(DB_PATH) as conn:
+        conn.executescript(schema_sql)
 
-def nlp_to_sql(nl_query, schema_ddl):
+# -----------------------------
+# Helper: Call HF Inference API
+# -----------------------------
+def nlp_to_sql(nl_query: str, schema_ddl: str) -> str:
     prompt = textwrap.dedent(f"""
-    ### Task
-    Translate the following natural language question into ONE valid SQLite SQL query.
+        ### Task
+        Translate the following natural language question into ONE valid SQLite SQL query.
 
-    ### Schema
-    {schema_ddl}
+        ### Schema
+        {schema_ddl}
 
-    ### Question
-    {nl_query}
+        ### Question
+        {nl_query}
 
-    ### SQL
+        ### SQL
     """)
-    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 256}}
-    response = requests.post(API_URL, headers=HEADERS, json=payload, timeout=45)
+    payload = {
+        "inputs": prompt,
+        "parameters": {"max_new_tokens": 256}
+    }
+
+    try:
+        resp = requests.post(API_URL, headers=HEADERS, json=payload, timeout=60)
+    except Exception as e:
+        return f"✖ Exception contacting API: {e}"
+
+    if resp.status_code != 200:
+        return f"✖ API error {resp.status_code}: {resp.text}"
+
     try:
-        sql = response.json()[0]["generated_text"].split("### SQL")[-1].strip()
-        return sql
-    except:
-        return "Error: " + str(response.json())
+        output = resp.json()
+    except ValueError:
+        return "✖ Non-JSON response from API."
 
-def run_pipeline(nl_query):
+    if not isinstance(output, list) or not output:
+        return "✖ Model returned empty output."
+
+    generated = output[0].get("generated_text", "")
+    if "### SQL" in generated:
+        sql = generated.split("### SQL")[-1].strip()
+    else:
+        sql = generated.strip()
+
+    return sql or "✖ Empty SQL string returned."
+
+# -----------------------------
+# Pipeline: NL → SQL → Execute
+# -----------------------------
+def run_pipeline(nl_query: str):
     trace = []
-    start = time.time()
+    t0 = time.time()
 
+    # build db if needed
     create_db_if_needed()
 
+    # load schema for prompt
     with open(SCHEMA_FILE) as f:
         schema_ddl = f.read()
-    trace.append(("Schema Loaded", f"{len(schema_ddl.splitlines())} lines"))
+    trace.append(("Schema loaded", f"{len(schema_ddl.splitlines())} lines"))
 
+    # convert NL → SQL
     sql_query = nlp_to_sql(nl_query, schema_ddl)
-    trace.append(("LLM Generated SQL", sql_query))
+    trace.append(("LLM output", sql_query))
 
+    # execute SQL
     try:
         with sqlite3.connect(DB_PATH) as conn:
-            cursor = conn.execute(sql_query)
-            rows = cursor.fetchall()
-            colnames = [desc[0] for desc in cursor.description] if cursor.description else []
-        result = {"columns": colnames, "rows": rows}
-        trace.append(("SQL Executed", f"{len(rows)} rows returned"))
+            cur = conn.execute(sql_query)
+            rows = cur.fetchall()
+            cols = [d[0] for d in cur.description] if cur.description else []
+        result = {"columns": cols, "rows": rows}
+        trace.append(("Execution", f"{len(rows)} rows"))
     except Exception as e:
         result = {"error": str(e)}
-        trace.append(("Execution Error", str(e)))
+        trace.append(("Execution error", str(e)))
 
-    trace.append(("Total Time", f"{time.time() - start:.2f}s"))
-    return sql_query, json.dumps(result, indent=2), "\n".join(f"{step}: {msg}" for step, msg in trace)
+    trace.append(("Total time", f"{time.time() - t0:.2f}s"))
 
-with gr.Blocks() as demo:
-    gr.Markdown("### NLP to SQL Query (Auto-Build SQLite, No Terminal Needed)")
-    input_box = gr.Textbox(label="Enter natural language query", placeholder="e.g., Show all employees in Engineering")
+    # stringify outputs for UI
+    result_json = json.dumps(result, indent=2, ensure_ascii=False)
+    trace_log   = "\n".join(f"{s}: {m}" for s, m in trace)
+
+    return sql_query, result_json, trace_log
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
+with gr.Blocks(title="NLP ➜ SQL (SQLite, HF Hub)") as demo:
+    gr.Markdown("## NLP → SQL Query Demo (auto-build SQLite, no terminal required)")
+    nl_in  = gr.Textbox(label="Natural-language question",
+                        placeholder="e.g. List all employees in Engineering hired after 2021")
     with gr.Row():
-        sql_box = gr.Code(label="Generated SQL")
-        result_box = gr.Code(label="Query Output")
-    trace_box = gr.Textbox(label="Trace Log", lines=6)
-    run_btn = gr.Button("Run Query")
-    run_btn.click(fn=run_pipeline, inputs=input_box, outputs=[sql_box, result_box, trace_box])
+        sql_out    = gr.Code(label="Generated SQL", language="sql")
+        result_out = gr.Code(label="Query result (JSON)")
+    trace_out = gr.Textbox(label="Trace", lines=6)
+
+    run_btn = gr.Button("Run")
+    run_btn.click(fn=run_pipeline, inputs=nl_in,
+                  outputs=[sql_out, result_out, trace_out])
 
 if __name__ == "__main__":
     demo.launch()