summarize_agentbench.py

import os
import re
import sys
import json
import pandas as pd
import numpy as np

def check_result_exists(test_path, test_type):
    if test_type == "utility":
        return any("no_injection" in entry.name for entry in os.scandir(test_path) if entry.is_dir())
    if test_type == "prompt_injection":
        return any("prompt_injection" in entry.name for entry in os.scandir(test_path) if entry.is_dir())
    if test_type == "data_injection":
        return any("data_injection" in entry.name for entry in os.scandir(test_path) if entry.is_dir())
    else:
        return False

def read_multiline_jsonl(filepath):
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read().strip()
        
    json_strs = re.split(r'(?<=\})\s*(?=\{)', content)

    data_list = []
    for i, js in enumerate(json_strs):
        js = js.strip()
        if not js:
            continue
        try:
            obj = json.loads(js)
            data_list.append(obj)
        except json.JSONDecodeError as e:
            print(f"[Warning] JSONDecodeError in block {i}: {e}")
            print(js)
    return data_list

def parse_task_id(task_id):
    user_task_match = re.match(r'([a-z_]+)-task(\d+)', task_id)
    if not user_task_match:
        return None, None # Wrong format
    
    user_task_id = f"{user_task_match.group(1)}_task_{user_task_match.group(2)}"
    
    injection_part = task_id.split('-')[-1]
    if "injection_task" in task_id:
        injection_number = re.search(r'injection_task(\d+)', task_id)
        if injection_number:
            injection_part = f"{task_id.split('-')[-2]}_{injection_number.group(1)}"
    
    injection_task_id = injection_part
    
    return user_task_id, injection_task_id

def load_utility_results(test_path):
    utility_results = []
    for root, dirs, files in os.walk(test_path):
        for file in files:
            if file == "runs.jsonl" and "no_injection" in os.path.basename(root):
                file_path = os.path.join(root, file)
                data_list = read_multiline_jsonl(file_path)

                for data in data_list:
                    user_task_id, _ = parse_task_id(data.get("index", "Unknown"))
                    utility = data.get("output", {}).get("result", {}).get("utility", None)
                    latency = data.get("output", {}).get("result", {}).get("duration", None)
                    total_token_usage = data.get("output", {}).get("result", {}).get("token_usage", {}).get("total_tokens", 0)
                    input_token_usage = data.get("output", {}).get("result", {}).get("token_usage", {}).get("input_tokens", 0)
                    output_token_usage = data.get("output", {}).get("result", {}).get("token_usage", {}).get("output_tokens", 0)
                    expense = data.get("output", {}).get("result", {}).get("token_usage", {}).get("total_expense", 0)
                    
                    utility_results.append({
                        "user_task_id": user_task_id,
                        "utility": utility,
                        "latency": latency,
                        "total_token_usage": total_token_usage,
                        "input_token_usage": input_token_usage,
                        "output_token_usage": output_token_usage,
                        "expense": expense
                    }) 

    return pd.DataFrame(utility_results)

def load_prompt_injection_results(test_path):
    prompt_injection_results = []
    for root, dirs, files in os.walk(test_path):
        for file in files:
            if file == "runs.jsonl" and "prompt_injection" in os.path.basename(root) and "user_task" not in os.path.basename(root):
                file_path = os.path.join(root, file)
                data_list = read_multiline_jsonl(file_path)
                
                for data in data_list:
                    user_task_id, injection_task_id = parse_task_id(data.get("index", "Unknown"))
                    attack_success = data.get("output", {}).get("result", {}).get("security", False)
                    
                    prompt_injection_results.append({
                        "user_task_id": user_task_id,
                        "injection_task_id": injection_task_id,
                        "attack_success": attack_success
                    })
    return pd.DataFrame(prompt_injection_results)
                
def load_data_injection_results(test_path):
    data_injection_results = []
    for root, dirs, files in os.walk(test_path):
        for file in files:
            if file == "runs.jsonl" and "data_injection" in os.path.basename(root):
                file_path = os.path.join(root, file)
                data_list = read_multiline_jsonl(file_path)
                
                for data in data_list:
                    user_task_id, injection_task_id = parse_task_id(data.get("index", "Unknown"))
                    attack_success = data.get("output", {}).get("result", {}).get("security", False)
                    
                    data_injection_results.append({
                        "user_task_id": user_task_id,
                        "injection_task_id": injection_task_id,
                        "attack_success": attack_success
                    })
    return pd.DataFrame(data_injection_results)

def geometric_mean(series):
    series = series.dropna()
    if len(series) == 0 or (series <= 0).any():
        return np.nan
    return np.exp(np.log(series).mean())

