Spaces:
Sleeping
Sleeping
Create app.py
Browse files
app.py
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import gradio as gr
|
| 2 |
+
from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, Float, insert, text, inspect
|
| 3 |
+
from smolagents import tool, CodeAgent, InferenceClientModel
|
| 4 |
+
|
| 5 |
+
# --- Setup SQLite database (in-memory) ---
|
| 6 |
+
engine = create_engine("sqlite:///:memory:")
|
| 7 |
+
metadata_obj = MetaData()
|
| 8 |
+
|
| 9 |
+
# Create receipts table
|
| 10 |
+
receipts = Table(
|
| 11 |
+
"receipts",
|
| 12 |
+
metadata_obj,
|
| 13 |
+
Column("receipt_id", Integer, primary_key=True),
|
| 14 |
+
Column("customer_name", String(16), primary_key=True),
|
| 15 |
+
Column("price", Float),
|
| 16 |
+
Column("tip", Float),
|
| 17 |
+
)
|
| 18 |
+
metadata_obj.create_all(engine)
|
| 19 |
+
|
| 20 |
+
# Insert sample data
|
| 21 |
+
rows = [
|
| 22 |
+
{"receipt_id": 1, "customer_name": "Alan Payne", "price": 12.06, "tip": 1.20},
|
| 23 |
+
{"receipt_id": 2, "customer_name": "Alex Mason", "price": 23.86, "tip": 0.24},
|
| 24 |
+
{"receipt_id": 3, "customer_name": "Woodrow Wilson", "price": 53.43, "tip": 5.43},
|
| 25 |
+
{"receipt_id": 4, "customer_name": "Margaret James", "price": 21.11, "tip": 1.00},
|
| 26 |
+
]
|
| 27 |
+
for row in rows:
|
| 28 |
+
stmt = insert(receipts).values(**row)
|
| 29 |
+
with engine.begin() as conn:
|
| 30 |
+
conn.execute(stmt)
|
| 31 |
+
|
| 32 |
+
# Create waiters table
|
| 33 |
+
waiters = Table(
|
| 34 |
+
"waiters",
|
| 35 |
+
metadata_obj,
|
| 36 |
+
Column("receipt_id", Integer, primary_key=True),
|
| 37 |
+
Column("waiter_name", String(16), primary_key=True),
|
| 38 |
+
)
|
| 39 |
+
metadata_obj.create_all(engine)
|
| 40 |
+
|
| 41 |
+
rows = [
|
| 42 |
+
{"receipt_id": 1, "waiter_name": "Corey Johnson"},
|
| 43 |
+
{"receipt_id": 2, "waiter_name": "Michael Watts"},
|
| 44 |
+
{"receipt_id": 3, "waiter_name": "Michael Watts"},
|
| 45 |
+
{"receipt_id": 4, "waiter_name": "Margaret James"},
|
| 46 |
+
]
|
| 47 |
+
for row in rows:
|
| 48 |
+
stmt = insert(waiters).values(**row)
|
| 49 |
+
with engine.begin() as conn:
|
| 50 |
+
conn.execute(stmt)
|
| 51 |
+
|
| 52 |
+
# --- Define SQL tool for the agent ---
|
| 53 |
+
@tool
|
| 54 |
+
def sql_engine(query: str) -> str:
|
| 55 |
+
"""
|
| 56 |
+
Executes SQL queries on the available tables: receipts and waiters.
|
| 57 |
+
|
| 58 |
+
Args:
|
| 59 |
+
query: SQL query string.
|
| 60 |
+
"""
|
| 61 |
+
try:
|
| 62 |
+
output = ""
|
| 63 |
+
with engine.connect() as con:
|
| 64 |
+
rows = con.execute(text(query))
|
| 65 |
+
for row in rows:
|
| 66 |
+
output += "\n" + str(row)
|
| 67 |
+
return output or "No results."
|
| 68 |
+
except Exception as e:
|
| 69 |
+
return f"⚠️ SQL Error: {str(e)}"
|
| 70 |
+
|
| 71 |
+
# Dynamically describe tables
|
| 72 |
+
updated_description = "Allows SQL queries on the following tables:\n"
|
| 73 |
+
inspector = inspect(engine)
|
| 74 |
+
for table in ["receipts", "waiters"]:
|
| 75 |
+
columns_info = [(col["name"], col["type"]) for col in inspector.get_columns(table)]
|
| 76 |
+
table_description = f"\n\nTable '{table}':\nColumns:\n" + "\n".join(
|
| 77 |
+
[f" - {name}: {col_type}" for name, col_type in columns_info]
|
| 78 |
+
)
|
| 79 |
+
updated_description += table_description
|
| 80 |
+
sql_engine.description = updated_description
|
| 81 |
+
|
| 82 |
+
# --- Create the agent ---
|
| 83 |
+
agent = CodeAgent(
|
| 84 |
+
tools=[sql_engine],
|
| 85 |
+
model=InferenceClientModel("meta-llama/Meta-Llama-3-8B-Instruct"),
|
| 86 |
+
)
|
| 87 |
+
|
| 88 |
+
# --- Define Gradio interface ---
|
| 89 |
+
def ask_agent(question):
|
| 90 |
+
"""Ask the AI agent a question."""
|
| 91 |
+
result = agent.run(question)
|
| 92 |
+
return result
|
| 93 |
+
|
| 94 |
+
demo = gr.Interface(
|
| 95 |
+
fn=ask_agent,
|
| 96 |
+
inputs=gr.Textbox(label="Ask a question (e.g., Who got the biggest tip?)"),
|
| 97 |
+
outputs=gr.Textbox(label="Agent response"),
|
| 98 |
+
title="🧠 Text-to-SQL Agent",
|
| 99 |
+
description="Ask natural-language questions about a restaurant receipts database. Powered by smolagents + Meta-Llama-3.",
|
| 100 |
+
)
|
| 101 |
+
|
| 102 |
+
demo.launch()
|