Update app.py

app.py

@@ -1,20 +1,44 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load a text-to-SQL model. Can be swapped out for better ones later.
+# Load text-to-SQL model
 generator = pipeline("text2text-generation", model="mrm8488/t5-base-finetuned-wikiSQL")
 
-def convert_to_sql(question):
-    prompt = f"translate English to SQL: {question}"
-    result = generator(prompt, max_length=256)[0]["generated_text"]
-    return result
+# Function to process full JSON input
+def convert_to_sql(payload):
+    question = payload.get("question", "")
+    schema = payload.get("schema", "")
+    sample_rows = payload.get("sample_rows", [])
 
+    # Turn sample data into a preview string
+    preview_rows = "\n".join([str(row) for row in sample_rows[:5]])
+
+    # Build enhanced prompt
+    prompt = f"""
+You are an AI that writes DuckDB-compatible SQL based on user questions.
+
+Table Schema:
+{schema}
+
+Sample Data:
+{preview_rows}
+
+Question:
+{question}
+
+Write only the SQL query — no explanations.
+""".strip()
+
+    result = generator(prompt, max_length=256, clean_up_tokenization_spaces=True)[0]["generated_text"]
+    return result.strip()
+
+# Gradio interface for testing (takes JSON input)
 demo = gr.Interface(
     fn=convert_to_sql,
-    inputs=gr.Textbox(lines=2, placeholder="Ask a question about your data..."),
+    inputs=gr.JSON(label="question + schema + sample_rows"),
     outputs="text",
-    title="Text-to-SQL Generator",
-    description="Enter a natural language question and get the SQL query for a table named 'sales_data'."
+    title="Smarter Text-to-SQL Generator",
+    description="Send in natural question, schema, and data — get back accurate SQL."
 )
 
 demo.launch()