added task heatmaps

app.py

@@ -6,7 +6,7 @@ from pathlib import Path
 import pandas as pd
 import os
 import json
-from utils.viz import create_scatter_plot, create_flow_chart, create_bar_chart
+from utils.viz import create_scatter_plot, create_flow_chart, create_bar_chart, create_task_success_heatmap
 from utils.processing import check_and_process_uploads
 from huggingface_hub import snapshot_download
 from apscheduler.schedulers.background import BackgroundScheduler
@@ -19,10 +19,8 @@ from apscheduler.schedulers.asyncio import AsyncIOScheduler
 import weave
 from utils.db import TracePreprocessor
 
-# Initialize the TracePreprocessor
 preprocessor = TracePreprocessor()
 
-
 from datetime import datetime
 
 abs_path = Path(__file__).parent
@@ -234,7 +232,7 @@ with gr.Blocks() as demo:
                             label="Select Columns to Display:",
                         ),
                         hide_columns=config.USACO_HIDE_COLUMNS,
-                        search_columns=config.USACO_SEARCH_COLUMNS,
+                        # search_columns=config.USACO_SEARCH_COLUMNS,
                         column_widths={"Agent Name": 40,
                                        "Accuracy": 20,
                                        "Total Cost": 20},
@@ -242,6 +240,17 @@ with gr.Blocks() as demo:
             with gr.Row():
                 scatter_plot = gr.Plot(create_scatter_plot(parse_json_files(os.path.join(abs_path, "evals_live"), 'usaco'), "Total Cost", "Accuracy", "Total Cost (in USD)", "Accuracy", ["Agent Name"]))
             
+            gr.Markdown("# Task Success Heatmap")
+            with gr.Row():
+                task_success_heatmap = gr.Plot()
+            demo.load(
+            lambda: create_task_success_heatmap(
+                preprocessor.get_task_success_data('usaco'),
+                'USACO'
+            ),
+            outputs=[task_success_heatmap]
+            )
+
             gr.Markdown("# Failure Report")
             with gr.Row():
                 with gr.Column(scale=1):
@@ -355,7 +364,7 @@ with gr.Blocks() as demo:
                             label="Select Columns to Display:",
                         ),
                         hide_columns=config.SWEBENCH_HIDE_COLUMNS,
-                        search_columns=config.SWEBENCH_SEARCH_COLUMNS,
+                        # search_columns=config.SWEBENCH_SEARCH_COLUMNS,
                         column_widths={"Agent Name": 40,
                                        "Accuracy": 20,
                                        "Total Cost": 20},
@@ -363,6 +372,17 @@ with gr.Blocks() as demo:
             with gr.Row():
                 scatter_plot = gr.Plot(create_scatter_plot(parse_json_files(os.path.join(abs_path, "evals_live"), 'swebench_verified'), "Total Cost", "Accuracy", "Total Cost (in USD)", "Accuracy", ["Agent Name"]))
             
+            gr.Markdown("# Task Success Heatmap")
+            with gr.Row():
+                task_success_heatmap = gr.Plot()
+            demo.load(
+            lambda: create_task_success_heatmap(
+                preprocessor.get_task_success_data('swebench_verified'),
+                'SWEBench Verified'
+            ),
+            outputs=[task_success_heatmap]
+            )
+            
             gr.Markdown("# Failure Report")
             with gr.Row():
                 with gr.Column(scale=1):
@@ -474,7 +494,7 @@ with gr.Blocks() as demo:
                             cant_deselect=["Agent Name"],
                             label="Select Columns to Display:",
                         ),
