Update app.py

app.py

@@ -1,154 +1,227 @@
 import os
 import gradio as gr
-from sqlalchemy import text, inspect
+from sqlalchemy import text
 from smolagents import tool, CodeAgent, HfApiModel
+import spaces
 import pandas as pd
-import tempfile
-from database import engine, initialize_database
-
-# Ensure the database initializes with placeholder data
-initialize_database()
-
-# SQL Execution Tool (FIXED - Defined BEFORE Use)
-@tool
-def sql_engine(query: str) -> str:
+from database import (
+    engine,
+    create_dynamic_table,
+    clear_database,
+    insert_rows_into_table,
+    get_table_schema
+)
+
+def process_uploaded_file(file):
     """
-    Executes an SQL SELECT query and returns the results.
-
-    Parameters:
-        query (str): The SQL query string that should be executed.
-
-    Returns:
-        str: The result of the SQL query formatted as a string.
+    Process the uploaded CSV file and load it into the database.
     
-    Example:
-        >>> sql_engine("SELECT * FROM users")
-        '1, Alice, 30, 100.50\\n2, Bob, 24, 250.75\\n3, Charlie, 35, 80.00'
-    """
-    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)}"
-
-
-# Function to execute an uploaded SQL script
-def execute_sql_script(file_path):
-    """
-    Executes an uploaded SQL file to initialize the database.
-
     Args:
-        file_path (str): Path to the SQL file.
-
+        file: Path to the uploaded file
+        
     Returns:
-        str: Success message or error description.
+        tuple: (success_flag, message)
     """
     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. Database updated."
+        if file is None:
+            return False, "Please upload a file."
+            
+        # Read the CSV file
+        df = pd.read_csv(file.name)
+        
+        if len(df.columns) == 0:
+            return False, "Error: File contains no columns"
+            
+        # Clear existing database and create new table
+        clear_database()
+        table = create_dynamic_table(df)
+        
+        # Convert DataFrame to list of dictionaries and insert
+        records = df.to_dict('records')
+        insert_rows_into_table(records, table)
+        
+        return True, "File successfully loaded! Proceeding to query interface..."
+        
     except Exception as e:
-        return f"Error: {str(e)}"
-
-# Function to get table names dynamically
-def get_table_names():
-    """
-    Returns a list of tables available in the database.
+        return False, f"Error processing file: {str(e)}"
 
-    Returns:
-        list: List of table names.
+def get_data_table():
     """
-    inspector = inspect(engine)
-    return inspector.get_table_names()
-
-# Function to get table schema dynamically
-def get_table_schema(table_name):
+    Fetches all data from the current table and returns it as a Pandas DataFrame.
     """
-    Returns a list of column names for a given table.
-
-    Args:
-        table_name (str): Name of the table.
+    try:
+        with engine.connect() as con:
+            result = con.execute(text("SELECT * FROM data_table"))
+            rows = result.fetchall()
+            
+            if not rows:
+                return pd.DataFrame()
+                
+            # Get column names from the result
+            columns = result.keys()
+            df = pd.DataFrame(rows, columns=columns)
+            return df
 
-    Returns:
-        list: List of column names.
-    """
-    inspector = inspect(engine)
-    columns = inspector.get_columns(table_name)
-    return [col["name"] for col in columns]
+    except Exception as e:
+        return pd.DataFrame({"Error": [str(e)]})
 
-# Function to fetch data dynamically from any table
-def get_table_data(table_name):
+@tool
+def sql_engine(query: str) -> str:
     """
-    Retrieves all rows from a specified table as a Pandas DataFrame.
-
-    Args:
-        table_name (str): Name of the table.
-
-    Returns:
-        pd.DataFrame: Table data or an error message.
+    Executes an SQL query and returns formatted results.
     """
     try:
         with engine.connect() as con:
-            result = con.execute(text(f"SELECT * FROM {table_name}"))
-            rows = result.fetchall()
-
-        columns = get_table_schema(table_name)
+            rows = con.execute(text(query)).fetchall()
 
         if not rows:
-            return pd.DataFrame(columns=columns)
+            return "No results found."
 
-        return pd.DataFrame(rows, columns=columns)
-    except Exception as e:
-        return pd.DataFrame({"Error": [str(e)]})
+        if len(rows) == 1 and len(rows[0]) == 1:
+            return str(rows[0][0])
 
-# Function to handle SQL file uploads and execute them
-def handle_file_upload(file):
-    """
-    Handles SQL file upload, executes SQL, and updates database schema.
+        return "\n".join([", ".join(map(str, row)) for row in rows])
 
-    Args:
-        file (File): Uploaded SQL file.
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-    Returns:
-        tuple: Execution result message and updated table data.
+def query_sql(user_query: str) -> str:
+    """
+    Converts natural language input to an SQL query using CodeAgent
+    and returns the execution results.
     """
-    temp_file_path = tempfile.mkstemp(suffix=".sql")[1]
-    with open(temp_file_path, "wb") as temp_file:
-        temp_file.write(file.read())
+    schema = get_table_schema()
+    if not schema:
+        return "Error: No data table exists. Please upload a file first."
+        
+    schema_info = (
+        "The database has a table named 'data_table' with the following schema:\n"
+        f"{schema}\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}")
+
+    if not isinstance(generated_sql, str):
+        return f"{generated_sql}"
+
+    if not generated_sql.strip().lower().startswith(("select", "show", "pragma")):
+        return "Error: Only SELECT queries are allowed."
+
+    result = sql_engine(generated_sql)
     
-    result = execute_sql_script(temp_file_path)
-    tables = get_table_names()
-
-    if tables:
-        table_data = {table: get_table_data(table) for table in tables}
-    else:
-        table_data = {"Error": ["No tables found after execution. Ensure your SQL file creates tables."]}
-
-    return result, table_data
+    try:
+        float_result = float(result)
+        return f"{float_result:.2f}"
+    except ValueError:
+        return result
+
+agent = CodeAgent(
+    tools=[sql_engine],
+    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
+)
+
+# Create the upload interface
+def create_upload_interface():
+    with gr.Blocks() as upload_interface:
+        gr.Markdown("""
+        # Data Query Interface
+
+        Upload your CSV file to begin.
+        
+        ### Requirements:
+        - File must be in CSV format
+        - First column will be used as the primary key
+        - All columns will be automatically typed based on their content
+        """)
+        
+        file_input = gr.File(
+            label="Upload CSV File",
+            file_types=[".csv"],
+            type="file"
+        )
+        status = gr.Textbox(label="Status", interactive=False)
+        
+        def handle_upload(file):
+            success, message = process_uploaded_file(file)
+            if success:
+                return message, gr.Blocks.update(visible=False), gr.Blocks.update(visible=True)
+            return message, gr.Blocks.update(visible=True), gr.Blocks.update(visible=False)
+        
+        file_input.upload(
+            fn=handle_upload,
+            inputs=[file_input],
+            outputs=[status, upload_interface, query_interface]
+        )
+    
+    return upload_interface
+
+# Create the query interface
+def create_query_interface():
+    with gr.Blocks(visible=False) as query_interface:
+        gr.Markdown("""
+        ## Data Query Interface
+        
+        Enter your questions about the data in natural language.
+        The AI will convert your questions into SQL queries.
+        """)
+        
+        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("### Current Data")
+                data_table = gr.Dataframe(
+                    value=get_data_table(),
+                    label="Data Table",
+                    interactive=False
+                )
+        
+        # Show current schema
+        schema_display = gr.Markdown(value="Loading schema...")
+        
+        def update_schema():
+            schema = get_table_schema()
+            if schema:
+                return f"### Current Schema:\n```\n{schema}\n```"
+            return "No data loaded"
+        
+        user_input.change(
+            fn=query_sql,
+            inputs=[user_input],
+            outputs=[query_output]
+        )
+        
+        # Add refresh buttons
+        with gr.Row():
+            refresh_table_btn = gr.Button("Refresh Table")
+            refresh_schema_btn = gr.Button("Refresh Schema")
+            
+        refresh_table_btn.click(
+            fn=get_data_table,
+            outputs=[data_table]
+        )
+        
+        refresh_schema_btn.click(
+            fn=update_schema,
+            outputs=[schema_display]
+        )
+        
+        # Update schema on load
+        query_interface.load(
+            fn=update_schema,
+            outputs=[schema_display]
+        )
+    
+    return query_interface
 
-# Gradio UI
+# Create the main interface
 with gr.Blocks() as demo:
-    gr.Markdown("## SQL Query Interface")
-
-    with gr.Row():
-        user_input = gr.Textbox(label="Ask a question about the data")
-        query_output = gr.Textbox(label="Result")
-
-    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")
-    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])
+    upload_interface = create_upload_interface()
+    query_interface = create_query_interface()
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, share=True)