Update app.py

app.py

@@ -63,15 +63,21 @@ def extract_leaderboard_stats(game_name: str) -> pd.DataFrame:
         
         if game_name == "Total Performance":
             query = "SELECT game_name, COUNT(DISTINCT episode) AS games_played, " \
-                    "AVG(generation_time) AS avg_gen_time, SUM(reward) AS total_rewards " \
+                    "SUM(reward) AS total_rewards " \
                     "FROM game_results GROUP BY game_name"
             df = pd.read_sql_query(query, conn)
         else:
             query = "SELECT COUNT(DISTINCT episode) AS games_played, " \
-                    "AVG(generation_time) AS avg_gen_time, SUM(reward) AS total_rewards " \
+                    "SUM(reward) AS total_rewards " \
                     "FROM game_results WHERE game_name = ?"
             df = pd.read_sql_query(query, conn, params=(game_name,))
         
+        # Fetch average generation time from moves table
+        gen_time_query = """
+            SELECT AVG(generation_time) FROM moves WHERE game_name = ?
+        """
+        avg_gen_time = conn.execute(gen_time_query, (game_name,)).fetchone()[0] or 0
+        
         # Calculate win rate against random bot
         vs_random_query = """
             SELECT COUNT(*) FROM game_results 
@@ -87,6 +93,7 @@ def extract_leaderboard_stats(game_name: str) -> pd.DataFrame:
         
         df["agent_name"] = model_name
         df["agent_type"] = agent_type
+        df["avg_generation_time"] = round(avg_gen_time, 2)
         df["vs_random"] = round(vs_random_rate, 2)
         
         all_stats.append(df)
@@ -95,7 +102,7 @@ def extract_leaderboard_stats(game_name: str) -> pd.DataFrame:
     leaderboard_df = pd.concat(all_stats, ignore_index=True) if all_stats else pd.DataFrame()
     
     if leaderboard_df.empty:
-        leaderboard_df = pd.DataFrame(columns=["LLM Model", "# games", "moves/game", "illegal-moves", "win-rate", "vs Random"])
+        leaderboard_df = pd.DataFrame(columns=["LLM Model", "# games", "total rewards", "avg gen time", "win-rate", "vs Random"])
     
     return leaderboard_df
 
@@ -113,7 +120,7 @@ with gr.Blocks() as interface:
         gr.Markdown("# LLM Model Leaderboard\nTrack performance across different games!")
         available_games = get_available_games()
         leaderboard_game_dropdown = gr.Dropdown(available_games, label="Select Game", value="Total Performance")
-        leaderboard_table = gr.Dataframe(headers=["LLM Model", "# games", "moves/game", "illegal-moves", "win-rate", "vs Random"])
+        leaderboard_table = gr.Dataframe(headers=["LLM Model", "# games", "total rewards", "avg gen time", "win-rate", "vs Random"])
         generate_button = gr.Button("Generate Leaderboard JSON")
         download_component = gr.File(label="Download Leaderboard JSON")
         refresh_button = gr.Button("Refresh Leaderboard")