Update app.py

app.py

@@ -3,56 +3,58 @@ import gradio as gr
 from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, Float, insert, text, inspect
 from smolagents import tool, CodeAgent, InferenceClientModel
 
-# --- Setup SQLite database (in-memory) ---
+# --- Setup SQLite database (persistent file-based) ---
 engine = create_engine("sqlite:///data.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),
-)
-metadata_obj.create_all(engine)
+# --- Create receipts table (if not exists) ---
+if not engine.dialect.has_table(engine.connect(), "receipts"):
+    receipts = Table(
+        "receipts",
+        metadata_obj,
+        Column("receipt_id", Integer, primary_key=True),
+        Column("customer_name", String(16)),
+        Column("price", Float),
+        Column("tip", Float),
+    )
+    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},
-]
-for row in rows:
-    stmt = insert(receipts).values(**row)
-    with engine.begin() as conn:
-        conn.execute(stmt)
+    # 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 rows:
+        stmt = insert(Table("receipts", metadata_obj, autoload_with=engine)).values(**row)
+        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)
+# --- Create waiters table (if not exists) ---
+if not engine.dialect.has_table(engine.connect(), "waiters"):
+    waiters = Table(
+        "waiters",
+        metadata_obj,
+        Column("receipt_id", Integer, primary_key=True),
+        Column("waiter_name", String(16)),
+    )
+    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)
-    with engine.begin() as conn:
-        conn.execute(stmt)
+    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(Table("waiters", metadata_obj, autoload_with=engine)).values(**row)
+        with engine.begin() as conn:
+            conn.execute(stmt)
 
 # --- Define SQL tool for the agent ---
 @tool
-def sql_engine(query: str) -> str:
+def sql_engine(query: str):
     """
     Executes SQL queries on the available tables: receipts and waiters.
 
@@ -60,16 +62,17 @@ def sql_engine(query: str) -> str:
         query: SQL query string.
     """
     try:
-        output = ""
+        print(f"🧩 Executing query: {query}")  # Debug log in Space console
         with engine.connect() as con:
             rows = con.execute(text(query))
-            for row in rows:
-                output += "\n" + str(row)
-        return output or "No results."
+            results = [tuple(row) for row in rows]
+        if not results:
+            return "No results."
+        return results
     except Exception as e:
         return f"⚠️ SQL Error: {str(e)}"
 
-# Dynamically describe tables
+# --- Describe tables for the model ---
 updated_description = "Allows SQL queries on the following tables:\n"
 inspector = inspect(engine)
 for table in ["receipts", "waiters"]:
@@ -80,26 +83,25 @@ for table in ["receipts", "waiters"]:
     updated_description += table_description
 sql_engine.description = updated_description
 
-# --- Create the agent ---
-# Use the HF_TOKEN secret stored in the Space
+# --- Create the AI agent ---
 agent = CodeAgent(
     tools=[sql_engine],
     model=InferenceClientModel(
         "meta-llama/Meta-Llama-3-8B-Instruct",
-        api_key=os.environ.get("HF_TOKEN")  # 👈 Use secret key
+        api_key=os.environ.get("HF_TOKEN"),  # Uses HF secret
     ),
 )
 
 # --- Define Gradio interface ---
 def ask_agent(question):
-    """Ask the AI agent a question."""
+    """Ask the AI agent a question and return the result."""
     result = agent.run(question)
     return result
 
 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"),
+    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.",
 )