Update app.py

app.py

@@ -1,20 +1,48 @@
 import os
 import gradio as gr
-from sqlalchemy import text
-from smolagents import tool, CodeAgent, HfApiModel
-import spaces
 import pandas as pd
-from database import (
-    engine,
-    create_dynamic_table,
-    clear_database,
-    insert_rows_into_table,
-    get_table_schema
-)
+from sqlalchemy import create_engine, text
+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())
+            )
 
 def get_data_table():
     """
-    Fetches all data from the current table and returns it as a Pandas DataFrame.
+    Get the current data table as a DataFrame.
     """
     try:
         # Get list of tables
@@ -28,18 +56,10 @@ def get_data_table():
             
         # Use the first table found
         table_name = tables[0][0]
-            
-        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
-
+        
+        # Read the table into a DataFrame
+        return pd.read_sql_table(table_name, engine)
+        
     except Exception as e:
         return pd.DataFrame({"Error": [str(e)]})
 
@@ -160,7 +180,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.
 
@@ -182,48 +202,10 @@ def sql_engine(query: str) -> str:
     except Exception as e:
         return f"Error: {str(e)}"
 
-agent = CodeAgent(
-    tools=[sql_engine],
-    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
-)
-
-def query_sql(user_query: str) -> str:
+def process_sql_result(generated_sql, table_name, column_names):
     """
-    Converts natural language input to an SQL query using CodeAgent.
+    Process and execute the generated SQL query.
     """
-    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(';')
 
@@ -255,12 +237,103 @@ def query_sql(user_query: str) -> str:
             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:
@@ -277,10 +350,7 @@ 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(
@@ -295,6 +365,9 @@ 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")
@@ -307,48 +380,6 @@ 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,
@@ -364,8 +395,15 @@ with gr.Blocks() as demo:
 
     user_input.change(
         fn=query_sql,
-        inputs=user_input,
-        outputs=query_output
+        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
     )
     
     refresh_btn.click(