analysis_src/extract_inference_data.py

#!/usr/bin/env python3
"""
Extract inference request and token usage data for all 'react with code' agents.

This script reads session.jsonl files to count inference requests and estimate token usage.
"""

import json
import sys
from pathlib import Path
from dataclasses import dataclass
import ast
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm import tqdm


PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from analysis_src.utils import (
    get_model_name,
    find_react_with_code_dirs,
    get_runs_stats,
    filter_scenarios_with_min_runs,
    find_latest_rollout_file
)

from analysis_src.model_styles import (
    get_model_style, MIN_FONT_SIZE, SINGLE_COLUMN_WIDTH, DOUBLE_COLUMN_WIDTH, get_color_palette, PLOT_PARAMETERS
)

# Paths
LEADERBOARD_DIR = PROJECT_ROOT / "ITBench-SRE-Agent" / "ITBench-Trajectories" / "ReAct-Agent-Trajectories"
RESULTS_JSON_DIR = LEADERBOARD_DIR / "results"
OUTPUT_DIR = PROJECT_ROOT / "ITBench-SRE-Agent" / "ITBench-Trajectories" / "output" / "inferences"

# Minimum runs per scenario required
MIN_RUNS_PER_SCENARIO = 3
MIN_QUALIFYING_SCENARIOS = 20

# Token estimation factor (chars per token)
CHARS_PER_TOKEN = 4

def extract_tokens_from_rollout(rollout_file: Path) -> dict:
    """
    Extract token counts and tool usage from a rollout file.
    
    Counts:
    - INPUT: system prompt + user messages + tool outputs
    - OUTPUT: assistant messages + tool call arguments
    - TOOLS: counts by tool name, including code execution
    """
    system_prompt_chars = 0
    user_input_chars = 0
    assistant_output_chars = 0
    tool_call_chars = 0
    tool_output_chars = 0
    
    assistant_msg_count = 0
    tool_call_count = 0
    tool_counts = {}  # tool_name -> count
    code_execution_count = 0  # Specifically track code/python execution
    
    # Tool names that indicate code execution
    CODE_TOOLS = ['execute_python', 'run_python', 'python', 'execute_code', 
                  'run_code', 'shell', 'bash', 'terminal', 'exec']
    
    try:
        with open(rollout_file) as f:
            for line in f:
                try:
                    d = json.loads(line)
                    msg_type = d.get('type', '')
                    payload = d.get('payload', {})
                    
                    if msg_type == 'session_meta':
                        # System prompt
                        instructions = payload.get('instructions', '')
                        system_prompt_chars += len(str(instructions))
                    
                    elif msg_type == 'response_item':
                        item_type = payload.get('type', '')
                        role = payload.get('role', '')
                        
                        if item_type == 'message':
                            content = payload.get('content', [])
                            if isinstance(content, list):
                                text = ' '.join([
                                    c.get('text', '') if isinstance(c, dict) else str(c) 
                                    for c in content
                                ])
                            else:
                                text = str(content)
                            
                            if role == 'user':
                                user_input_chars += len(text)
                            elif role == 'assistant':
                                assistant_output_chars += len(text)
                                assistant_msg_count += 1
                        
                        elif item_type == 'function_call':
                            # Tool call (output)
                            name = payload.get('name', '')
                            arguments = payload.get('arguments', '')
                            tool_call_chars += len(str(name)) + len(str(arguments))
                            tool_call_count += 1
                            
                            # Track tool usage
                            tool_counts[name] = tool_counts.get(name, 0) + 1
                            
                            # Check if it's code execution
                            name_lower = name.lower()
                            if any(code_tool in name_lower for code_tool in CODE_TOOLS):
                                code_execution_count += 1
                            # Also check if arguments contain python code patterns
                            args_str = str(arguments).lower()
                            if 'python' in name_lower or ('def ' in args_str or 'import ' in args_str):
                                code_execution_count += 1
                        
                        elif item_type == 'function_call_output':
                            # Tool output (input to model)
                            output = payload.get('output', '')
                            tool_output_chars += len(str(output))
                    
                except json.JSONDecodeError:
                    continue
    except Exception as e:
        return None
    
    # INPUT = system + user + tool outputs (fed back to model)
    input_chars = system_prompt_chars + user_input_chars + tool_output_chars
    # OUTPUT = assistant responses + tool call arguments
    output_chars = assistant_output_chars + tool_call_chars
    
    return {
        'system_prompt_chars': system_prompt_chars,
        'user_input_chars': user_input_chars,
        'assistant_output_chars': assistant_output_chars,
        'tool_call_chars': tool_call_chars,
        'tool_output_chars': tool_output_chars,
        'input_chars': input_chars,
        'output_chars': output_chars,
        'input_tokens': input_chars // CHARS_PER_TOKEN,
        'output_tokens': output_chars // CHARS_PER_TOKEN,
        'assistant_msg_count': assistant_msg_count,
        'tool_call_count': tool_call_count,
        'tool_counts': tool_counts,
        'code_execution_count': code_execution_count,
    }


