Update app.py

app.py

@@ -1,4 +1,6 @@
+# app.py
 import os
+import re
 import gradio as gr
 from huggingface_hub import InferenceClient
 
@@ -8,59 +10,148 @@ client = InferenceClient(
     token=HF_TOKEN,
 )
 
-SYSTEM_PROMPT = """You are a SQL expert. Given a database schema and a question in English, generate a DuckDB-compatible SQL query that answers the question.
-
-Rules:
-- Return ONLY the SQL query, no explanation, no markdown, no code fences
-- Use exact table and column names from the schema
-- Use DuckDB SQL syntax
-- Add LIMIT 100 unless the user asks for a specific count or all rows
-- Keep queries simple: use the fewest tables and JOINs possible
-- For text search, always use ILIKE '%term%' (case-insensitive) instead of exact match
-- Never use LOWER() or UPPER() for comparison, use ILIKE instead
-- If the question is not related to querying the database (e.g. personal questions, general knowledge, chitchat), respond with exactly: NOT_A_DATA_QUESTION
-- Only generate SQL for questions that can be answered by querying the provided database schema"""
-
-
-def generate_sql(question: str, schema_ddl: str) -> str:
-    prompt = f"Database schema:\n{schema_ddl}\n\nQuestion: {question}\n\nSQL:"
-
-    sql = ""
-    for token in client.chat_completion(
-        messages=[
-            {"role": "system", "content": SYSTEM_PROMPT},
-            {"role": "user", "content": prompt},
-        ],
-        max_tokens=500,
-        temperature=0.1,
-        stream=True,
-    ):
-        chunk = token.choices[0].delta.content or ""
-        sql += chunk
-        yield sql.strip()
-
-    # Clean up final result
-    sql = sql.strip()
-    if sql.startswith("```"):
-        sql = sql.split("\n", 1)[1] if "\n" in sql else sql[3:]
-    if sql.endswith("```"):
-        sql = sql[:-3]
-    sql = sql.strip()
-    if sql.lower().startswith("sql"):
-        sql = sql[3:].strip()
-
-    yield sql
+# ─────────────────────────────────────────────
+# SYSTEM PROMPT  (strict, few-shot)
+# ─────────────────────────────────────────────
+SYSTEM_PROMPT = """You are a strict SQL code generator for DuckDB.
 
+YOUR ONLY JOB is to output a single, valid DuckDB SQL query.
 
+ABSOLUTE OUTPUT RULES — violating any rule makes the output wrong:
+1. Output ONLY raw SQL. No markdown, no code fences, no backticks, no explanations.
+2. Never prefix with "sql", "SQL:", "Here is", or any natural language.
+3. Never output anything after the semicolon.
+4. If the question cannot be answered from the schema, output exactly: NOT_A_DATA_QUESTION
+5. NOT_A_DATA_QUESTION also applies to: greetings, general knowledge, math unrelated to the schema, anything not about querying the provided tables.
+
+SQL RULES:
+- Use ONLY table and column names that appear in the schema — never invent names.
+- Use DuckDB syntax exclusively.
+- Text matching: always use ILIKE '%term%'. Never use LOWER() or UPPER() for comparison.
+- For SELECT queries, default to LIMIT 100 unless the user asks for all rows or a specific count.
+- Prefer the fewest JOINs and subqueries needed to answer the question.
+- Never use SELECT * — always name the columns you need.
+- Age filters: use a numeric comparison on the age column directly (e.g. age > 50).
+- Counts: use COUNT(*) or COUNT(column). Alias it clearly, e.g. AS num_patients.
+- INSERT, UPDATE, DELETE, CREATE, DROP, ALTER are all allowed — the user owns their database.
+
+FEW-SHOT EXAMPLES:
+
+Schema:
+CREATE TABLE patients (patient_id INT, age INT, gender VARCHAR, diagnosis VARCHAR, died BOOLEAN);
+
+Q: How many patients above 50 have asthma?
+A: SELECT COUNT(*) AS num_patients FROM patients WHERE age > 50 AND diagnosis ILIKE '%asthma%';
+
+Q: Show me all patients who died during their hospital stay.
+A: SELECT patient_id, age, gender, diagnosis FROM patients WHERE died = true LIMIT 100;
+
+Q: What is the average age of female patients?
+A: SELECT AVG(age) AS avg_age FROM patients WHERE gender ILIKE '%female%';
+
+Q: Hello, how are you?
+A: NOT_A_DATA_QUESTION
+
+Q: What is the capital of France?
+A: NOT_A_DATA_QUESTION
+
+Now answer the user's question using ONLY the schema they provide."""
+
+# ─────────────────────────────────────────────
+# HELPERS
+# ─────────────────────────────────────────────
+VALID_SQL_STARTS = ("SELECT", "WITH", "INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER")
+
+
+def clean_sql(raw: str) -> str:
+    """Remove markdown fences, leading 'sql' keyword, and extra whitespace."""
+    sql = raw.strip()
+    sql = re.sub(r"^```[a-zA-Z]*\n?", "", sql)
+    sql = re.sub(r"```$", "", sql)
+    sql = re.sub(r"(?i)^sql\s+", "", sql)
+    return sql.strip()
+
+
+def validate_sql(sql: str) -> str:
+    """
+    Light sanity check on the generated SQL.
+    Returns the SQL unchanged if it looks valid, or an error string.
+    """
+    upper = sql.upper().strip()
+
+    if upper == "NOT_A_DATA_QUESTION":
+        return "⚠️ That question doesn't appear to be about the database. Try asking something that can be answered by querying the schema."
+
+    if not upper.startswith(VALID_SQL_STARTS):
+        return (
+            f"⚠️ The model returned an unexpected response instead of SQL:\n\n{sql}\n\n"
+            "Try rephrasing your question to be more specific about the data."
+        )
+
+    return sql  # looks good
+
+
+# ─────────────────────────────────────────────
+# MAIN GENERATOR
+# ─────────────────────────────────────────────
+def generate_sql(question: str, schema_ddl: str):
+    if not question.strip():
+        yield "⚠️ Please enter a question."
+        return
+    if not schema_ddl.strip():
+        yield "⚠️ Please provide your schema DDL."
+        return
+
+    prompt = (
+        f"Database schema:\n{schema_ddl.strip()}\n\n"
+        f"Question: {question.strip()}\n\n"
+        "SQL:"
+    )
+
+    accumulated = ""
+    try:
+        for token in client.chat_completion(
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": prompt},
+            ],
+            max_tokens=500,
+            temperature=0.0,
+            stream=True,
+        ):
+            chunk = token.choices[0].delta.content or ""
+            accumulated += chunk
+            yield accumulated  # stream raw while typing
+
+    except Exception as e:
+        yield f"❌ Error calling model: {e}"
+        return
+
+    # Final: clean then validate
+    final = validate_sql(clean_sql(accumulated))
+    yield final
+
+
+# ─────────────────────────────────────────────
+# GRADIO UI
+# ─────────────────────────────────────────────
 demo = gr.Interface(
     fn=generate_sql,
     inputs=[
-        gr.Textbox(label="Question", placeholder="Show me all patients who died during their hospital stay"),
-        gr.Textbox(label="Schema DDL", lines=10, placeholder="CREATE TABLE patients (...)"),
+        gr.Textbox(
+            label="Question",
+            placeholder="Show me how many patients above 50 have asthma",
+        ),
+        gr.Textbox(
+            label="Schema DDL",
+            lines=10,
+            placeholder="CREATE TABLE patients (...)",
+        ),
     ],
-    outputs=gr.Textbox(label="Generated SQL"),
+    outputs=gr.Textbox(label="Generated SQL", show_copy_button=True),
     title="TinyEHR Text-to-SQL",
     description="Generate SQL queries for the TinyEHR dataset from natural language.",
+    flagging_mode="never",
 )
 
-demo.launch()
+demo.launch()