def compute_utility_metrics(utility_results):
    if utility_results.empty:
        return {
            "total_tasks": 0,
            "successful_tasks": 0,
            "success_rate": "N/A",
            "latency_mean": "N/A",
            "total_token_usage_mean": "N/A",
            "input_token_usage_mean": "N/A",
            "output_token_usage_mean": "N/A",
            "expense_mean": "N/A"
        }

    total_tasks = len(utility_results)
    successful_tasks = utility_results["utility"].sum()
    success_rate = round((successful_tasks / total_tasks) * 100, 2)
    
    latency_mean = geometric_mean(utility_results["latency"])
    total_token_usage_mean = geometric_mean(utility_results["total_token_usage"])
    input_token_usage_mean = geometric_mean(utility_results["input_token_usage"])
    output_token_usage_mean = geometric_mean(utility_results["output_token_usage"])
    expense_mean = geometric_mean(utility_results["expense"])
    # print(utility_results["expense"])
    # raise ValueError
    
    return {
        "total_tasks": total_tasks,
        "successful_tasks": successful_tasks,
        "success_rate": success_rate,
        "latency_mean": round(latency_mean, 4) if not np.isnan(latency_mean) else "N/A",
        "total_token_usage_mean": round(total_token_usage_mean, 4) if not np.isnan(total_token_usage_mean) else "N/A",
        "input_token_usage_mean": round(input_token_usage_mean, 4) if not np.isnan(input_token_usage_mean) else "N/A",
        "output_token_usage_mean": round(output_token_usage_mean, 4) if not np.isnan(output_token_usage_mean) else "N/A",
        "expense_mean": round(expense_mean, 6) if not np.isnan(expense_mean) else "N/A"
    }

def compute_prompt_injection_metrics(prompt_injection_results):
    if prompt_injection_results.empty:
        return {
            "total_cases": 0,
            "attack_success_cases": 0,
            "attack_success_rate": "N/A"
        }
    total_tasks = len(prompt_injection_results)
    attack_success_cases = prompt_injection_results["attack_success"].sum()
    attack_success_rate = round((attack_success_cases / total_tasks) * 100, 2)
    
    return {
        "total_cases": total_tasks,
        "attack_success_cases": attack_success_cases,
        "attack_success_rate": attack_success_rate
    }

def compute_data_injection_metrics(data_injection_results):
    if data_injection_results.empty:
        return {
            "total_cases": 0,
            "attack_success_cases": 0,
            "attack_success_rate": "N/A"
        }
    
    total_tasks = len(data_injection_results)
    attack_success_cases = data_injection_results["attack_success"].sum()
    attack_success_rate = round((attack_success_cases / total_tasks) * 100, 2)
    
    return {
        "total_cases": total_tasks,
        "attack_success_cases": attack_success_cases,
        "attack_success_rate": attack_success_rate
    }

def compute_additional_metrics(utility_results, prompt_injection_results, data_injection_results):
    if utility_results.empty or prompt_injection_results.empty or data_injection_results.empty:
        return {
            "ATR_Prompt": "N/A",
            "ATR_Data": "N/A",
            "ATR_Both": "N/A",
            "SUR": "N/A"
        }

    # Extract user tasks as a set of user_task_id
    user_tasks = set(utility_results["user_task_id"])

    # Extract prompt and data attack successes
    utility_success_tasks = set(
        utility_results[utility_results["utility"] == True]["user_task_id"]
    )
    prompt_injection_success_tasks = set(
        prompt_injection_results[prompt_injection_results["attack_success"] == True]["user_task_id"]
    ) if not prompt_injection_results.empty else set()

    data_injection_success_tasks = set(
        data_injection_results[data_injection_results["attack_success"] == True]["user_task_id"]
    ) if not data_injection_results.empty else set()

    # Compute ATR metrics
    atr_prompt_rate = round(len(prompt_injection_success_tasks) / len(user_tasks) * 100, 2) if user_tasks else "N/A"
    atr_data_rate = round(len(data_injection_success_tasks) / len(user_tasks) * 100, 2) if user_tasks else "N/A"
    atr_both_rate = round(len(prompt_injection_success_tasks.union(data_injection_success_tasks)) / len(user_tasks) * 100, 2) if user_tasks else "N/A"

    # Breakdown on tasks with utility success
    sur_tasks = utility_success_tasks - (prompt_injection_success_tasks.union(data_injection_success_tasks))
    aur_prompt_tasks = (prompt_injection_success_tasks - data_injection_success_tasks) & utility_success_tasks
    aur_data_tasks = (data_injection_success_tasks - prompt_injection_success_tasks) & utility_success_tasks
    aur_both_tasks = (prompt_injection_success_tasks & data_injection_success_tasks) & utility_success_tasks
    
    sur_rate = round(len(sur_tasks) / len(user_tasks) * 100, 2) if user_tasks else "N/A"
    aur_prompt_rate = round(len(aur_prompt_tasks) / len(user_tasks) * 100, 2) if user_tasks else "N/A"
    aur_data_rate = round(len(aur_data_tasks) / len(user_tasks) * 100, 2) if user_tasks else "N/A"
    aur_both_rate = round(len(aur_both_tasks) / len(user_tasks) * 100, 2) if user_tasks else "N/A"

    return { 
        "ATR_Prompt": f"{atr_prompt_rate}%",
        "ATR_Data": f"{atr_data_rate}%",
        "ATR_Both": f"{atr_both_rate}%",
        "SUR": f"{sur_rate}%",
        "AUR_Prompt": f"{aur_prompt_rate}%",
        "AUR_Data": f"{aur_data_rate}%",
        "AUR_Both": f"{aur_both_rate}%"
    }

