Update app.py

app.py

@@ -211,17 +211,20 @@ 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
     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"{generated_sql}"
+        return f"Error: Invalid query generated"
 
-    # Normalize the SQL for checking
-    normalized_sql = generated_sql.strip().lower()
-    
-    # Check if it's a valid SELECT query
-    if not (normalized_sql.startswith("select") and "from" in normalized_sql):
-        return "Error: Only SELECT queries are allowed."
+    # 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(';')
 
     # Fix table names
     for wrong_name in ['table_name', 'customers', 'main']:
@@ -234,14 +237,18 @@ def query_sql(user_query: str) -> str:
             if col in generated_sql and f'"{col}"' not in generated_sql and f'`{col}`' not in generated_sql:
                 generated_sql = generated_sql.replace(col, f'"{col}"')
 
-    # Execute the query
-    result = sql_engine(generated_sql)
-    
     try:
-        float_result = float(result)
-        return f"{float_result:.2f}"
-    except ValueError:
-        return result
+        # Execute the query
+        result = sql_engine(generated_sql)
+        
+        # Try to format as number if possible
+        try:
+            float_result = float(result)
+            return f"{float_result:,.0f}"  # Format with commas, no decimals
+        except ValueError:
+            return result
+    except Exception as e:
+        return f"Error executing query: {str(e)}"
 
 # Create the Gradio interface
 with gr.Blocks() as demo: