Spaces:
Running
Running
import os | |
from dotenv import load_dotenv | |
import gradio as gr | |
from sqlalchemy import ( | |
create_engine, | |
MetaData, | |
Table, | |
Column, | |
String, | |
Integer, | |
Float, | |
insert, | |
text, | |
) | |
from smolagents import tool, CodeAgent, HfApiModel | |
# Load Hugging Face token from environment variables | |
load_dotenv(override=True) | |
hf_token = os.getenv("HF_TOKEN") | |
# Initialize in-memory SQLite database | |
engine = create_engine("sqlite:///:memory:") | |
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) | |
# Function to insert data | |
def insert_rows_into_table(rows, table): | |
with engine.begin() as connection: | |
connection.execute(insert(table), rows) | |
# 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_rows_into_table(rows, receipts) | |
# SQL Execution function | |
def sql_engine(query: str) -> str: | |
""" | |
Allows you to perform SQL queries on the table. Returns a string representation of the result. | |
The table is named 'receipts'. Its description is as follows: | |
Columns: | |
- receipt_id: INTEGER | |
- customer_name: VARCHAR(16) | |
- price: FLOAT | |
- tip: FLOAT | |
Args: | |
query: The query to perform. This should be correct SQL. | |
""" | |
output = "" | |
try: | |
with engine.connect() as con: | |
rows = con.execute(text(query)) | |
for row in rows: | |
output += "\n" + str(row) | |
except Exception as e: | |
output = f"Error: {str(e)}" | |
return output.strip() | |
# Set up the Hugging Face agent | |
agent = CodeAgent( | |
tools=[sql_engine], | |
model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct", token=hf_token), | |
) | |
# Gradio function | |
def query_sql(user_query): | |
return sql_engine(user_query) | |
# Define Gradio interface | |
iface = gr.Interface( | |
fn=query_sql, | |
inputs=gr.Textbox(label="Enter your SQL Query"), | |
outputs=gr.Textbox(label="Query Result"), | |
title="SQL Query Executor", | |
description="Enter SQL queries to interact with an in-memory SQLite database.", | |
allow_flagging="never", | |
) | |
# Launch the app | |
if __name__ == "__main__": | |
iface.launch() | |