app.py

import os
import gradio as gr
from sqlalchemy import text
from smolagents import tool, CodeAgent, HfApiModel
import spaces

# Import the persistent database
from database import engine, receipts

@tool
def sql_engine(query: str) -> str:
    """
    Executes an SQL query on the 'receipts' table and returns formatted results.
    
    Args:
        query: The SQL query to execute.

    Returns:
        Query result as a formatted string.
    """
    try:
        with engine.connect() as con:
            rows = con.execute(text(query)).fetchall()

        if not rows:
            return "No results found."

        # Convert results into a readable string format
        return "\n".join([", ".join(map(str, row)) for row in rows])

    except Exception as e:
        return f"Error: {str(e)}"

@tool
def query_sql(user_query: str) -> str:
    """
    Converts natural language input to an SQL query using CodeAgent
    and returns the execution results.

    Args:
        user_query: The user's request in natural language.

    Returns:
        The query result from the database as a formatted string.
    """
    # Generate SQL from natural language
    generated_sql = agent.run(f"Convert this request into SQL: {user_query}")

    # Log the generated SQL for debugging
    print(f"Generated SQL: {generated_sql}")

    # Ensure we only execute valid SQL queries
    if not generated_sql.lower().startswith(("select", "show", "pragma")):
        return "Error: Only SELECT queries are allowed."

    # Execute the SQL query and return the result
    result = sql_engine(generated_sql)

    # Log the result for debugging
    print(f"SQL Query Result: {result}")

    # If the result is empty, return a friendly message
    if not result.strip():
        return "No results found."

    return result  # Return clean result, avoiding passing raw SQL back

# Initialize CodeAgent to generate SQL queries from natural language
agent = CodeAgent(
    tools=[sql_engine],  # Ensure sql_engine is properly registered
    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
)

# Define Gradio interface
demo = gr.Interface(
    fn=query_sql,
    inputs=gr.Textbox(label="Enter your query in plain English"),
    outputs=gr.Textbox(label="Query Result"),
    title="Natural Language to SQL Executor",
    description="Enter a plain English request, and the AI will generate an SQL query and return the results.",
    flagging_mode="never",
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)