Update app.py

app.py

@@ -191,13 +191,11 @@ def query_sql(user_query: str) -> str:
     """
     Converts natural language input to an SQL query using CodeAgent.
     """
-    # Get current table information
     table_name, column_names, column_info = get_table_info()
     
     if not table_name:
         return "Error: No data table exists. Please upload a file first."
     
-    # Create schema information with actual column names
     schema_info = (
         f"The database has a table named '{table_name}' with the following columns:\n"
         + "\n".join([
@@ -211,12 +209,15 @@ def query_sql(user_query: str) -> str:
         "DO NOT explain your reasoning, and DO NOT return anything other than the SQL query itself."
     )
 
-    # Get the SQL query from the agent
+    # Get SQL from the agent
     generated_sql = agent.run(f"{schema_info} Convert this request into SQL: {user_query}")
 
-    # Clean up and validate the SQL
     if not isinstance(generated_sql, str):
-        return f"Error: Invalid query generated"
+        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()]
@@ -247,7 +248,11 @@ def query_sql(user_query: str) -> str:
             return f"{float_result:,.0f}"  # Format with commas, no decimals
         except ValueError:
             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
         return f"Error executing query: {str(e)}"
 
 # Create the Gradio interface