Update app.py

app.py

@@ -2,84 +2,79 @@ import os
 import gradio as gr
 from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, Float, insert, text, inspect
 from smolagents import tool, CodeAgent, InferenceClientModel
+import pandas as pd
 
-# --- Setup SQLite database (persistent in Space) ---
-engine = create_engine("sqlite:///data.db")
+# --- Setup SQLite database ---
+engine = create_engine("sqlite:///students.db")
 metadata_obj = MetaData()
 
-# --- Create receipts table ---
-receipts = Table(
-    "receipts",
-    metadata_obj,
-    Column("receipt_id", Integer, primary_key=True),
-    Column("customer_name", String(16), primary_key=True),
-    Column("price", Float),
-    Column("tip", Float),
+# --- Create tables ---
+students = Table(
+    "students", metadata_obj,
+    Column("student_id", Integer, primary_key=True),
+    Column("name", String(30)),
+    Column("age", Integer)
 )
+
+marks = Table(
+    "marks", metadata_obj,
+    Column("student_id", Integer, primary_key=True),
+    Column("subject", String(20), primary_key=True),
+    Column("marks", Float),
+    Column("grade", String(2))
+)
+
 metadata_obj.create_all(engine)
 
-# Insert sample data
-rows = [
-    {"receipt_id": 1, "customer_name": "Alan Payne", "price": 12.06, "tip": 1.20},
-    {"receipt_id": 2, "customer_name": "Alex Mason", "price": 23.86, "tip": 0.24},
-    {"receipt_id": 3, "customer_name": "Woodrow Wilson", "price": 53.43, "tip": 5.43},
-    {"receipt_id": 4, "customer_name": "Margaret James", "price": 21.11, "tip": 1.00},
+# --- Insert sample data safely ---
+students_data = [
+    {"student_id": 1, "name": "Alice Johnson", "age": 16},
+    {"student_id": 2, "name": "Bob Smith", "age": 17},
+    {"student_id": 3, "name": "Charlie Lee", "age": 16},
+    {"student_id": 4, "name": "Diana King", "age": 17},
 ]
-for row in rows:
-    stmt = insert(receipts).values(**row)
+
+marks_data = [
+    {"student_id": 1, "subject": "Math", "marks": 88, "grade": "A"},
+    {"student_id": 1, "subject": "English", "marks": 75, "grade": "B"},
+    {"student_id": 2, "subject": "Math", "marks": 92, "grade": "A"},
+    {"student_id": 2, "subject": "English", "marks": 81, "grade": "A"},
+    {"student_id": 3, "subject": "Math", "marks": 64, "grade": "C"},
+    {"student_id": 3, "subject": "English", "marks": 70, "grade": "B"},
+    {"student_id": 4, "subject": "Math", "marks": 79, "grade": "B"},
+    {"student_id": 4, "subject": "English", "marks": 85, "grade": "A"},
+]
+
+for row in students_data:
+    stmt = insert(students).values(**row).prefix_with("OR IGNORE")
     with engine.begin() as conn:
         conn.execute(stmt)
 
-# --- Create waiters table ---
-waiters = Table(
-    "waiters",
-    metadata_obj,
-    Column("receipt_id", Integer, primary_key=True),
-    Column("waiter_name", String(16), primary_key=True),
-)
-metadata_obj.create_all(engine)
-
-rows = [
-    {"receipt_id": 1, "waiter_name": "Corey Johnson"},
-    {"receipt_id": 2, "waiter_name": "Michael Watts"},
-    {"receipt_id": 3, "waiter_name": "Michael Watts"},
-    {"receipt_id": 4, "waiter_name": "Margaret James"},
-]
-for row in rows:
-    stmt = insert(waiters).values(**row)
+for row in marks_data:
+    stmt = insert(marks).values(**row).prefix_with("OR IGNORE")
     with engine.begin() as conn:
         conn.execute(stmt)
 
-# --- Define SQL tool for the agent ---
+# --- Define SQL tool ---
 @tool
 def sql_engine(query: str) -> list:
-    """
-    Executes SQL queries on the available tables: receipts and waiters.
-
-    Args:
-        query: SQL query string.
-
-    Returns:
-        List of tuples with query results.
-    """
     try:
-        print(f"🧩 Executing query: {query}")  # Debug log
         with engine.connect() as con:
             rows = con.execute(text(query))
-            results = [tuple(row) for row in rows]  # Keep results as tuples
-        return results or []
+            return [tuple(row) for row in rows] or []
     except Exception as e:
         return [f"⚠️ SQL Error: {str(e)}"]
 
-# Dynamically describe tables
-updated_description = "Allows SQL queries on the following tables:\n"
+# Table descriptions
 inspector = inspect(engine)
-for table in ["receipts", "waiters"]:
+updated_description = "Ask natural language questions about the student marks database.\n"
+for table in ["students", "marks"]:
     columns_info = [(col["name"], col["type"]) for col in inspector.get_columns(table)]
-    table_description = f"\n\nTable '{table}':\nColumns:\n" + "\n".join(
+    table_description = f"\n\nTable '{table}':\n" + "\n".join(
         [f"  - {name}: {col_type}" for name, col_type in columns_info]
     )
     updated_description += table_description
+
 sql_engine.description = updated_description
 
 # --- Create the agent ---
@@ -87,29 +82,61 @@ agent = CodeAgent(
     tools=[sql_engine],
     model=InferenceClientModel(
         "meta-llama/Meta-Llama-3-8B-Instruct",
-        api_key=os.environ.get("HF_TOKEN")  # Use secret from Space
+        api_key=os.environ.get("HF_TOKEN")
     ),
 )
 
-# --- Define Gradio interface ---
+# --- Gradio functions ---
 def ask_agent(question: str) -> str:
-    """Ask the AI agent a question and return its answer."""
     try:
         result = agent.run(question)
-        # Convert tuples to readable string if result is a list of tuples
         if isinstance(result, list):
-            result_str = "\n".join(str(r) for r in result)
-            return result_str or "No results."
+            df = pd.DataFrame(result)
+            return df.to_string(index=False)
         return str(result)
     except Exception as e:
         return f"⚠️ Error: {str(e)}"
 
-demo = gr.Interface(
-    fn=ask_agent,
-    inputs=gr.Textbox(label="Ask a question (e.g., Who got the biggest tip?)"),
-    outputs=gr.Textbox(label="Agent response"),
-    title="🧠 Text-to-SQL Agent",
-    description="Ask natural-language questions about a restaurant receipts database. Powered by smolagents + Meta-Llama-3.",
-)
+# Sample prompts
+sample_prompts = [
+    "Which student has the highest marks in Math?",
+    "List all students with their subjects and grades.",
+    "What is the average marks of each student?",
+    "Show students who scored above 80 in English."
+]
+
+def set_prompt(prompt):
+    return prompt
+
+# Display database tables
+def get_tables():
+    with engine.connect() as con:
+        students_df = pd.read_sql_table("students", con)
+        marks_df = pd.read_sql_table("marks", con)
+    return students_df, marks_df
+
+students_df, marks_df = get_tables()
+
+# --- Gradio App ---
+with gr.Blocks() as demo:
+    gr.Markdown("## Text-TO-SQL Agent")
+    with gr.Row():
+        gr.Column():
+            gr.Markdown("### Students Table")
+            gr.Dataframe(students_df, interactive=False)
+        gr.Column():
+            gr.Markdown("### Marks Table")
+            gr.Dataframe(marks_df, interactive=False)
+
+    gr.Markdown("### Ask the AI Agent")
+    question = gr.Textbox(label="Ask a question")
+    output = gr.Textbox(label="Agent response")
+    ask_btn = gr.Button("Ask")
+    ask_btn.click(ask_agent, inputs=question, outputs=output)
+
+    gr.Markdown("### Sample Prompts")
+    for sp in sample_prompts:
+        btn = gr.Button(sp)
+        btn.click(set_prompt, inputs=gr.State(sp), outputs=question)
 
 demo.launch()