Update app.py

app.py

@@ -3,63 +3,64 @@ 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 (persistent file-based) ---
+# --- Setup SQLite database (persistent file) ---
 engine = create_engine("sqlite:///data.db")
 metadata_obj = MetaData()
 
-# --- 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)
+# 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
-    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)
+# 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)
 
-# --- 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)
+# 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(Table("waiters", metadata_obj, autoload_with=engine)).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(waiters).values(**row)
+    with engine.begin() as conn:
+        conn.execute(stmt)
 
 # --- Define SQL tool for the agent ---
 @tool
-def sql_engine(query: str):
+def sql_engine(query: str) -> str:
     """
     Executes SQL queries on the available tables: receipts and waiters.
 
     Args:
         query: SQL query string.
+
+    Returns:
+        A string representing the query results, or an error message.
     """
     try:
         print(f"🧩 Executing query: {query}")  # Debug log in Space console
@@ -68,11 +69,11 @@ def sql_engine(query: str):
             results = [tuple(row) for row in rows]
         if not results:
             return "No results."
-        return results
+        return "\n".join(str(r) for r in results)
     except Exception as e:
         return f"⚠️ SQL Error: {str(e)}"
 
-# --- Describe tables for the model ---
+# Dynamically describe tables
 updated_description = "Allows SQL queries on the following tables:\n"
 inspector = inspect(engine)
 for table in ["receipts", "waiters"]:
@@ -83,25 +84,26 @@ for table in ["receipts", "waiters"]:
     updated_description += table_description
 sql_engine.description = updated_description
 
-# --- Create the AI agent ---
+# --- Create the agent ---
+# Use the HF_TOKEN secret stored in the Space
 agent = CodeAgent(
     tools=[sql_engine],
     model=InferenceClientModel(
         "meta-llama/Meta-Llama-3-8B-Instruct",
-        api_key=os.environ.get("HF_TOKEN"),  # Uses HF secret
+        api_key=os.environ.get("HF_TOKEN")  # 👈 Hugging Face secret
     ),
 )
 
 # --- Define Gradio interface ---
-def ask_agent(question):
-    """Ask the AI agent a question and return the result."""
+def ask_agent(question: str) -> str:
+    """Ask the AI agent a question."""
     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.",
 )