Spaces:
Running
Running
File size: 2,659 Bytes
6a0ec6a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
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
@tool
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()
|