Update app.py

app.py

@@ -6,20 +6,20 @@ import pandas as pd
 import tempfile
 from database import engine, initialize_database
 
-# Ensure the database initializes (won't crash if empty)
+# Ensure the database initializes with placeholder data
 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.
+    Executes an SQL SELECT query on the database.
 
     Args:
-        query (str): The SQL query string to execute. Only SELECT queries are allowed.
+        query (str): The SQL query to be executed. Only SELECT queries are allowed.
 
     Returns:
-        str: A formatted string containing the query results, or an error message if the query fails.
+        str: A formatted string with query results, or an error message if the query fails.
     """
     try:
         with engine.connect() as con:
@@ -30,7 +30,7 @@ def sql_engine(query: str) -> str:
     except Exception as e:
         return f"Error: {str(e)}"
 
-# Initialize CodeAgent for SQL query generation (Moved Below `sql_engine`)
+# Initialize CodeAgent for SQL query generation
 agent = CodeAgent(
     tools=[sql_engine],
     model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
@@ -132,49 +132,6 @@ def handle_file_upload(file):
 
     return result, table_data
 
-# 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:
     gr.Markdown("## SQL Query Interface")
@@ -183,7 +140,7 @@ with gr.Blocks() as demo:
         user_input = gr.Textbox(label="Ask a question about the data")
         query_output = gr.Textbox(label="Result")
 
-    user_input.change(fn=handle_query, inputs=user_input, outputs=query_output)
+    user_input.change(fn=query_sql, inputs=user_input, outputs=query_output)
 
     gr.Markdown("## Upload SQL File to Execute")
     file_upload = gr.File(label="Upload SQL File")