Update app.py

app.py

@@ -1,144 +1,124 @@
 import os
 import gradio as gr
-from sqlalchemy import text
+from sqlalchemy import text, create_engine, inspect
 from smolagents import tool, CodeAgent, HfApiModel
-import spaces
-import pandas as pd
-from database import engine, receipts
 import pandas as pd
+import tempfile
+from database import engine
+
+# Function to execute SQL script from uploaded file
+def execute_sql_script(file_path):
+    try:
+        with engine.connect() as con:
+            with open(file_path, "r") as f:
+                sql_script = f.read()
+                con.execute(text(sql_script))
+        return "SQL file executed successfully."
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-def get_receipts_table():
-    """
-    Fetches all data from the 'receipts' table and returns it as a Pandas DataFrame.
+# Function to fetch table names dynamically
+def get_table_names():
+    inspector = inspect(engine)
+    return inspector.get_table_names()
 
-    Returns:
-        A Pandas DataFrame containing all receipt data.
-    """
+# Function to fetch table schema dynamically
+def get_table_schema(table_name):
+    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):
     try:
         with engine.connect() as con:
-            result = con.execute(text("SELECT * FROM receipts"))
+            result = con.execute(text(f"SELECT * FROM {table_name}"))
             rows = result.fetchall()
-        
-        if not rows:
-            return pd.DataFrame(columns=["receipt_id", "customer_name", "price", "tip"])
 
-        # Convert rows into a DataFrame
-        df = pd.DataFrame(rows, columns=["receipt_id", "customer_name", "price", "tip"])
-        return df
+        columns = get_table_schema(table_name)
 
+        if not rows:
+            return pd.DataFrame(columns=columns)
+
+        return pd.DataFrame(rows, columns=columns)
     except Exception as e:
-        return pd.DataFrame({"Error": [str(e)]})  # Return error message in DataFrame format
+        return pd.DataFrame({"Error": [str(e)]})
 
+# SQL Execution Tool
 @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."
-
-        if len(rows) == 1 and len(rows[0]) == 1:
-            return str(rows[0][0])  # Convert numerical result to string
-
         return "\n".join([", ".join(map(str, row)) for row in rows])
-
     except Exception as e:
         return f"Error: {str(e)}"
 
+# Function to generate and execute SQL queries dynamically
 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.
-    """
-
-    schema_info = (
-        "The database has a table named 'receipts' with the following schema:\n"
-        "- receipt_id (INTEGER, primary key)\n"
-        "- customer_name (VARCHAR(16))\n"
-        "- price (FLOAT)\n"
-        "- tip (FLOAT)\n"
-        "Generate a valid SQL SELECT query using ONLY these column names.\n"
-        "DO NOT explain your reasoning, and DO NOT return anything other than the SQL query itself."
-    )
-
-    generated_sql = agent.run(f"{schema_info} Convert this request into SQL: {user_query}")
+    # Get schema details dynamically
+    tables = get_table_names()
+    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"
 
-    if not isinstance(generated_sql, str):
-        return f"{generated_sql}"  # Handle unexpected numerical result
+    schema_info += "Generate a valid SQL SELECT query using ONLY these column names. DO NOT return anything other than the SQL query itself."
 
-    print(f"{generated_sql}")
+    generated_sql = agent.run(f"{schema_info} Convert this request into SQL: {user_query}")
 
-    if not generated_sql.strip().lower().startswith(("select", "show", "pragma")):
+    if not isinstance(generated_sql, str) or not generated_sql.strip().lower().startswith(("select", "show", "pragma")):
         return "Error: Only SELECT queries are allowed."
 
-    result = sql_engine(generated_sql)
-
-    print(f"{result}")
-    
-    try:
-        float_result = float(result)
-        return f"{float_result:.2f}"
-    except ValueError:
-        return result 
+    return sql_engine(generated_sql)
 
+# Function to handle query input
 def handle_query(user_input: str) -> str:
-    """
-    Calls query_sql, captures the output, and directly returns it to the UI.
+    return query_sql(user_input)
 
-    Args:
-        user_input: The user's natural language question.
+# Function to handle SQL file uploads
+def handle_file_upload(file):
+    temp_file_path = tempfile.mkstemp(suffix=".sql")[1]
+    with open(temp_file_path, "wb") as temp_file:
+        temp_file.write(file.read())
+    
+    result = execute_sql_script(temp_file_path)
+    tables = get_table_names()
 
-    Returns:
-        The SQL query result as a plain string to be displayed in the UI.
-    """
-    return query_sql(user_input)
+    if tables:
+        table_data = {table: get_table_data(table) for table in tables}
+    else:
+        table_data = {}
 
+    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"),
 )
 
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("""
-## Plain Text Query Interface
-
-This tool allows you to query a receipts database using natural language. Simply type your question into the input box, and the tool will generate and execute an SQL query to retrieve relevant data. The results will be displayed in the output box.
-
-### How to Use:
-1. Enter a question related to the receipts data in the text box.
-2. The tool will convert your question into an SQL query and execute it.
-3. The result will be displayed in the output box.
-4. The current receipts table is also displayed for reference.
-""")
+    gr.Markdown("## SQL Query Interface")
 
     with gr.Row():
-        with gr.Column(scale=1):
-            user_input = gr.Textbox(label="Ask a question about the data")
-            query_output = gr.Textbox(label="Result")
-        
-        with gr.Column(scale=2):
-            gr.Markdown("### Receipts Table")
-            receipts_table = gr.Dataframe(value=get_receipts_table(), label="Receipts Table")
+        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)
 
-    demo.load(fn=get_receipts_table, outputs=receipts_table)
+    gr.Markdown("## Upload SQL File to Execute")
+    file_upload = gr.File(label="Upload SQL File")
+    upload_output = gr.Textbox(label="Execution Result")
+
+    # Dynamic table display
+    table_output = gr.Dataframe(label="Database Tables (Dynamic)")
+
+    file_upload.change(fn=handle_file_upload, inputs=file_upload, outputs=[upload_output, table_output])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, share=True)