-                        search_columns=config.SWEBENCH_SEARCH_COLUMNS,
+                        # search_columns=config.SWEBENCH_SEARCH_COLUMNS,
                         hide_columns=config.SWEBENCH_HIDE_COLUMNS,
                         column_widths={"Agent Name": 40,
                                        "Accuracy": 20,
@@ -483,6 +503,17 @@ with gr.Blocks() as demo:
             with gr.Row():
                 scatter_plot = gr.Plot(create_scatter_plot(parse_json_files(os.path.join(abs_path, "evals_live"), 'swebench_lite'), "Total Cost", "Accuracy", "Total Cost (in USD)", "Accuracy", ["Agent Name"]))
 
+            gr.Markdown("# Task Success Heatmap")
+            with gr.Row():
+                task_success_heatmap = gr.Plot()
+            demo.load(
+            lambda: create_task_success_heatmap(
+                preprocessor.get_task_success_data('swebench_lite'),
+                'SWEBench Lite'
+            ),
+            outputs=[task_success_heatmap]
+            )
+            
             gr.Markdown("# Failure Report")
             with gr.Row():
                 with gr.Column(scale=1):
@@ -594,7 +625,7 @@ with gr.Blocks() as demo:
                             cant_deselect=["Agent Name"],
                             label="Select Columns to Display:",
                         ),
