|
import os |
|
from dotenv import load_dotenv |
|
import gradio as gr |
|
from sqlalchemy import text |
|
from smolagents import tool, CodeAgent, HfApiModel |
|
|
|
|
|
from database import engine, receipts |
|
|
|
|
|
load_dotenv(override=True) |
|
hf_token = os.getenv("HF_TOKEN") |
|
|
|
|
|
@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() |
|
|
|
|
|
agent = CodeAgent( |
|
tools=[sql_engine], |
|
model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct", token=hf_token), |
|
) |
|
|
|
|
|
def query_sql(user_query): |
|
return sql_engine(user_query) |
|
|
|
|
|
demo = 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", |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
demo.launch(server_name="0.0.0.0", server_port=7860, share=True) |
|
|