def extract_session_stats(session_file: Path) -> dict:
    """
    Extract inference stats from session.jsonl and rollout files.
    
    Uses the latest rollout file for accurate token counting.
    """
    if not session_file.exists():
        return None
    
    trial_dir = session_file.parent
    
    # Count inference requests from session.jsonl
    inference_count = 0
    try:
        with open(session_file) as f:
            for line in f:
                try:
                    d = json.loads(line)
                    if d.get('type') == 'response_item':
                        inference_count += 1
                except json.JSONDecodeError:
                    continue
    except Exception as e:
        print(f"    Warning: Error reading {session_file}: {e}")
        return None
    
    # First check stdout.log for real token counts (OpenAI models)
    stdout_log = trial_dir / "traces" / "stdout.log"
    has_real_tokens = False
    input_tokens = 0
    output_tokens = 0
    cached_input_tokens = 0
    
    if stdout_log.exists():
        try:
            with open(stdout_log) as f:
                for line in f:
                    try:
                        d = json.loads(line)
                        if d.get('type') == 'turn.completed':
                            usage = d.get('usage', {})
                            input_tokens = usage.get('input_tokens', 0)
                            output_tokens = usage.get('output_tokens', 0)
                            cached_input_tokens = usage.get('cached_input_tokens', 0)
                            if input_tokens > 0 or output_tokens > 0:
                                has_real_tokens = True
                            break
                    except json.JSONDecodeError:
                        continue
        except Exception:
            pass
    
    # Extract from latest rollout file for tokens (if needed) and tool counts
    tool_call_count = 0
    tool_counts = {}
    code_execution_count = 0
    
    latest_rollout = find_latest_rollout_file(trial_dir)
    if latest_rollout:
        rollout_stats = extract_tokens_from_rollout(latest_rollout)
        if rollout_stats:
            # Use rollout tokens if no real API token data
            if not has_real_tokens:
                input_tokens = rollout_stats['input_tokens']
                output_tokens = rollout_stats['output_tokens']
            
            # Always use rollout for tool counts
            tool_call_count = rollout_stats['tool_call_count']
            tool_counts = rollout_stats['tool_counts']
            code_execution_count = rollout_stats['code_execution_count']
    
    return {
        'inference_count': inference_count,
        'input_tokens': input_tokens,
        'cached_input_tokens': cached_input_tokens,
        'output_tokens': output_tokens,
        'total_tokens': input_tokens + output_tokens,
        'has_real_tokens': has_real_tokens,
        'tool_call_count': tool_call_count,
        'tool_counts': tool_counts,
        'code_execution_count': code_execution_count,
    }


def read_agent_stats(agent_dir: Path) -> dict[str, list[dict]]:
    """
    Read session stats from all scenarios/trials for an agent.
    
    Returns:
        Dict mapping scenario_id -> list of stats (one per trial)
    """
    scenario_data = {}
    # Check if directory contains Scenario folders directly, or if we need to go one level deeper
    # (e.g., agent_dir/sre/Scenario-1, agent_dir/finops/Scenario-1, etc.)
    has_scenarios = any(d.name.startswith("Scenario") for d in agent_dir.iterdir() if d.is_dir())

    if not has_scenarios:
        # Look for subdirectories that might contain scenarios (sre, finops, etc.)
        subdirs = [d for d in agent_dir.iterdir() if d.is_dir() and not d.name.startswith(".")]
        if len(subdirs) == 1:
            # If there's exactly one subdirectory, use it
            agent_dir = subdirs[0]
        elif len(subdirs) > 1:
            # If there are multiple, try to find one with Scenario folders
            for subdir in subdirs:
                if any(d.name.startswith("Scenario") for d in subdir.iterdir() if d.is_dir()):
                    agent_dir = subdir
                    break


    
    for scenario_dir in agent_dir.iterdir():
        if not scenario_dir.is_dir() or not scenario_dir.name.startswith("Scenario"):
            continue
        
        scenario_id = scenario_dir.name
        trials = []
        
        for trial_dir in sorted(scenario_dir.iterdir()):
            if not trial_dir.is_dir():
                continue
            
            session_file = trial_dir / "session.jsonl"
            stats = extract_session_stats(session_file)
            if stats:
                trials.append(stats)
        
        if trials:
            scenario_data[scenario_id] = trials
    
    return scenario_data

