Update app.py

app.py

@@ -1,49 +1,20 @@
 import os
 import gradio as gr
+from sqlalchemy import text
+from smolagents import tool, CodeAgent, HfApiModel
+import spaces
 import pandas as pd
-from sqlalchemy import create_engine, text
-from langchain.tools import tool
-from code_agent import CodeAgent
-from hf_api_model import HfApiModel
-
-# Initialize SQLite database engine
-engine = create_engine('sqlite:///data.db')
-
-def clear_database():
-    """
-    Clear all tables from the database.
-    """
-    with engine.connect() as con:
-        # Get all table names
-        tables = con.execute(text(
-            "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
-        )).fetchall()
-        
-        # Drop each table
-        for table in tables:
-            con.execute(text(f"DROP TABLE IF EXISTS {table[0]}"))
-
-def create_dynamic_table(df):
-    """
-    Create a table dynamically based on DataFrame structure.
-    """
-    df.to_sql('data_table', engine, index=False, if_exists='replace')
-    return 'data_table'
-
-def insert_rows_into_table(records, table_name):
-    """
-    Insert records into the specified table.
-    """
-    with engine.begin() as conn:
-        for record in records:
-            conn.execute(
-                text(f"INSERT INTO {table_name} ({', '.join(record.keys())}) VALUES ({', '.join(['?' for _ in record])})")
-                .bindparams(*record.values())
-            )
+from database import (
+    engine,
+    create_dynamic_table,
+    clear_database,
+    insert_rows_into_table,
+    get_table_schema
+)
 
 def get_data_table():
     """
-    Get the current data table as a DataFrame.
+    Fetches all data from the current table and returns it as a Pandas DataFrame.
     """
     try:
         # Get list of tables
@@ -57,10 +28,18 @@ def get_data_table():
             
         # Use the first table found
         table_name = tables[0][0]
-        
-        # Read the table into a DataFrame
-        return pd.read_sql_table(table_name, engine)
-        
+            
+        with engine.connect() as con:
+            result = con.execute(text(f"SELECT * FROM {table_name}"))
+            rows = result.fetchall()
+            
+            if not rows:
+                return pd.DataFrame()
+                
+            columns = result.keys()
+            df = pd.DataFrame(rows, columns=columns)
+            return df
+
     except Exception as e:
         return pd.DataFrame({"Error": [str(e)]})
 
@@ -181,7 +160,7 @@ def process_uploaded_file(file):
 def sql_engine(query: str) -> str:
     """
     Executes an SQL query and returns formatted results.
-    
+
     Args:
         query: The SQL query string to execute on the database. Must be a valid SELECT query.
 
@@ -203,10 +182,48 @@ def sql_engine(query: str) -> str:
     except Exception as e:
         return f"Error: {str(e)}"
 
-def process_sql_result(generated_sql, table_name, column_names):
+agent = CodeAgent(
+    tools=[sql_engine],
+    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
+)
+
+def query_sql(user_query: str) -> str:
     """
-    Process and execute the generated SQL query.
+    Converts natural language input to an SQL query using CodeAgent.
     """
+    table_name, column_names, column_info = get_table_info()
+    
+    if not table_name:
+        return "Error: No data table exists. Please upload a file first."
+    
+    schema_info = (
+        f"The database has a table named '{table_name}' with the following columns:\n"
+        + "\n".join([
+            f"- {col} ({info['type']}){' primary key' if info['is_primary'] else ''}"
+            for col, info in column_info.items()
+        ])
+        + "\n\nGenerate a valid SQL SELECT query using ONLY these column names.\n"
+        "The table name is '" + table_name + "'.\n"
+        "If column names contain spaces, they must be quoted.\n"
+        "You can use aggregate functions like COUNT, AVG, SUM, etc.\n"
+        "DO NOT explain your reasoning, and DO NOT return anything other than the SQL query itself."
+    )
+
+    # Get SQL from the agent
+    generated_sql = agent.run(f"{schema_info} Convert this request into SQL: {user_query}")
+
+    if not isinstance(generated_sql, str):
+        return "Error: Invalid query generated"
+
+    # Clean up the SQL
+    if generated_sql.isnumeric():  # If the agent returned just a number
+        return generated_sql
+
+    # Extract just the SQL query if there's additional text
+    sql_lines = [line for line in generated_sql.split('\n') if 'select' in line.lower()]
+    if sql_lines:
+        generated_sql = sql_lines[0]
+
     # Remove any trailing semicolons
     generated_sql = generated_sql.strip().rstrip(';')
 
@@ -238,103 +255,12 @@ def process_sql_result(generated_sql, table_name, column_names):
             return generated_sql
         return f"Error executing query: {str(e)}"
 
-def query_sql(user_query: str, show_full: bool) -> tuple:
-    """
-    Converts natural language input to an SQL query using CodeAgent.
-    Returns both short and full responses based on switch state.
-    """
-    table_name, column_names, column_info = get_table_info()
-    
-    if not table_name:
-        return "Error: No data table exists. Please upload a file first.", ""
-
-    schema_info = (
-        f"The database has a table named '{table_name}' with the following columns:\n"
-        + "\n".join([
-            f"- {col} ({info['type']}{' primary key' if info['is_primary'] else ''})"
-            for col, info in column_info.items()
-        ])
-        + "\n\nGenerate a valid SQL SELECT query using ONLY these column names.\n"
-        "The table name is '" + table_name + "'.\n"
-        "If column names contain spaces, they must be quoted.\n"
-        "You can use aggregate functions like COUNT, AVG, SUM, etc.\n"
-        "DO NOT explain your reasoning, and DO NOT return anything other than the SQL query itself."
-    )
-
-    # Get full response from the agent
-    full_response = agent.run(f"{schema_info} Convert this request into SQL: {user_query}")
-    
-    # Process the short response as before
-    if not isinstance(full_response, str):
-        return "Error: Invalid query generated", ""
-
-    # Extract and process SQL for short response
-    generated_sql = full_response
-    if generated_sql.isnumeric():
-        short_response = generated_sql
-    else:
-        sql_lines = [line for line in generated_sql.split('\n') if 'select' in line.lower()]
-        if sql_lines:
-            generated_sql = sql_lines[0]
-        
-        # Process the SQL query and get the short result
-        short_response = process_sql_result(generated_sql, table_name, column_names)
-
-    return short_response, full_response
-
-def handle_upload(file_obj):
-    if file_obj is None:
-        return (
-            "Please upload a file.",
-            None,
-            "No schema available",
-            gr.update(visible=True),
-            gr.update(visible=False)
-        )
-    
-    success, message = process_uploaded_file(file_obj)
-    if success:
-        df = get_data_table()
-        _, _, column_info = get_table_info()
-        schema = "\n".join([
-            f"- {col} ({info['type']}){'primary key' if info['is_primary'] else ''}"
-            for col, info in column_info.items()
-        ])
-        return (
-            message,
-            df,
-            f"### Current Schema:\n```\n{schema}\n```",
-            gr.update(visible=False),
-            gr.update(visible=True)
-        )
-    return (
-        message,
-        None,
-        "No schema available",
-        gr.update(visible=True),
-        gr.update(visible=False)
-    )
-
-def refresh_data():
-    df = get_data_table()
-    _, _, column_info = get_table_info()
-    schema = "\n".join([
-        f"- {col} ({info['type']}){'primary key' if info['is_primary'] else ''}"
-        for col, info in column_info.items()
-    ])
-    return df, f"### Current Schema:\n```\n{schema}\n```"
-
-# Initialize the CodeAgent
-agent = CodeAgent(
-    tools=[sql_engine],
-    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
-)
-
 # Create the Gradio interface
 with gr.Blocks() as demo:
     with gr.Group() as upload_group:
         gr.Markdown("""
         # CSVAgent
+
         Upload your data file to begin.
         
         ### Supported File Types:
@@ -351,7 +277,10 @@ with gr.Blocks() as demo:
         https://tableconvert.com/sql-to-csv
         - Will work on the handling of SQL files soon.
         
+        
         ### Try it out! Upload a CSV file and then ask a question about the data!
+        - There is issues with the UI displaying the answer correctly, some questions such as "How many Customers are located in Korea?" 
+        The right answer will appear in the logs, but throws an error on the "Results" section.
         """)
         
         file_input = gr.File(
@@ -366,9 +295,6 @@ with gr.Blocks() as demo:
             with gr.Column(scale=1):
                 user_input = gr.Textbox(label="Ask a question about the data")
                 query_output = gr.Textbox(label="Result")
-                # Add the switch and secondary result box
-                full_response_switch = gr.Switch(label="Show Full Response", value=False)
-                full_response_output = gr.Textbox(label="Full Response", visible=False)
             
             with gr.Column(scale=2):
                 gr.Markdown("### Current Data")
@@ -381,6 +307,48 @@ with gr.Blocks() as demo:
         schema_display = gr.Markdown(value="Loading schema...")
         refresh_btn = gr.Button("Refresh Data")
 
+    def handle_upload(file_obj):
+        if file_obj is None:
+            return (
+                "Please upload a file.",
+                None,
+                "No schema available",
+                gr.update(visible=True),
+                gr.update(visible=False)
+            )
+        
+        success, message = process_uploaded_file(file_obj)
+        if success:
+            df = get_data_table()
+            _, _, column_info = get_table_info()
+            schema = "\n".join([
+                f"- {col} ({info['type']}){' primary key' if info['is_primary'] else ''}"
+                for col, info in column_info.items()
+            ])
+            return (
+                message,
+                df,
+                f"### Current Schema:\n```\n{schema}\n```",
+                gr.update(visible=False),
+                gr.update(visible=True)
+            )
+        return (
+            message,
+            None,
+            "No schema available",
+            gr.update(visible=True),
+            gr.update(visible=False)
+        )
+
+    def refresh_data():
+        df = get_data_table()
+        _, _, column_info = get_table_info()
+        schema = "\n".join([
+            f"- {col} ({info['type']}){' primary key' if info['is_primary'] else ''}"
+            for col, info in column_info.items()
+        ])
+        return df, f"### Current Schema:\n```\n{schema}\n```"
+
     # Event handlers
     file_input.upload(
         fn=handle_upload,
@@ -396,15 +364,8 @@ with gr.Blocks() as demo:
 
     user_input.change(
         fn=query_sql,
-        inputs=[user_input, full_response_switch],
-        outputs=[query_output, full_response_output]
-    )
-    
-    # Add switch change event to control visibility of full response
-    full_response_switch.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=full_response_switch,
-        outputs=full_response_output
+        inputs=user_input,
+        outputs=query_output
     )
     
     refresh_btn.click(