Update app.py

app.py

@@ -1,4 +1,3 @@
-# app.py
 import os
 import re
 import gradio as gr
@@ -13,8 +12,8 @@ client = InferenceClient(
 # ─────────────────────────────────────────────
 # SYSTEM PROMPT  (strict, few-shot)
 # ─────────────────────────────────────────────
-SYSTEM_PROMPT = """You are a strict SQL code generator for DuckDB.
-
+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:
@@ -26,18 +25,23 @@ ABSOLUTE OUTPUT RULES — violating any rule makes the output wrong:
 
 SQL RULES:
 - Use ONLY table and column names that appear in the schema — never invent names.
-- Use DuckDB syntax exclusively.
+- Use DuckDB syntax exclusively. Never use SQLite or MySQL syntax.
 - 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.
-...
+- Date arithmetic: NEVER use julianday(). Use datediff('day', start_col, end_col) for days between two timestamps. Use epoch(end_col - start_col) / 86400 for interval-to-days.
+- Identifier quoting: wrap table and column names in double quotes if they start with a digit or contain special characters (e.g. "2b_concept", "my-column").
+- String concatenation: use || operator, never CONCAT().
+- Current date/time: use current_date or current_timestamp, never NOW().
+
 FEW-SHOT EXAMPLES:
 
 Schema:
 CREATE TABLE patients (patient_id INT, age INT, gender VARCHAR, diagnosis VARCHAR, died BOOLEAN);
+CREATE TABLE admissions (subject_id INT, admittime TIMESTAMP, dischtime TIMESTAMP, admission_type VARCHAR);
 
 Q: How many patients above 50 have asthma?
 A: SELECT COUNT(*) AS num_patients FROM patients WHERE age > 50 AND diagnosis ILIKE '%asthma%';
@@ -48,6 +52,9 @@ A: SELECT patient_id, age, gender, diagnosis FROM patients WHERE died = true LIM
 Q: What is the average age of female patients?
 A: SELECT AVG(age) AS avg_age FROM patients WHERE gender ILIKE '%female%';
 
+Q: Who are the top 10 patients with the longest hospital stay?
+A: SELECT a.subject_id, datediff('day', a.admittime, a.dischtime) AS stay_days FROM admissions a WHERE a.dischtime IS NOT NULL ORDER BY stay_days DESC LIMIT 10;
+
 Q: Hello, how are you?
 A: NOT_A_DATA_QUESTION
 
@@ -63,11 +70,19 @@ VALID_SQL_STARTS = ("SELECT", "WITH", "INSERT", "UPDATE", "DELETE", "CREATE", "D
 
 
 def clean_sql(raw: str) -> str:
-    """Remove markdown fences, leading 'sql' keyword, and extra whitespace."""
+    """Remove markdown fences, leading 'sql' keyword, thinking tags, and extra whitespace."""
     sql = raw.strip()
+    # Strip <think>...</think> blocks (some models emit these)
+    sql = re.sub(r"<think>.*?</think>", "", sql, flags=re.DOTALL)
+    # Strip markdown code fences
     sql = re.sub(r"^```[a-zA-Z]*\n?", "", sql)
     sql = re.sub(r"```$", "", sql)
+    # Strip leading "sql" keyword
     sql = re.sub(r"(?i)^sql\s+", "", sql)
+    # Strip any trailing text after the semicolon
+    semi_match = re.search(r";", sql)
+    if semi_match:
+        sql = sql[: semi_match.end()]
     return sql.strip()
 
 
@@ -79,7 +94,10 @@ def validate_sql(sql: str) -> str:
     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."
+        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 (