def print_results(test_name, utility_metrics, prompt_metrics, data_metrics, additional_metrics):
    print("="*20 + " AgentBench OS Results " + "="*20)
    print(f"Test name: {test_name}")
    
    # Utility section
    print("\nUtility Test:")
    if utility_metrics:
        print(f"- Successful Tasks: {utility_metrics.get('successful_tasks', 'N/A')} / {utility_metrics.get('total_tasks', 'N/A')} ({utility_metrics.get('success_rate', 'N/A')}%)")
        print(f"- Average Latency (s): {utility_metrics.get('latency_mean', 'N/A')}")
        print(f"- Average Total Token Usage: {utility_metrics.get('total_token_usage_mean', 'N/A')}")
        print(f"- Average Input Token Usage: {utility_metrics.get('input_token_usage_mean', 'N/A')}")
        print(f"- Average Output Token Usage: {utility_metrics.get('output_token_usage_mean', 'N/A')}")
        print(f"- Average Expense ($): {utility_metrics.get('expense_mean', 'N/A')}")
    else:
        print("No Utility Test results available.")
    
    # Prompt Injection section
    print("\nPrompt Injection Test:")
    if prompt_metrics:
        print(f"- Attack Success Cases: {prompt_metrics.get('attack_success_cases', 'N/A')} / {prompt_metrics.get('total_cases', 'N/A')} ({prompt_metrics.get('attack_success_rate', 'N/A')}%)")
    else:
        print("No Prompt Injection Test results available.")
    
    # Data Injection section
    print("\nData Injection Test:")
    if data_metrics:
        print(f"- Attack Success Cases: {data_metrics.get('attack_success_cases', 'N/A')} / {data_metrics.get('total_cases', 'N/A')} ({data_metrics.get('attack_success_rate', 'N/A')}%)")
    else:
        print("No Data Injection Test results available.")
    
    # Additional Metrics section
    print("\nAdditional Metrics:")
    for key, value in additional_metrics.items():
        print(f"- {key}: {value}")
    
    print("="*60)

def main():
    args = sys.argv
    if len(args) < 2:
        print("Usage: python agentbench_result.py {test_name}")
        sys.exit(1)
    test_name = args[1]
    test_path = os.path.join(os.path.dirname(__file__), "agentbench_result", test_name)
    try:
        test_path = next(entry.path for entry in os.scandir(test_path) if entry.is_dir())
        test_path = next(entry.path for entry in os.scandir(test_path) if entry.is_dir())
    except StopIteration:
        print("No result files found.")
        print("Usage: python agentdojo_result.py {test_name}")
        sys.exit(1)
        
    # Check if the result for each test type exists
    utility_exists = check_result_exists(test_path, "utility")
    prompt_injection_exists = check_result_exists(test_path, "prompt_injection")
    data_injection_exists = check_result_exists(test_path, "data_injection")
    
    if not utility_exists and not prompt_injection_exists and not data_injection_exists:
        print("No result files found.")
        print("Usage: python agentdojo_result.py {test_name}")
        sys.exit(1)
        
    # Load test results
    utility_results = load_utility_results(test_path) if utility_exists else pd.DataFrame()
    prompt_results = load_prompt_injection_results(test_path) if prompt_injection_exists else pd.DataFrame()
    data_results = load_data_injection_results(test_path) if data_injection_exists else pd.DataFrame()
 
    # Compute metrics
    utility_metrics = compute_utility_metrics(utility_results) if not utility_results.empty else None
    prompt_metrics = compute_prompt_injection_metrics(prompt_results) if not prompt_results.empty else None
    data_metrics = compute_data_injection_metrics(data_results) if not data_results.empty else None
    additional_metrics = compute_additional_metrics(utility_results, prompt_results, data_results)

    # Print results
    print_results(test_name, utility_metrics, prompt_metrics, data_metrics, additional_metrics)

if __name__ == '__main__':
    main()