Update app.py

app.py

@@ -22,18 +22,24 @@ def get_data_table():
             tables = con.execute(text(
                 "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
             )).fetchall()
+            
         if not tables:
             return pd.DataFrame()
+            
         # 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
+
     except Exception as e:
         return pd.DataFrame({"Error": [str(e)]})
 
@@ -49,13 +55,17 @@ def get_table_info():
             tables = con.execute(text(
                 "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
             )).fetchall()
+            
         if not tables:
             return None, [], {}
+            
         # Use the first table found
         table_name = tables[0][0]
+        
         # Get column information
         with engine.connect() as con:
             columns = con.execute(text(f"PRAGMA table_info({table_name})")).fetchall()
+            
         # Extract column names and types
         column_names = [col[1] for col in columns]
         column_info = {
@@ -65,7 +75,9 @@ def get_table_info():
             }
             for col in columns
         }
+        
         return table_name, column_names, column_info
+        
     except Exception as e:
         print(f"Error getting table info: {str(e)}")
         return None, [], {}
@@ -78,18 +90,24 @@ def process_sql_file(file_path):
         # Read the SQL file
         with open(file_path, 'r') as file:
             sql_content = file.read()
+            
         # Replace AUTO_INCREMENT with AUTOINCREMENT for SQLite compatibility
         sql_content = sql_content.replace('AUTO_INCREMENT', 'AUTOINCREMENT')
+            
         # Split into individual statements
         statements = [stmt.strip() for stmt in sql_content.split(';') if stmt.strip()]
+        
         # Clear existing database
         clear_database()
+        
         # Execute each statement
         with engine.begin() as conn:
             for statement in statements:
                 if statement.strip():
                     conn.execute(text(statement))
+                    
         return True, "SQL file successfully executed!"
+        
     except Exception as e:
         return False, f"Error processing SQL file: {str(e)}"
 
@@ -100,15 +118,20 @@ def process_csv_file(file_path):
     try:
         # Read the CSV file
         df = pd.read_csv(file_path)
+        
         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, "CSV file successfully loaded!"
+        
     except Exception as e:
         return False, f"Error processing CSV file: {str(e)}"
 
@@ -119,14 +142,17 @@ def process_uploaded_file(file):
     try:
         if file is None:
             return False, "Please upload a file."
+            
         # Get file extension
         file_ext = os.path.splitext(file)[1].lower()
+        
         if file_ext == '.sql':
             return process_sql_file(file)
         elif file_ext == '.csv':
             return process_csv_file(file)
         else:
             return False, "Error: Unsupported file type. Please upload either a .sql or .csv file."
+            
     except Exception as e:
         return False, f"Error processing file: {str(e)}"
 
@@ -134,19 +160,25 @@ 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.
+
     Returns:
         str: The formatted query results as a 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:  # Fixed the comparison operators here
+
+        if len(rows) == 1 and len(rows[0]) == 1:
             return str(rows[0][0])
+
         return "\n".join([", ".join(map(str, row)) for row in rows])
+
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -155,15 +187,15 @@ agent = CodeAgent(
     model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
 )
 
-def query_sql(user_query: str, show_full: bool) -> tuple:
+def query_sql(user_query: str) -> str:
     """
     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.", ""
-
+        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([
@@ -177,16 +209,15 @@ def query_sql(user_query: str, show_full: bool) -> tuple:
         "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}")
-    
-    if not isinstance(full_response, str):
-        return "Error: Invalid query generated", ""
+    # 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
-    generated_sql = full_response
     if generated_sql.isnumeric():  # If the agent returned just a number
-        return generated_sql, full_response
+        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()]
@@ -200,7 +231,7 @@ def query_sql(user_query: str, show_full: bool) -> tuple:
     for wrong_name in ['table_name', 'customers', 'main']:
         if wrong_name in generated_sql:
             generated_sql = generated_sql.replace(wrong_name, table_name)
-
+    
     # Add quotes around column names that need them
     for col in column_names:
         if ' ' in col:  # If column name contains spaces
@@ -210,40 +241,48 @@ def query_sql(user_query: str, show_full: bool) -> tuple:
     try:
         # Execute the query
         result = sql_engine(generated_sql)
+        
         # Try to format as number if possible
         try:
             float_result = float(result)
-            short_response = f"{float_result:,.0f}"  # Format with commas, no decimals
+            return f"{float_result:,.0f}"  # Format with commas, no decimals
         except ValueError:
-            short_response = result
-        return short_response, full_response
+            return result
+            
     except Exception as e:
         if str(e).startswith("(sqlite3.OperationalError) near"):
             # If it's a SQL syntax error, return the raw result
-            return generated_sql, full_response
-        return f"Error executing query: {str(e)}", full_response
+            return generated_sql
+        return f"Error executing query: {str(e)}"
 
 # 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:
         - CSV (.csv): CSV file with headers that will be automatically converted to a table
+        
         ### CSV Requirements:
         - Must include headers
         - First column will be used as the primary key
         - Column types will be automatically detected
         - Sample CSV Files: https://github.com/datablist/sample-csv-files
         ### Based on ZennyKenny's SqlAgent
+
         ### SQL to CSV File Conversion
         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?"
+        - 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(
             label="Upload Data File",
             file_types=[".csv", ".sql"],
@@ -256,8 +295,7 @@ 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")
-                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")
                 data_table = gr.Dataframe(
@@ -265,6 +303,7 @@ with gr.Blocks() as demo:
                     label="Data Table",
                     interactive=False
                 )
+        
         schema_display = gr.Markdown(value="Loading schema...")
         refresh_btn = gr.Button("Refresh Data")
 
@@ -277,10 +316,11 @@ with gr.Blocks() as demo:
                 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()
+            _, _, 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()
@@ -302,7 +342,7 @@ with gr.Blocks() as demo:
 
     def refresh_data():
         df = get_data_table()
-        _,_ , column_info = get_table_info()
+        _, _, 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()
@@ -324,14 +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]
-    )
-    
-    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(