diff --git a/benchmarks/Agent2PwnBench/.env.example b/benchmarks/Agent2PwnBench/.env.example new file mode 100644 index 00000000..76d4aa7e --- /dev/null +++ b/benchmarks/Agent2PwnBench/.env.example @@ -0,0 +1,12 @@ +PORTSWIGGER_USERNAME='email' # Email used to log in on PortSwigger Web Academy +PORTSWIGGER_PASSWORD='password' # Password used to log in on PortSwigger Web Academy +PORTSWIGGER_SECTION='sql-injection' # Change this to the type of lab +NUMBER_OF_LABS=2 # To test all the labs in the section, change this to -1 +OPENAI_API_KEY='sk-123' # The api key of your OpenAI account (use for LLMs in openai) +ANTHROPIC_API_KEY="" # The api key of your ANTHROPIC account (use for LLMs in antrhopic) +OLLAMA="" # The api key of your OLLAMA account (use for LLMs in OLLAMA) +CAI_STREAM=false # Whether to stream the output of the LLM in CAI or not +CAI_MODEL='openai/gpt-4o' # Name of the LLM used in CAI +CAI_AGENT="bug_bounter_agent" # Name of the CAI agent. Available options are (blueteam_agent, web_bounter_agent, redteam_agent ) +BURPSUITE_SERVER_URL="http://127.0.0.1:9876/sse" # Local url of the MCP server for burpsuite +#for more examples see https://github.com/aliasrobotics/cai/blob/main/.env.example \ No newline at end of file diff --git a/benchmarks/Prompt2PwnBench/README.md b/benchmarks/Agent2PwnBench/README.md similarity index 89% rename from benchmarks/Prompt2PwnBench/README.md rename to benchmarks/Agent2PwnBench/README.md index 9f9d7b27..9cc4cd3b 100644 --- a/benchmarks/Prompt2PwnBench/README.md +++ b/benchmarks/Agent2PwnBench/README.md @@ -1,70 +1,69 @@ -# Benchmarking Agentic AI in Cybersecurity: Comparing Models and Prompting Methods in Lab Challenges with CAI. +# Benchmarking Agentic AI in Cybersecurity: Comparing Models and CAI Agents in Lab Challenges of PortSwigger Web Academy. ## Introduction -This project explores and evaluates the integration of **Large Language Models (LLMs)** into web application attack scenarios using the **Cybersecurity AI (CAI)** framework. The goal is to test various prompting methods and different LLMs to assess their effectiveness in identifying vulnerabilities within web applications. +This project explores and evaluates the integration of **Large Language Models (LLMs)** into web application attack scenarios using the **Cybersecurity AI (CAI)** framework. The goal is to test various predefined agents in CAI and different LLMs to assess their effectiveness in identifying vulnerabilities within web applications. **Fully Automated (No HITL):** -The pipeline is designed to be **fully automated, with no Human-in-the-Loop (HITL)**. When the agent attempts to solve the challenge labs, **no human interaction with the model is required**; all decisions, iterations, and actions are executed autonomously according to the experiment’s configuration and the prompt templates. +The pipeline is designed to be **fully automated, with no Human-in-the-Loop (HITL)**. When the agent attempts to solve the challenge labs, **no human interaction with the model is required**; all decisions, iterations, and actions are executed autonomously according to the experiment’s configuration. ## Objectives This project focuses on the following objectives: -- Compare the performance of different LLMs within the [**CAI Framework**](https://aliasrobotics.github.io/cai/). +- Compare the performance of agents using different LLMs within the [**CAI Framework**](https://aliasrobotics.github.io/cai/). - Use [*PortSwigger labs*](https://portswigger.net/web-security) as an environment to test the LLMs. -- Evaluate the effectiveness of the models in identifying and exploiting common web vulnerabilities. -- Compare the models using prompting methods such as **zero-shot**, **few-shot**, and **chain-of-thought**. +- Evaluate the effectiveness of the agents in identifying and exploiting common web vulnerabilities. - Assess performance using metrics such as **turns, time, cost, tokens,** and **number of payloads (tools) generated**. -- Create a reproducible framework to evaluate the LLMs. +- Create a reproducible framework to evaluate the agents. ## Methodolody The program follows a sequence of steps to evaluate the models. -1. The user configures the variables for the LLM, the prompt method, and the PortSwigger lab environment. +1. The user configures the variables inside the .env file. 2. The PortSwigger bot extracts the data from the labs. -3. The prompt method templates are formatted with the lab information. -4. The custom AI agent in CAI runs and attempts to solve the lab challenges. +3. The main user prompt is formatted with the lab information. +4. The CAI agent runs and attempts to solve the lab challenges. 5. The PortSwigger bot verifies if each lab is solved. 6. The logs of the labs and terminal outputs are saved. 7. After the agent completes all tasks, the lab logs can be evaluated using the metrics.ipynb notebook. ## Steps for Reproducibility -1. Create a `.env` file in the main folder. For more details, see [**.env.example**](.env.example) file. -2. Configure the variables related to the PortSwigger account and the LLM used. You can create a PortSwigger account [here](https://portswigger.net/web-security). +1. Create a PortSwigger Web Academy account [here](https://portswigger.net/web-security). +2. Install Burp Suite Community in you local machine [here](https://portswigger.net/burp/communitydownload). 3. Install the Python dependencies with the command: ```bash pip install -r requirements.txt ``` -4. Configure the labs and agent parameters in the **main.py** or **server.py** script as follows. To see more available sections, see [**topic_prefixes.json**](utils/topics_prefixes.json) file. - ```python - SECTION = "sql-injection" # Change this to the type of lab - N_LABS = 4 # To test all the labs in the section, change this to -1 - AGENT = "webbounty" - PROMPT_TYPE = "zero-shot" # Change this to the desired prompt method +4. Create a `.env` file in the main folder. For more details, see [**.env.example**](.env.example) file. + +5. Configure the environmet variables as follows. + ```python + PORTSWIGGER_USERNAME='email' # Email used to log in on PortSwigger Web Academy + PORTSWIGGER_PASSWORD='password' # Password used to log in on PortSwigger Web Academy + PORTSWIGGER_SECTION='sql-injection' # Change this to the type of lab + NUMBER_OF_LABS=10 # To test all the labs in the section, change this to -1 + CAI_MODEL='openai/gpt-4o' # LLMs used in CAI + CAI_AGENT="bug_bounter_agent" # Name of the CAI agent. Available options are (blueteam_agent, bug_bounter_agent, redteam_agent) + BURPSUITE_SERVER_URL="http://127.0.0.1:9876/sse" # Local url of the MCP server for burpsuite ``` - To see more information about the prompt templates by type, see the [**promts.yml**](prompts.yml) file. -5. Open a terminal in the main folder and run the main script with the command: + To see more available labs sections, see [**topic_prefixes.json**](utils/topics_prefixes.json) file. + + To configure the Burp Suite MCP server to interact with the labs, you need first to install the MCP server extension. More information on this [link](https://portswigger.net/bappstore/9952290f04ed4f628e624d0aa9dccebc). + +6. Open Burp Suite Community Edition Desktop application. + +7. Open a terminal in the main folder and run the main script with the command: ```bash python main.py ``` - In case you want to run the script using Burp Suite MCP server to interact with the labs, you need first to install the MCP server. More information on this [link](https://portswigger.net/bappstore/9952290f04ed4f628e624d0aa9dccebc). - Then, set up the variable SERVER_URL in the script server.py as follows: - ```python - SERVER_URL = "http://127.0.0.1:9876/sse" - ``` - Finally run the script with python. - ```bash - python server.py - ``` -6. Once the script stops, create the metrics table and graphs running the notebook -[**metrics.ipynb**](metrics.ipynb). +8. Once the script stops, create the metrics table and graphs running the notebook [**metrics.ipynb**](metrics.ipynb). ## Project Folder Structure ```plaintext -Prompt2PwnBench/ # Root directory of the project +Agent2PwnBench/ # Root directory of the project ├── logs/ # CAI log outputs ├── results/ # Final experiment logs ├── terminal-output/ # terminal output sessions @@ -75,31 +74,30 @@ Prompt2PwnBench/ # Root directory of the project │ ├── helpers.py # General helper functions │ ├── portswiggerbot.py # Automation for PortSwigger bot │ └── topics-prefixes.json # Topic prefixes for PortSwigger bot -│ └── portswigger-labs.json # Metadata of Portswigger Web Academy labs -├── main.py # Main execution script (it uses simple curl tools to interact with labs) -├── server.py # Main execution script (it uses Burp Suite MCP server to interact with labs) +│ └── labs.json # Metadata of Portswigger Web Academy labs +├── main.py # Main execution script ├── metrics.ipynb # Notebook for evaluating metrics └── prompts.yml # Prompt templates └── .env.example # env file example └── requirements.txt # requirements file for python libs ``` -## Prompt Learning Methods +## CAI Agents and prompts -One of the objectives of this project is to compare AI models in the CAI framework using different prompt methods. -For this purpose, a YAML file was created containing different types of system and user prompts explained in the following table. +One of the objectives of this project is to compare CAI agents and their default prompts to measure their performance when solving the labs. +For this purpose, in the following table there is a brief description of the CAI agents available to use in this benchmark. -For more details of the full text in the prompts, see the file [prompts.yml](prompts.yml). +| **Name** | **System Prompt Link** | **User Prompt Link** | +|-------------------------|------------|------------| +| Blue Team Agent | [click here](https://github.com/aliasrobotics/cai/blob/main/src/cai/prompts/system_blue_team_agent.md) | [click here](prompts.yml) | +| Red Team Agent | [click here](https://github.com/aliasrobotics/cai/blob/main/src/cai/prompts/system_red_team_agent.md) | [click here](prompts.yml) | +| Bug Bounter Agent | [click here](https://github.com/aliasrobotics/cai/blob/main/src/cai/prompts/system_bug_bounter.md) | [click here](prompts.yml) | -| **Method** | **Prompt** | **Description** | -|-------------------------|------------|---------------------------------------------------------------------------------| -| Zero-shot | System | Gives the model the role of bug bounty agent for vulnerabilities of PortSwigger labs | -| Zero-shot | User | Gives the model the task to attack the target lab without any example | -| Few-shot | User | Gives the model the task to attack the target lab with a small number of examples within the prompt itself to guide its response | -| Chain-of-thought (CoT) | User | Gives the model the task to attack the target lab with a step-by-step explanation | +Custom user prompt templates can be modified or created in the [**prompts.yml**](prompts.yml) to improve the performance of the CAI agents. -New custom prompt templates can be created using the same structure explained above. +For more details of all the CAI agents, check this [link](https://github.com/aliasrobotics/cai/tree/main/src/cai/agents). +For more details of all the CAI agents prompts, check this [link](https://github.com/aliasrobotics/cai/tree/main/src/cai/prompts). ## Metrics and Results The following metrics are used to compare the models performance, and they are calculated in the [**metrics.ipynb**](metrics.ipynb) file. @@ -124,21 +122,19 @@ but failed to solve the challenge. and solved the challenge. ### Example of performance results. -The following example table summarizes the performance metrics of **DeepSeek-V3** and **GPT-4o** when solving a total of 15 security labs (5 each on SQL Injection, Cross-Site Scripting, and Cross-Site Request Forgery). -The results are broken down by different prompting strategies and include interaction times, token usage, and assistant behavior statistics. For more examples with graphs and tables you can check the [**metrics.ipynb**](metrics.ipynb) file. +The following example table summarizes the performance metrics of **GPT-4o** when solving a total of 2 labs on SQL Injection, using 3 different CAI Agents. -| prompt | model | avg_turns | avg_active_seconds | avg_idle_seconds | avg_total_seconds | avg_prompt_tokens | avg_completion_tokens | avg_total_tokens | avg_interaction_costs | avg_total_assistant_messages | avg_total_assistant_tools | -|------------------|------------------------|-----------|--------------------|------------------|-------------------|-------------------|-----------------------|------------------|-----------------------|-----------------------------|---------------------------| -| chain-of-thought | deepseek-deepseek-chat | 2.7 | 645.5 | 149.9 | 795.5 | 23578.5 | 1674.0 | 25252.5 | 0.0 | 2.7 | 1.7 | -| chain-of-thought | openai-gpt-4o | 1.2 | 70.0 | 150.9 | 220.9 | 8774.1 | 1034.3 | 9808.5 | 0.0 | 1.1 | 0.2 | -| few-shot | deepseek-deepseek-chat | 2.1 | 668.7 | 88.6 | 757.3 | 24301.0 | 1779.3 | 26080.3 | 0.0 | 2.2 | 1.2 | -| few-shot | openai-gpt-4o | 1.9 | 167.9 | 222.3 | 390.1 | 24134.3 | 780.7 | 24914.9 | 0.0 | 1.3 | 0.9 | -| zero-shot | deepseek-deepseek-chat | 2.7 | 634.1 | 209.0 | 843.1 | 16071.9 | 1392.5 | 17464.3 | 0.0 | 2.7 | 1.7 | -| zero-shot | openai-gpt-4o | 2.9 | 812.9 | 163.4 | 976.3 | 23446.7 | 872.8 | 24319.5 | 0.0 | 1.8 | 2.1 | +The results are broken down by different CAI agents and include interaction times, token usage, and assistant behavior statistics. For more examples with graphs and tables you can check the [**metrics.ipynb**](metrics.ipynb) file. + +| agent | section | model | avg_turns | avg_active_seconds | avg_idle_seconds | avg_total_seconds | avg_prompt_tokens | avg_completion_tokens | avg_total_tokens | avg_interaction_costs | avg_total_assistant_messages | avg_total_assistant_tools | total_interrupted | total_not_solved | total_solved | +|:------------------|:--------------|:--------------|------------:|---------------------:|-------------------:|--------------------:|--------------------:|------------------------:|-------------------:|------------------------:|-------------------------------:|----------------------------:|--------------------:|-------------------:|---------------:| +| blueteam_agent | sql-injection | openai-gpt-4o | 1 | 32.5 | 86 | 118.5 | 2251.5 | 548.5 | 2800 | 0 | 1 | 0 | 0 | 2 | 0 | +| bug_bounter_agent | sql-injection | openai-gpt-4o | 4.5 | 253.5 | 259 | 512.5 | 41288 | 631.5 | 41919.5 | 0 | 2 | 3.5 | 0 | 1 | 1 | +| redteam_agent | sql-injection | openai-gpt-4o | 2 | 64 | 292 | 356 | 9930.5 | 691 | 10621.5 | 0 | 1.5 | 1 | 0 | 0 | 2 | ## Portswigger Web Academy labs -This project allows you to perform evaluations with any of the following labs: +This project allows you to perform evaluations with any of the following labs. For more details, check the [**labs.json**](utils/labs.json). | Section | Lab Title | URL | |---------|-----------|-----| diff --git a/benchmarks/Prompt2PwnBench/server.py b/benchmarks/Agent2PwnBench/main.py similarity index 74% rename from benchmarks/Prompt2PwnBench/server.py rename to benchmarks/Agent2PwnBench/main.py index e7499923..a686e9f1 100644 --- a/benchmarks/Prompt2PwnBench/server.py +++ b/benchmarks/Agent2PwnBench/main.py @@ -1,58 +1,41 @@ -from cai.sdk.agents import Agent, Runner, gen_trace_id, trace, OpenAIChatCompletionsModel -from cai.sdk.agents.mcp import MCPServer, MCPServerSse -from cai.sdk.agents.model_settings import ModelSettings - -#cai tools -from cai.tools.reconnaissance.generic_linux_command import generic_linux_command -from cai.tools.web.search_web import make_web_search_with_explanation -from cai.tools.reconnaissance.exec_code import execute_code -from cai.tools.command_and_control.sshpass import run_ssh_command_with_credentials -from cai.tools.reconnaissance.shodan import shodan_search, shodan_host_info -from cai.tools.web.search_web import make_google_search -from cai.tools.misc.reasoning import think -from cai.tools.web.google_search import google_search -from cai.tools.network.capture_traffic import capture_remote_traffic,remote_capture_session - - #Portswigger web scraper bot import utils.portswiggerbot as pb + #Other utilities import utils.helpers as hp -from openai import AsyncOpenAI import asyncio import os import yaml import sys import shutil from datetime import datetime -import subprocess -import time -from typing import Any from dotenv import load_dotenv -load_dotenv(override=True) +load_dotenv('.env') + + +from cai.sdk.agents import Agent, Runner, gen_trace_id, trace, OpenAIChatCompletionsModel +from cai.sdk.agents.mcp import MCPServer, MCPServerSse +from cai.sdk.agents.model_settings import ModelSettings + +#cai agents +from cai.agents.blue_teamer import blueteam_agent +from cai.agents.bug_bounter import bug_bounter_agent +from cai.agents.red_teamer import redteam_agent #MAIN PARAMETERS USERNAME = os.getenv("PORTSWIGGER_USERNAME") PASSWORD = os.getenv("PORTSWIGGER_PASSWORD") MODEL = os.getenv("CAI_MODEL") -SERVER_URL = "http://127.0.0.1:9876/sse" -REMOVE_LABS_WITH_EXPLOIT_SERVER = True # Set to True to remove labs in portswigger that require an additional exploit server to be solved +SERVER_URL = os.getenv("BURPSUITE_SERVER_URL") +SECTION = os.getenv("PORTSWIGGER_SECTION") +N_LABS = int(os.getenv("NUMBER_OF_LABS")) +AGENT = os.getenv("CAI_AGENT") + +REMOVE_LABS_WITH_EXPLOIT_SERVER = False # Set to True to remove labs in portswigger that require an additional exploit server to be solved -SECTION = "sql-injection" - - # Change this to the desired section or type of portswigger lab -N_LABS = 10 # If you want to test all the labs in the section, change this to -1 - -AGENT = "webbounty" -PROMPT_TYPE = "zero-shot" # Change this to the desired prompt method -TOOLS = [generic_linux_command, - execute_code] - - -#function to save command terminal logs def setup_tee_logging(log_dir="terminal_output",model="openai/gpt-4o",log_name_prefix="console_log"): """ Save command line outputs of the experiments in folder. @@ -79,8 +62,6 @@ def setup_tee_logging(log_dir="terminal_output",model="openai/gpt-4o",log_name_p return log_path - -#function to run the agent async def run_agent(agent,user_prompt): """ Run the agent with the provided user prompt. @@ -93,7 +74,6 @@ async def run_agent(agent,user_prompt): response = await Runner.run(agent, user_prompt) return response - def delete_files(folder="logs"): """ Delete all files in the specified folder. @@ -110,15 +90,13 @@ def delete_files(folder="logs"): os.remove(file_path) print(f"Deleted file: {filename}") - -def create_folder_and_move_logs(lab,section,agent,prompt_type,model,lab_status): +def create_folder_and_move_logs(lab,section,agent,model,lab_status): """ Create a folder structure based on the lab information and move the cai logs to that folder. Args: lab (dict): The lab information containing title, url, and other metadata. section (str): The section or type of lab. agent (str): The name of the agent used. - prompt_type (str): The type of prompt used. model (str): The model used for the agent. lab_status (str): The status of the lab (e.g., solved, not-solved, interrupted, unknown). """ @@ -127,15 +105,12 @@ def create_folder_and_move_logs(lab,section,agent,prompt_type,model,lab_status): model = model.replace("/","-") lab_name = lab['url'].split("/")[-1] if lab_name == 'lab-html-context-nothing-encoded': - lab_name = f"{lab_name}-{lab['url'].split("/")[-2]}" - - - + lab_name = f"{lab_name}-{lab['url'].split("/")[-2]}" if lab_status: lab_status = lab_status.lower().replace(" ","-") - destination_folder = os.path.join("results",model, agent, prompt_type, section, lab_status, lab_name) + destination_folder = os.path.join("results",model, agent, section, lab_status, lab_name) else: - destination_folder = os.path.join("results",model, agent, prompt_type, section, lab_name) + destination_folder = os.path.join("results",model, agent, section, lab_name) os.makedirs(destination_folder, exist_ok=True) @@ -145,6 +120,23 @@ def create_folder_and_move_logs(lab,section,agent,prompt_type,model,lab_status): destination_file = os.path.join(destination_folder, logs_file) shutil.move(source_file, destination_file) +def get_agent(agent_name: str): + """ + Get the agent class based on the agent name. + Args: + agent_name (str): The name of the agent. + Returns: + agent_class (class): The class of the agent. + """ + if agent_name == "blueteam_agent": + return blueteam_agent + elif agent_name == "bug_bounter_agent": + return bug_bounter_agent + elif agent_name == "redteam_agent": + return redteam_agent + else: + raise ValueError(f"Unknown agent name: {agent_name}") + async def run(mcp_server: MCPServer): @@ -153,7 +145,7 @@ async def run(mcp_server: MCPServer): #configuration to write terminal output to a file log_file_path = "terminal_output" - log_name_prefix = f"{SECTION}.{AGENT}.{PROMPT_TYPE}" + log_name_prefix = f"{SECTION}.{AGENT}" _ = setup_tee_logging(log_file_path, model= MODEL, log_name_prefix=log_name_prefix) print(f"[INFO] Console output is being saved to: {log_file_path}") @@ -164,23 +156,12 @@ async def run(mcp_server: MCPServer): with open('prompts.yml', 'r') as f: prompts = yaml.load(f, Loader=yaml.SafeLoader) - #Create CAI Agent - custom_agent = Agent( - name="Custom Agent", - instructions=prompts['prompts']['system'][AGENT], - description="""Agent for AI project to test models with different prompts types on PortSwigger Labs.""", - tools=TOOLS, - mcp_servers=[mcp_server], - model_settings=ModelSettings(tool_choice="required"), - model=OpenAIChatCompletionsModel( - model=os.getenv('CAI_MODEL', "openai/gpt-4o"), - openai_client=AsyncOpenAI(), - ) - ) + #Set agent + custom_agent = get_agent(AGENT) print(f"Getting PortSwigger Labs of {SECTION}.....") #Initialize bot - bot = pb.Bot(headless=False) + bot = pb.Bot(headless=True) #Login to Portswigger Academy bot.login(USERNAME,PASSWORD) @@ -188,7 +169,6 @@ async def run(mcp_server: MCPServer): #Get labs urls by section or type of lab topics = bot.choose_topic(SECTION) - #Get lab metadata #To remove labs that interact with an internal exploit server to be solved if REMOVE_LABS_WITH_EXPLOIT_SERVER: @@ -216,7 +196,7 @@ async def run(mcp_server: MCPServer): print(f"\n####################################") print(f"Trying to hack Lab: {lab['title']}") - user_prompt = prompts['prompts']['user'][PROMPT_TYPE][AGENT] + user_prompt = prompts['prompts']['user'][AGENT] user_prompt = (user_prompt .replace("**LAB_TITLE**", lab['title']) .replace("**LAB_DESCRIPTION**", lab['description']) @@ -249,7 +229,7 @@ async def run(mcp_server: MCPServer): print(f"Lab Status: {LAB_STATUS}") - create_folder_and_move_logs(lab, SECTION, AGENT, PROMPT_TYPE, MODEL, LAB_STATUS) + create_folder_and_move_logs(lab, SECTION, AGENT, MODEL, LAB_STATUS) #delete logs files and close bot browser diff --git a/benchmarks/Agent2PwnBench/metrics.ipynb b/benchmarks/Agent2PwnBench/metrics.ipynb new file mode 100755 index 00000000..729640b4 --- /dev/null +++ b/benchmarks/Agent2PwnBench/metrics.ipynb @@ -0,0 +1,491 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "9edb6753", + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd \n", + "import json\n", + "from collections import defaultdict\n", + "import os\n", + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7cebb746", + "metadata": {}, + "outputs": [], + "source": [ + "#UTIL Funcitons to calculate metrics\n", + "\n", + "def read_results(main_folder=\"results\",model=None):\n", + " \"\"\"\n", + " Reads the results from the specified main folder and organizes them into a nested dictionary structure.\n", + " \n", + " Args:\n", + " main_folder (str): The path to the main folder containing the results.\n", + " \n", + " Returns:\n", + " list: list with dictionaries with agents, sections, and their corresponding logs.\n", + " \"\"\"\n", + " results = defaultdict(lambda: defaultdict(dict))\n", + " list_data = []\n", + " agents = os.listdir(os.path.join(main_folder,model))\n", + " for agent in agents:\n", + " sections = os.listdir(os.path.join(main_folder,model,agent))\n", + " for section in sections:\n", + " statuses = os.listdir(os.path.join(main_folder,model,agent,section))\n", + " for status in statuses:\n", + " labs = os.listdir(os.path.join(main_folder,model,agent,section,status))\n", + " for lab in labs:\n", + " try:\n", + " file = os.listdir(os.path.join(main_folder,model,agent,section,status,lab))[0]\n", + " except IndexError:\n", + " print(os.listdir(os.path.join(main_folder,model,agent,section,status,lab)))\n", + " with open(os.path.join(main_folder,model,agent,section,status,lab,file)) as f:\n", + " logs = [json.loads(line) for line in f]\n", + " \n", + " data = {\n", + " 'agent':agent,\n", + " 'section':section,\n", + " 'model':model,\n", + " 'lab title':lab,\n", + " 'status':status,\n", + " 'logs':logs\n", + " } \n", + " list_data.append(data)\n", + " return list_data\n", + "\n", + "\n", + "def get_metrics(labs):\n", + " \"\"\"\n", + " Extracts metrics from the provided list of lab results.\n", + " \n", + " Args:\n", + " labs (list): A list of dictionaries containing lab results, where each dictionary includes logs and metadata.\n", + " \n", + " Returns:\n", + " list: A list of dictionaries containing calculated metrics for each lab\n", + " \"\"\"\n", + " \n", + " results = []\n", + " for lab in labs:\n", + " \n", + " #------- DATA EXTRACTION --------\n", + " completitions = [log for log in lab['logs'] if log.get('object') == 'chat.completion']\n", + " user_messages = [log for log in lab['logs'] if log.get(\"event\") == \"user_message\" ]\n", + " assistant_messages = [log for log in lab['logs'] if log.get(\"event\") == \"assistant_message\" ]\n", + " model_metadata = [log for log in lab['logs'] if \"model\" in log ]\n", + "\n", + " #model\n", + " model = model_metadata[0]['model']\n", + "\n", + " #assistant messages \n", + " assistant_contents = [\n", + " choice['message']['content']\n", + " for co in completitions\n", + " for choice in co['choices']\n", + " ]\n", + "\n", + " #assistant tools\n", + " assistant_tools_calls = [\n", + " tool['function']\n", + " for co in completitions\n", + " for choice in co['choices']\n", + " for tool in choice['message']['tool_calls']\n", + " ]\n", + "\n", + " #finish reason\n", + " finish_reasons = [\n", + " choice['finish_reason']\n", + " for co in completitions\n", + " for choice in co['choices']\n", + " ]\n", + "\n", + " #integration of finish reason, assistant_contents, and assistant_tools_calls\n", + " assistant_outputs = [{\"message\":a, \"finish_reason\":b,\"tool\":c} for a, b, c in zip(assistant_contents, finish_reasons,assistant_tools_calls)]\n", + "\n", + "\n", + " #------- METRICS CALCULATION --------\n", + " #turns\n", + " total_turns = len(user_messages)\n", + "\n", + " #time\n", + " active_seconds = [ac['timing']['active_seconds'] for ac in completitions]\n", + " idle_seconds = [ac['timing']['idle_seconds'] for ac in completitions]\n", + " total_active_seconds = sum(active_seconds)\n", + " total_idle_seconds = sum(idle_seconds) \n", + " total_seconds = total_active_seconds + total_idle_seconds\n", + "\n", + " #tokens\n", + " prompt_tokens = [ac['usage']['prompt_tokens'] for ac in completitions]\n", + " completion_tokens = [ac['usage']['completion_tokens'] for ac in completitions]\n", + " total_prompt_tokens = sum(prompt_tokens)\n", + " total_completion_tokens = sum(completion_tokens)\n", + " total_tokens = total_prompt_tokens + total_completion_tokens\n", + "\n", + " #costs\n", + " interaction_costs = [ac['cost']['interaction_cost'] for ac in completitions]\n", + " total_interaction_costs = sum(interaction_costs)\n", + "\n", + " #assistant outputs\n", + " total_assistant_messages = len([x for x in assistant_contents if x is not None])\n", + "\n", + " #assistant tools\n", + " total_assistant_tools = len([x for x in assistant_tools_calls])\n", + "\n", + " metrics = {\n", + " \"agent\": lab['agent'],\n", + " \"section\": lab['section'],\n", + " \"model\": lab['model'],\n", + " \"lab_title\": lab['lab title'],\n", + " \"status\": lab['status'],\n", + " \"turns\": total_turns,\n", + " \"active_seconds\": total_active_seconds,\n", + " \"idle_seconds\": total_idle_seconds,\n", + " \"total_seconds\": total_seconds,\n", + " \"prompt_tokens\": total_prompt_tokens,\n", + " \"completion_tokens\": total_completion_tokens,\n", + " \"total_tokens\": total_tokens,\n", + " \"interaction_costs\": total_interaction_costs,\n", + " \"total_assistant_messages\": total_assistant_messages,\n", + " \"total_assistant_tools\": total_assistant_tools,\n", + " \"assistant_outputs\": json.dumps(assistant_outputs) \n", + " }\n", + " results.append(metrics)\n", + " return results " + ] + }, + { + "cell_type": "markdown", + "id": "766a3f90", + "metadata": {}, + "source": [ + "# Define the model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "73c2a8b5", + "metadata": {}, + "outputs": [], + "source": [ + "from dotenv import load_dotenv\n", + "load_dotenv('.env')\n", + "\n", + "MODEL = os.getenv(\"CAI_MODEL\").replace('/','-')\n", + "MODEL" + ] + }, + { + "cell_type": "markdown", + "id": "d69e2850", + "metadata": {}, + "source": [ + "