Update app.py

app.py

@@ -4,10 +4,14 @@ from sqlalchemy import text, create_engine, inspect
 from smolagents import tool, CodeAgent, HfApiModel
 import pandas as pd
 import tempfile
-from database import engine
+from database import engine, initialize_database
+
+# Ensure the database initializes
+initialize_database()
 
 # Function to execute SQL script from uploaded file
 def execute_sql_script(file_path):
+    """Executes an uploaded SQL file to initialize the database."""
     try:
         with engine.connect() as con:
             with open(file_path, "r") as f:
@@ -19,17 +23,20 @@ def execute_sql_script(file_path):
 
 # Function to fetch table names dynamically
 def get_table_names():
+    """Returns a list of all tables in the database."""
     inspector = inspect(engine)
     return inspector.get_table_names()
 
 # Function to fetch table schema dynamically
 def get_table_schema(table_name):
+    """Returns a list of column names for a given table."""
     inspector = inspect(engine)
     columns = inspector.get_columns(table_name)
     return [col["name"] for col in columns]
 
 # Function to fetch table data dynamically
 def get_table_data(table_name):
+    """Retrieves all rows from the specified table as a Pandas DataFrame."""
     try:
         with engine.connect() as con:
             result = con.execute(text(f"SELECT * FROM {table_name}"))
@@ -47,6 +54,7 @@ def get_table_data(table_name):
 # SQL Execution Tool
 @tool
 def sql_engine(query: str) -> str:
+    """Executes an SQL SELECT query and returns formatted results."""
     try:
         with engine.connect() as con:
             rows = con.execute(text(query)).fetchall()
@@ -58,8 +66,11 @@ def sql_engine(query: str) -> str:
 
 # Function to generate and execute SQL queries dynamically
 def query_sql(user_query: str) -> str:
-    # Get schema details dynamically
+    """Processes a user’s natural language query and generates an SQL query dynamically."""
     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:
@@ -77,10 +88,12 @@ def query_sql(user_query: str) -> str:
 
 # Function to handle query input
 def handle_query(user_input: str) -> str:
+    """Handles user input and returns the SQL query result."""
     return query_sql(user_input)
 
 # Function to handle SQL file uploads
 def handle_file_upload(file):
+    """Handles file upload, executes SQL, and updates database schema dynamically."""
     temp_file_path = tempfile.mkstemp(suffix=".sql")[1]
     with open(temp_file_path, "wb") as temp_file:
         temp_file.write(file.read())
@@ -91,7 +104,7 @@ def handle_file_upload(file):
     if tables:
         table_data = {table: get_table_data(table) for table in tables}
     else:
-        table_data = {}
+        table_data = {"Error": ["No tables found after execution. Ensure your SQL file creates tables."]}
 
     return result, table_data
 
@@ -122,3 +135,4 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+