Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import gradio as gr
-from sqlalchemy import text, inspect, create_engine
+from sqlalchemy import text, inspect
 from smolagents import tool, CodeAgent, HfApiModel
 import pandas as pd
 import tempfile
@@ -9,6 +9,33 @@ from database import engine, initialize_database
 # Ensure the database initializes (won't crash if empty)
 initialize_database()
 
+# SQL Execution Tool (FIXED - Defined BEFORE Use)
+@tool
+def sql_engine(query: str) -> str:
+    """
+    Executes an SQL SELECT query and returns the results.
+
+    Args:
+        query (str): The SQL query string to execute. Only SELECT queries are allowed.
+
+    Returns:
+        str: A formatted string containing the query results, or an error message if the query fails.
+    """
+    try:
+        with engine.connect() as con:
+            rows = con.execute(text(query)).fetchall()
+        if not rows:
+            return "No results found."
+        return "\n".join([", ".join(map(str, row)) for row in rows])
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Initialize CodeAgent for SQL query generation (Moved Below `sql_engine`)
+agent = CodeAgent(
+    tools=[sql_engine],
+    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
+)
+
 # Function to execute an uploaded SQL script
 def execute_sql_script(file_path):
     """
@@ -105,11 +132,48 @@ def handle_file_upload(file):
 
     return result, table_data
 
-# Initialize CodeAgent for SQL query generation
-agent = CodeAgent(
-    tools=[sql_engine],
-    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
-)
+# Function to handle natural language to SQL conversion
+def query_sql(user_query: str) -> str:
+    """
+    Converts a user's natural language query into an SQL query.
+
+    Args:
+        user_query (str): The question asked by the user.
+
+    Returns:
+        str: The results of the executed SQL query.
+    """
+    tables = get_table_names()
+    if not tables:
+        return "Error: No tables found. Please upload an SQL file first."
+
+    schema_info = "Available tables and columns:\n"
+    
+    for table in tables:
+        columns = get_table_schema(table)
+        schema_info += f"Table '{table}' has columns: {', '.join(columns)}.\n"
+
+    schema_info += "Generate a valid SQL SELECT query using ONLY these column names. DO NOT return anything other than the SQL query itself."
+
+    generated_sql = agent.run(f"{schema_info} Convert this request into SQL: {user_query}")
+
+    if not isinstance(generated_sql, str) or not generated_sql.strip().lower().startswith(("select", "show", "pragma")):
+        return "Error: Only SELECT queries are allowed."
+
+    return sql_engine(generated_sql)
+
+# Function to handle query input
+def handle_query(user_input: str) -> str:
+    """
+    Handles user input and returns the SQL query result.
+
+    Args:
+        user_input (str): User's natural language query.
+
+    Returns:
+        str: The query result or error message.
+    """
+    return query_sql(user_input)
 
 # Gradio UI
 with gr.Blocks() as demo:
@@ -132,5 +196,3 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
-
-