sqlAgent / app.py
ZennyKenny's picture
remove iface instantiation
2d636d1 verified
raw
history blame
1.72 kB
import os
from dotenv import load_dotenv
import gradio as gr
from sqlalchemy import text
from smolagents import tool, CodeAgent, HfApiModel
# Import the database setup from database.py
from database import engine, receipts
# Load Hugging Face token from environment variables
load_dotenv(override=True)
hf_token = os.getenv("HF_TOKEN")
# 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.",
flagging_mode="never",
)
# Launch the app
if __name__ == "__main__":
launch(server_name="0.0.0.0", server_port=7860, share=True)