def load_performance_data() -> pd.DataFrame:
    """Load performance data from the consistency analysis."""
    perf_file = PROJECT_ROOT / "ITBench-SRE-Agent" / "ITBench-Trajectories" / "output" / "consistency" / "performance_data.csv"
    if perf_file.exists():
        df = pd.read_csv(perf_file)
        return df[df["metric_raw"] == "root_cause_entity_f1"][["model", "performance"]]
    return pd.DataFrame()

def extract_all_data() -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    Extract inference data for all agents.
    
    Returns:
        - summary_df: Aggregated stats per model
        - detail_df: Per-scenario stats
    """
    agent_dirs = find_react_with_code_dirs(LEADERBOARD_DIR)
    print(f"Found {len(agent_dirs)} 'react with code' agent directories")
    
    summary_records = []
    detail_records = []
    
    for agent_dir in tqdm(agent_dirs, desc="Processing agents"):
        model_name = get_model_name(agent_dir.name)

        print(f"\nReading: {agent_dir.name}")
        scenario_data = read_agent_stats(agent_dir)

        n_scenarios, min_runs, max_runs, n_qualifying = get_runs_stats(scenario_data, MIN_RUNS_PER_SCENARIO)

        if n_scenarios == 0:
            print(f"  SKIPPING {model_name}: No session data found")
            continue

        if n_qualifying < MIN_QUALIFYING_SCENARIOS:
            print(f"  SKIPPING {model_name}: Only {n_qualifying}/{n_scenarios} scenarios have {MIN_RUNS_PER_SCENARIO}+ runs")
            continue

        # Filter scenarios
        scenario_data = filter_scenarios_with_min_runs(scenario_data, MIN_RUNS_PER_SCENARIO)
        n_scenarios_filtered = len(scenario_data)

        print(f"  Processing: {model_name} ({n_scenarios_filtered} scenarios)")

        # Aggregate across all scenarios and trials
        all_inference_counts = []
        all_input_tokens = []
        all_output_tokens = []
        all_total_tokens = []
        all_cached_tokens = []
        all_tool_call_counts = []
        all_code_execution_counts = []
        aggregated_tool_counts = {}

        for scenario_id, trials in tqdm(scenario_data.items(), desc=f"  {model_name} scenarios", leave=False):
            for trial in trials:
                all_inference_counts.append(trial['inference_count'])
                all_input_tokens.append(trial['input_tokens'])
                all_output_tokens.append(trial['output_tokens'])
                all_total_tokens.append(trial['total_tokens'])
                all_cached_tokens.append(trial.get('cached_input_tokens', 0))
                all_tool_call_counts.append(trial.get('tool_call_count', 0))
                all_code_execution_counts.append(trial.get('code_execution_count', 0))
                
                # Aggregate tool counts
                for tool_name, count in trial.get('tool_counts', {}).items():
                    aggregated_tool_counts[tool_name] = aggregated_tool_counts.get(tool_name, 0) + count
                
                detail_records.append({
                    'model': model_name,
                    'scenario': scenario_id,
                    'inference_count': trial['inference_count'],
                    'input_tokens': trial['input_tokens'],
                    'cached_input_tokens': trial.get('cached_input_tokens', 0),
                    'output_tokens': trial['output_tokens'],
                    'total_tokens': trial['total_tokens'],
                    'tool_call_count': trial.get('tool_call_count', 0),
                    'code_execution_count': trial.get('code_execution_count', 0),
                })
        
        # Summary stats
        summary_records.append({
            'model': model_name,
            'n_scenarios': n_scenarios_filtered,
            'n_trials': len(all_inference_counts),
            'avg_inference_count': np.mean(all_inference_counts),
            'std_inference_count': np.std(all_inference_counts),
            'avg_input_tokens': np.mean(all_input_tokens),
            'avg_cached_tokens': np.mean(all_cached_tokens),
            'avg_output_tokens': np.mean(all_output_tokens),
            'avg_total_tokens': np.mean(all_total_tokens),
            'total_inference_count': sum(all_inference_counts),
            'total_tokens': sum(all_total_tokens),
            'avg_tool_call_count': np.mean(all_tool_call_counts) if all_tool_call_counts else 0,
            'total_tool_calls': sum(all_tool_call_counts),
            'avg_code_execution_count': np.mean(all_code_execution_counts) if all_code_execution_counts else 0,
            'total_code_executions': sum(all_code_execution_counts),
            'top_tools': dict(sorted(aggregated_tool_counts.items(), key=lambda x: -x[1])[:10]),
        })
    
    summary_df = pd.DataFrame(summary_records)
    detail_df = pd.DataFrame(detail_records)
    
    # Merge with performance data
    perf_df = load_performance_data()
    if len(perf_df) > 0:
        summary_df = pd.merge(summary_df, perf_df, on='model', how='left')
    
    return summary_df, detail_df


def save_data(summary_df: pd.DataFrame, detail_df: pd.DataFrame):
    """Save extracted data to CSV files."""
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    
    summary_path = OUTPUT_DIR / "inference_summary.csv"
    detail_path = OUTPUT_DIR / "inference_detail.csv"
    
    summary_df.to_csv(summary_path, index=False)
    detail_df.to_csv(detail_path, index=False)
    
    print(f"\nData saved to:")
    print(f"  - {summary_path}")
    print(f"  - {detail_path}")


def print_summary(summary_df: pd.DataFrame):
    """Print summary table."""
    print("\n" + "="*80)
    print("Inference Summary")
    print("="*80)
    
    summary_df = summary_df.sort_values("avg_inference_count", ascending=False)
    
    print(f"\n{'Model':<25} {'Avg Infer':>10} {'Avg Tokens':>12} {'Avg In':>10} {'Avg Out':>10}")
    print("-" * 70)
    for _, row in summary_df.iterrows():
        print(f"{row['model']:<25} {row['avg_inference_count']:>10.1f} {row['avg_total_tokens']:>12.0f} {row['avg_input_tokens']:>10.0f} {row['avg_output_tokens']:>10.0f}")

def plot_tool_usage(summary_df: pd.DataFrame):
    """
    Figure: Tool usage per model - total tool calls and code execution.
    """
    plt.rcParams.update(PLOT_PARAMETERS)

    if 'avg_tool_call_count' not in summary_df.columns:
        print("Skipping tool usage: no tool data")
        return
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(SINGLE_COLUMN_WIDTH * 2, 2.5))
    
    data = summary_df.sort_values("avg_tool_call_count", ascending=True)

    color_palette = get_color_palette(len(data))
    colors = [color_palette[i % len(color_palette)] for i in range(len(data))]
    
    # Left: Total tool calls
    bars1 = ax1.barh(data["model"], data["avg_tool_call_count"], color=colors,
                     edgecolor='black', linewidth=0.5)
    ax1.set_xlabel("Avg. Tool Calls per Scenario")
    
    for bar, val in zip(bars1, data["avg_tool_call_count"]):
        ax1.text(val + 1, bar.get_y() + bar.get_height()/2,
                 f'{val:.0f}', va='center', ha='left', fontsize=MIN_FONT_SIZE - 1)
    
    ax1.set_xlim(0, data["avg_tool_call_count"].max() * 1.15)
    
    # Right: Code executions
    bars2 = ax2.barh(data["model"], data["avg_code_execution_count"], color=colors,
                     edgecolor='black', linewidth=0.5)
    ax2.set_xlabel("Avg. Code Executions per Scenario")
    
    for bar, val in zip(bars2, data["avg_code_execution_count"]):
        if val > 0:
            ax2.text(val + 0.5, bar.get_y() + bar.get_height()/2,
                     f'{val:.0f}', va='center', ha='left', fontsize=MIN_FONT_SIZE - 1)
    
    ax2.set_xlim(0, max(data["avg_code_execution_count"].max() * 1.3, 1))
    ax2.set_yticklabels([])
    
    plt.title("Tool Call Distribution")

    plt.tight_layout()
    plt.show()
    fig.savefig(OUTPUT_DIR / "fig_tool_usage.png")
    plt.close(fig)
    print("Saved: fig_tool_usage.png")

def plot_inference_vs_performance(summary_df: pd.DataFrame):
    """
    Figure 3: Inference count vs Performance scatter.
    """
    if 'performance' not in summary_df.columns:
        print("Skipping inference vs performance: no performance data")
        return

    plt.rcParams.update(PLOT_PARAMETERS)

    fig, ax = plt.subplots(figsize=(SINGLE_COLUMN_WIDTH, SINGLE_COLUMN_WIDTH))

    data = summary_df.dropna(subset=['performance'])

    # Manual label offsets to avoid overlap
    label_offsets = {
        "GPT-5.1": (-5, -8, "right", "top"),
        "o4-mini": (5, -8, "left", "top"),
        "GPT-OSS-120B": (5, 3, "left", "bottom"),
        "Gemini-2.5-Pro": (-5, 3, "right", "bottom"),
        "Gemini-3-Flash": (5, 3, "left", "bottom"),
        "gemini-3-pro-preview": (5, 3, "left", "bottom"),
        "Kimi-K2": (5, 3, "left", "bottom"),
    }

    # Get color palette
    color_palette = get_color_palette(len(data))

    # Scatter plot
    for i, (_, row) in enumerate(data.iterrows()):
        ax.scatter(row["avg_inference_count"], row["performance"],
                   c=[color_palette[i % len(color_palette)]], s=60, edgecolors='black',
                   linewidth=0.5, zorder=10)
        
        # Label with custom offset
        offset = label_offsets.get(row["model"], (5, 3, "left", "bottom"))
        ax.annotate(row["model"], 
                    (row["avg_inference_count"], row["performance"]),
                    xytext=(offset[0], offset[1]), textcoords='offset points',
                    fontsize=MIN_FONT_SIZE - 1, ha=offset[2], va=offset[3])
    
    ax.set_xlabel("Avg. Inference Requests")
    ax.set_ylabel("Performance (RC Entity F1)")
    ax.set_xlim(0, data["avg_inference_count"].max() * 1.2)
    ax.set_ylim(0, 0.7)
    
    plt.title("Inference Requests vs. Performance")

    plt.tight_layout()
    plt.show()
    fig.savefig(OUTPUT_DIR / "fig_inference_vs_performance.png")
    plt.close(fig)
    print("Saved: fig_inference_vs_performance.png")

def plot_tool_breakdown_heatmap(summary_df: pd.DataFrame):
    """
    Generate a heatmap showing which tools each agent uses most.
    """
    
    # Parse the stringified dict of top_tools
    tool_usage = []
    
    for _, row in summary_df.iterrows():
        if pd.isna(row.get('top_tools')):
            print("pd.isna")
            continue
            
        tools = row['top_tools']
        total_calls = row['total_tool_calls']
        if total_calls == 0:
            print("No tool calls")
            continue
                       
        for tool, count in tools.items():
            tool_usage.append({
                'model': row['model'],
                'tool': tool,
                'count': count,
                'avg_per_scenario': count / row['n_scenarios']
            })
            
            
    df = pd.DataFrame(tool_usage)
    if len(df) == 0:
        print("No tool usage data found")
        return

    # Pivot for heatmap
    pivot_df = df.pivot(index='model', columns='tool', values='avg_per_scenario').fillna(0)
    
    # Filter to top 10 most used tools across all models
    # top_tools = pivot_df.sum().sort_values(ascending=False).head(10).index
    top_tools = pivot_df.sum().sort_values(ascending=False).index
    pivot_df = pivot_df[top_tools]
    
    # Sort models by total tool usage
    pivot_df['total'] = pivot_df.sum(axis=1)
    pivot_df = pivot_df.sort_values('total', ascending=False).drop('total', axis=1)
    
    # Plot
    PLOT_PARAMETERS['font.size'] = 8
    plt.rcParams.update(PLOT_PARAMETERS)
    
    fig, ax = plt.subplots(figsize=(SINGLE_COLUMN_WIDTH * 2, 4))
    
    sns.heatmap(pivot_df, annot=True, fmt='.1f', cmap='YlOrRd', ax=ax,
                cbar_kws={'label': 'Avg. Calls per Scenario'})
    
    ax.set_xlabel("")
    ax.set_ylabel("")
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    
    plt.title("Tool Call Distribution")

    plt.tight_layout()
    plt.show()
    fig.savefig(OUTPUT_DIR / "fig_tool_usage_heatmap.png")
    plt.close(fig)
    print("Saved: fig_tool_usage_heatmap.png")

def main():
    print("Extracting inference data for 'react with code' agents...")
    print(f"Reading from directories: {LEADERBOARD_DIR}")
    print(f"Output directory: {OUTPUT_DIR}")
    
    summary_df, detail_df = extract_all_data()
    
    if len(summary_df) == 0:
        print("No data extracted!")
        return
    
    save_data(summary_df, detail_df)
    print_summary(summary_df)


if __name__ == "__main__":
    main()