Update app.py

app.py

@@ -207,6 +207,7 @@ def query_sql(user_query: str) -> str:
         + "\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."
     )
 
@@ -215,7 +216,11 @@ def query_sql(user_query: str) -> str:
     if not isinstance(generated_sql, str):
         return f"{generated_sql}"
 
-    if not generated_sql.strip().lower().startswith(("select", "show", "pragma")):
+    # 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."
 
     # Fix table names
@@ -229,6 +234,7 @@ 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: