Update app.py

app.py

@@ -1,41 +1,57 @@
+from dotenv import load_dotenv
+import os
 import gradio as gr
-from transformers import T5Tokenizer, T5ForConditionalGeneration
-
-# Load FLAN-T5-small
-model_name = "google/flan-t5-small"
-tokenizer = T5Tokenizer.from_pretrained(model_name)
-model = T5ForConditionalGeneration.from_pretrained(model_name)
-
-def generate_sql(payload):
-    question = payload.get("question", "")
-    schema = payload.get("schema", "")
-    sample_rows = payload.get("sample_rows", [])
-
-    # Convert sample rows to a flat string
-    rows_text = " ".join([str(row) for row in sample_rows]) if sample_rows else ""
-
-    # Construct the prompt
-    prompt = (
-        f"You are a SQL expert.\n"
-        f"Schema: {schema}\n"
-        f"Sample Rows: {rows_text}\n"
-        f"Question: {question}\n"
-        f"Generate SQL:"
+from groq import Groq
+
+load_dotenv()
+api = os.getenv("groq_api_key")
+
+def create_prompt(user_query, table_metadata):
+    system_prompt = """
+You are a SQL query generator specialized in generating SQL queries for a single table at a time.
+Your task is to accurately convert natural language queries into SQL statements based on the user's intent and the provided table metadata.
+
+Rules:
+- Single Table Only: Use only the table in the metadata.
+- Metadata-Based Validation: Use only columns in the metadata.
+- User Intent: Support filters, grouping, sorting, etc.
+- SQL Syntax: Use standard SQL (DuckDB compatible).
+- Output only valid SQL. No extra commentary.
+
+Input:
+User Query: {user_query}
+Table Metadata: {table_metadata}
+
+Output:
+SQL Query (on a single line, nothing else).
+    """
+    return system_prompt.strip(), f"User Query: {user_query}\nTable Metadata: {table_metadata}"
+
+def generate_output(system_prompt, user_prompt):
+    client = Groq(api_key=api)
+    chat_completion = client.chat.completions.create(
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt}
+        ],
+        model="llama3-70b-8192"
     )
+    response = chat_completion.choices[0].message.content.strip()
+    return response if response.lower().startswith("select") else "Can't perform the task at the moment."
 
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-    outputs = model.generate(input_ids, max_length=256, temperature=0.7)
-    sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    return sql
+# NEW: accepts user_query and dynamic table_metadata string
+def response(payload):
+    user_query = payload.get("question", "")
+    table_metadata = payload.get("schema", "")
+    system_prompt, user_prompt = create_prompt(user_query, table_metadata)
+    return generate_output(system_prompt, user_prompt)
 
 demo = gr.Interface(
-    fn=generate_sql,
-    inputs=gr.JSON(label="JSON (question, schema, sample_rows)"),
+    fn=response,
+    inputs=gr.JSON(label="Input JSON (question, schema)"),
     outputs="text",
-    title="FLAN-T5 Text-to-SQL",
-    description="Using FLAN-T5 to generate SQL from natural language and tabular schema.",
-    allow_flagging="never"
+    title="SQL Generator (Groq + LLaMA3)",
+    description="Input: question & table metadata. Output: SQL using dynamic schema."
 )
 
 demo.launch()