-                        search_columns=config.MLAGENTBENCH_SEARCH_COLUMNS,
+                        # search_columns=config.MLAGENTBENCH_SEARCH_COLUMNS,
                         hide_columns=config.MLAGENTBENCH_HIDE_COLUMNS,
                         column_widths={"Agent Name": 40,
                                        "Overall Score": 20,
@@ -714,12 +745,14 @@ with gr.Blocks() as demo:
 
 async def main():
     # Preprocess traces
+    # preprocessor = TracePreprocessor()
     # preprocessor.preprocess_traces('evals_live')
+    # preprocessor = TracePreprocessor()
     
-    # # # Download the results from the Hugging Face Hub
+    # Download the results from the Hugging Face Hub
     # await asyncio.to_thread(download_latest_results)
 
-    # # Check for new uploads and process them
+    # Check for new uploads and process them
     # await check_and_process_uploads()
     
     scheduler = AsyncIOScheduler()

utils/db.py

@@ -5,6 +5,7 @@ import pickle
 from functools import lru_cache
 import threading
 import pandas as pd
+import ast
 
 class TracePreprocessor:
     def __init__(self, db_path='preprocessed_traces.db'):
@@ -40,6 +41,8 @@ class TracePreprocessor:
                     benchmark_name TEXT,
                     agent_name TEXT,
                     date TEXT,
+                    successful_tasks TEXT,
+                    failed_tasks TEXT,
                     total_cost REAL,
                     accuracy REAL,
                     precision REAL,
@@ -95,12 +98,14 @@ class TracePreprocessor:
                 with self.get_conn() as conn:
                     conn.execute('''
                         INSERT OR REPLACE INTO parsed_results 
-                        (benchmark_name, agent_name, date, total_cost, accuracy, precision, recall, f1_score, auc, overall_score, vectorization_score, fathomnet_score, feedback_score, house_price_score, spaceship_titanic_score, amp_parkinsons_disease_progression_prediction_score, cifar10_score, imdb_score) 
-                        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                        (benchmark_name, agent_name, date, successful_tasks, failed_tasks, total_cost, accuracy, precision, recall, f1_score, auc, overall_score, vectorization_score, fathomnet_score, feedback_score, house_price_score, spaceship_titanic_score, amp_parkinsons_disease_progression_prediction_score, cifar10_score, imdb_score) 
+                        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                     ''', (
                         benchmark_name,
                         agent_name,
                         config['date'],
+                        str(results.get('successful_tasks')),
+                        str(results.get('failed_tasks')),
                         results.get('total_cost'),
                         results.get('accuracy'),
                         results.get('precision'),
@@ -161,6 +166,8 @@ class TracePreprocessor:
             if column in df.columns:
                 df[column] = df[column].round(3)
 
+        df = df.drop(columns=['successful_tasks', 'failed_tasks'], axis=1)
+
         # Rename columns
         df = df.rename(columns={
             'agent_name': 'Agent Name',
@@ -183,6 +190,45 @@ class TracePreprocessor:
         })
 
         return df
+    
+    def get_task_success_data(self, benchmark_name):
+        with self.get_conn() as conn:
+            query = '''
+                SELECT agent_name, successful_tasks, failed_tasks
+                FROM parsed_results 
+                WHERE benchmark_name = ?
+            '''
+            df = pd.read_sql_query(query, conn, params=(benchmark_name,))
+        
+        # Get all unique task IDs
+        task_ids = set()
+        for tasks in df['successful_tasks']:
+            if ast.literal_eval(tasks) is not None:
+                task_ids.update(ast.literal_eval(tasks))
+        for tasks in df['failed_tasks']:
+            if ast.literal_eval(tasks) is not None:
+                task_ids.update(ast.literal_eval(tasks))
+
+        # Create a DataFrame with agent_name, task_ids, and success columns
+        data_list = []
+        for _, row in df.iterrows():
+            agent_name = row['agent_name']
+            for task_id in task_ids:
+                success = 1 if task_id in row['successful_tasks'] else 0
+                data_list.append({
+                    'agent_name': agent_name,
+                    'task_id': task_id,
+                    'success': success
+                })
+        df = pd.DataFrame(data_list)
+
+        df = df.rename(columns={
+            'agent_name': 'Agent Name',
+            'task_id': 'Task ID',
+            'success': 'Success'
+        })
+
+        return df
 
 if __name__ == '__main__':
     preprocessor = TracePreprocessor()

utils/viz.py

@@ -4,6 +4,65 @@ from utils.pareto import Agent, compute_pareto_frontier
 import plotly.graph_objects as go
 import textwrap
 
+def create_task_success_heatmap(df, benchmark_name):
+    # Pivot the dataframe to create a matrix of agents vs tasks
+    pivot_df = df.pivot(index='Agent Name', columns='Task ID', values='Success')
+    
+    # Create the heatmap
+    fig = go.Figure(data=go.Heatmap(
+        z=pivot_df.values,
+        y=pivot_df.index,
+        x=pivot_df.columns,
+        colorscale=[[0, 'white'], [1, '#1b9e77']],  # White for failed, green for success
+        showscale=False,
+        hovertemplate='<b>Agent:</b> %{y}<br>' +
+                      '<b>Task:</b> %{x}<br>' +
+                      '<b>Status:</b> %{z}<extra></extra>'
+    ))
+    
+    # Update the layout
+    fig.update_layout(
+        xaxis_title='Task ID',
+        height=600,
+        width=1300,
+        yaxis=dict(
+            autorange='reversed',
+            showticklabels=True,  # Show y-axis tick labels (agent names)
+            showline=True,
+            linecolor='black',
+            showgrid=False
+        ),
+        xaxis=dict(
+            side='top',
+            showticklabels=False,  # Hide x-axis tick labels (task IDs)
+            showline=True,
+            linecolor='black',
+            showgrid=False
+        ),
+        plot_bgcolor='white',
+        paper_bgcolor='white',
+        hoverlabel=dict(
+            bgcolor="white", 
+            font_size=12, 
+            font_family="Arial"
+        ),
+        modebar=dict(
+            activecolor='#1f77b4',
+            orientation='h',
+            bgcolor='rgba(255,255,255,0.8)',
+            color='#777',
+            add=['pan2d'],
+            remove=[
+                'zoom2d', 'zoomIn2d', 'zoomOut2d', 'resetScale2d',
+                'hoverClosestCartesian', 'hoverCompareCartesian',
+                'toggleSpikelines', 'lasso2d', 'lasso', 'select2d', 'select'
+            ]
+        ),
+        dragmode='pan'
+    )
+    
+    return fig
+
 def create_bar_chart(categories, values, x_label, y_label, title):
     # Sort categories and values based on values in descending order
     sorted_data = sorted(zip(categories, values), key=lambda x: x[1], reverse=True)