Upload 3 files

app.py

@@ -1,15 +1,21 @@
-import os, json, sqlite3, time, textwrap, requests
+import os, sqlite3, textwrap, requests, json, time
 import gradio as gr
-from dotenv import load_dotenv
-load_dotenv()
 
-HF_TOKEN   = os.getenv("HF_TOKEN")  # store in Space Secrets
-MODEL_ID   = "defog/sqlcoder-7b-2"
-API_URL    = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
-HEADERS    = {"Authorization": f"Bearer {HF_TOKEN}"}
-DB_PATH    = "company.db"
+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}"}
 
-def nlp_to_sql(nl_query: str, schema_ddl: str) -> str:
+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)
+
+def nlp_to_sql(nl_query, schema_ddl):
     prompt = textwrap.dedent(f"""
     ### Task
     Translate the following natural language question into ONE valid SQLite SQL query.
@@ -24,48 +30,48 @@ def nlp_to_sql(nl_query: str, schema_ddl: str) -> str:
     """)
     payload = {"inputs": prompt, "parameters": {"max_new_tokens": 256}}
     response = requests.post(API_URL, headers=HEADERS, json=payload, timeout=45)
-    sql = response.json()[0]["generated_text"].split("### SQL")[-1].strip()
-    return sql
+    try:
+        sql = response.json()[0]["generated_text"].split("### SQL")[-1].strip()
+        return sql
+    except:
+        return "Error: " + str(response.json())
 
 def run_pipeline(nl_query):
     trace = []
     start = time.time()
 
-    # 1. load schema
-    with open("schema.sql") as f:
+    create_db_if_needed()
+
+    with open(SCHEMA_FILE) as f:
         schema_ddl = f.read()
-    trace.append(("Load Schema", f"{len(schema_ddl.splitlines())} lines loaded"))
+    trace.append(("Schema Loaded", f"{len(schema_ddl.splitlines())} lines"))
 
-    # 2. NL ➜ SQL
     sql_query = nlp_to_sql(nl_query, schema_ddl)
-    trace.append(("LLM Output", sql_query))
+    trace.append(("LLM Generated SQL", sql_query))
 
-    # 3. Execute SQL
     try:
         with sqlite3.connect(DB_PATH) as conn:
             cursor = conn.execute(sql_query)
             rows = cursor.fetchall()
-            colnames = [d[0] for d in cursor.description] if cursor.description else []
+            colnames = [desc[0] for desc in cursor.description] if cursor.description else []
         result = {"columns": colnames, "rows": rows}
-        trace.append(("Execution", f"{len(rows)} rows returned"))
+        trace.append(("SQL Executed", f"{len(rows)} rows returned"))
     except Exception as e:
         result = {"error": str(e)}
         trace.append(("Execution Error", str(e)))
 
-    # 4. timing
-    trace.append(("Latency", f"{time.time() - start:0.2f}s"))
-
-    return sql_query, json.dumps(result, indent=2), "\n".join([f"{s}: {m}" for s, m in trace])
+    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)
 
-with gr.Blocks(title="NLP ➜ SQL Demo") as demo:
-    gr.Markdown("### NLP to SQL Query • SQLite • Trace Visibility")
-    nl_input  = gr.Textbox(label="Natural-Language Question", placeholder="e.g. List employees in Engineering hired after 2021")
+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")
     with gr.Row():
-        sql_out   = gr.Code(label="Generated SQL")
-        result_out= gr.Code(label="Query Result")
-    trace_out = gr.Textbox(label="Trace", lines=6)
-    submit = gr.Button("Run")
-    submit.click(fn=run_pipeline, inputs=nl_input, outputs=[sql_out, result_out, trace_out])
+        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])
 
 if __name__ == "__main__":
     demo.launch()