summarize_agentdojo.py

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

def check_result_exists(test_path, test_type):
    if test_type == "utility":
        sample_path = os.path.join(test_path, "slack", "user_task_0", "none", "none.json")
    elif test_type == "prompt_injection":
        sample_path = os.path.join(test_path, "slack", "user_task_0", "important_instructions", "injection_task_1.json")
    elif test_type == "data_injection":
        sample_path = os.path.join(test_path, "slack", "user_task_0", "data_injection", "injection_task_1.json")
    else:
        sample_path = None

    return sample_path and os.path.exists(sample_path)

def load_utility_results(test_path):
    utility_results = []
    for root, dirs, files in os.walk(test_path):
        for file in files:
            if file == "none.json":
                if "injection_task" in root:
                    continue
                file_path = os.path.join(root, file)
                with open(file_path, 'r') as f:
                    data = json.load(f)
                
                suite = data.get("suite_name", "Unknown")
                user_task_id = data.get("user_task_id", "Unknown")    
                utility = data.get("utility", None)
                latency = data.get("duration", None)
                total_token_usage = data.get("token_usage", {}).get("total_tokens", 0)
                input_token_usage = data.get("token_usage", {}).get("input_tokens", 0)
                output_token_usage = data.get("token_usage", {}).get("output_tokens", 0)
                expense = data.get("token_usage", {}).get("total_expense", 0)
                
                utility_results.append({
                    "suite": suite,
                    "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):
        if os.path.basename(root) == "important_instructions":
            for file in files:
                if file.startswith("injection_task_") and file.endswith(".json"):
                    file_path = os.path.join(root, file)
                    with open(file_path, 'r') as f:
                        data = json.load(f)
                    
                    suite = data.get("suite_name", "Unknown")
                    user_task_id = data.get("user_task_id", "Unknown")
                    injection_task_id = data.get("injection_task_id", "Unknown")
                    attack_success = data.get("security", False)
                    prompt_injection_results.append({
                        "suite": suite,
                        "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):
        if os.path.basename(root) == "data_injection":
            for file in files:
                if file.startswith("injection_task_") and file.endswith(".json"):
                    file_path = os.path.join(root, file)
                    with open(file_path, 'r') as f:
                        data = json.load(f)
                    
                    suite = data.get("suite_name", "Unknown")
                    user_task_id = data.get("user_task_id", "Unknown")
                    injection_task_id = data.get("injection_task_id", "Unknown")
                    attack_success = data.get("security", False)
                    data_injection_results.append({
                        "suite": suite,
                        "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"])
    
    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 tuples (suite, user_task_id)
    user_tasks = set(zip(utility_results["suite"], utility_results["user_task_id"]))

    utility_success_tasks = set(
        zip(utility_results[utility_results["utility"] == True]["suite"],
            utility_results[utility_results["utility"] == True]["user_task_id"])
    )
    prompt_injection_success_tasks = set(
        prompt_injection_results[prompt_injection_results["attack_success"] == True][["suite", "user_task_id"]].itertuples(index=False, name=None)
    ) if not prompt_injection_results.empty else set()
    data_injection_success_tasks = set(
        data_injection_results[data_injection_results["attack_success"] == True][["suite", "user_task_id"]].itertuples(index=False, name=None)
    ) 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 + " Agentdojo 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 agentdojo_result.py {test_name}")
        sys.exit(1)
    test_name = args[1]
    test_path = os.path.join(os.path.dirname(__file__), "agentdojo_result", test_name)

    # 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()