Update app.py

app.py

@@ -1,83 +1,79 @@
 import os
 import gradio as gr
-from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, Float, insert, text, inspect
+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 ---
-engine = create_engine("sqlite:///students.db")
+# --- Setup SQLite database (persistent in Space) ---
+engine = create_engine("sqlite:///data.db")
 metadata_obj = MetaData()
 
-# --- 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))
+# --- 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),
 )
-
 metadata_obj.create_all(engine)
 
-# --- 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},
-]
-
-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"},
+# 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},
 ]
-
-for row in students_data:
-    stmt = insert(students).values(**row).prefix_with("OR IGNORE")
-    with engine.begin() as conn:
+with engine.begin() as conn:
+    for row in rows:
+        stmt = insert(receipts).values(**row)
         conn.execute(stmt)
 
-for row in marks_data:
-    stmt = insert(marks).values(**row).prefix_with("OR IGNORE")
-    with engine.begin() as conn:
+# --- 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"},
+]
+with engine.begin() as conn:
+    for row in rows:
+        stmt = insert(waiters).values(**row)
         conn.execute(stmt)
 
-# --- Define SQL tool ---
+# --- Define SQL tool for the agent ---
 @tool
-def sql_engine(query: str) -> dict:
+def sql_engine(query: str) -> list:
+    """Executes SQL queries on the available tables: receipts and waiters."""
     try:
+        print(f"🧩 Executing query: {query}")  # Debug log
         with engine.connect() as con:
             rows = con.execute(text(query))
-            columns = rows.keys()
-            data = [tuple(row) for row in rows]
-            return {"columns": columns, "data": data} if data else {"columns": [], "data": []}
+            results = [tuple(row) for row in rows]
+        return results or []
     except Exception as e:
-        return {"columns": ["Error"], "data": [[str(e)]]}
-
+        return [f"⚠️ SQL Error: {str(e)}"]
 
-# --- Table descriptions ---
+# Dynamically describe tables
+updated_description = "Allows SQL queries on the following tables:\n"
 inspector = inspect(engine)
-updated_description = "Ask natural language questions about the student marks database.\n"
-for table in ["students", "marks"]:
+for table in ["receipts", "waiters"]:
     columns_info = [(col["name"], col["type"]) for col in inspector.get_columns(table)]
-    table_description = f"\n\nTable '{table}':\n" + "\n".join(
-        [f"  - {name}: {col_type}" for name, col_type in columns_info]
+    table_description = f"\n\nTable '{table}':\nColumns:\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 ---
@@ -85,60 +81,28 @@ agent = CodeAgent(
     tools=[sql_engine],
     model=InferenceClientModel(
         "meta-llama/Meta-Llama-3-8B-Instruct",
-        api_key=os.environ.get("HF_TOKEN")
+        api_key=os.environ.get("HF_TOKEN")  # Use secret from Space
     ),
 )
 
-# --- Gradio functions ---
-def ask_agent(question: str):
-    result = agent.run(question)
-    if isinstance(result, dict) and "columns" in result and "data" in result:
-        if not result["data"]:
-            return pd.DataFrame(columns=result["columns"])
-        return pd.DataFrame(result["data"], columns=result["columns"])
-    return pd.DataFrame([[str(result)]], columns=["Result"])
-
-# --- 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("## 🎓 Student Marks Database Explorer")
-    
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("### Students Table")
-            gr.Dataframe(students_df, interactive=False)
-        with 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.Dataframe(headers=None)  # Result displayed as DataFrame
-    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)
+# --- Define Gradio interface ---
+def ask_agent(question: str) -> str:
+    """Ask the AI agent a question and return its answer."""
+    try:
+        result = agent.run(question)
+        if isinstance(result, list):
+            result_str = "\n".join(str(r) for r in result)
+            return result_str or "No results."
+        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.",
+)
 
 demo.launch()