diff --git a/aisploit/redteam/__init__.py b/aisploit/redteam/__init__.py index 9bd5f7f..66f581a 100644 --- a/aisploit/redteam/__init__.py +++ b/aisploit/redteam/__init__.py @@ -1,8 +1,11 @@ from .job import RedTeamJob +from .report import RedTeamReport, RedTeamReportEntry from .task import RedTeamTask, RedTeamEndTokenTask, RedTeamClassifierTask __all__ = [ "RedTeamJob", + "RedTeamReport", + "RedTeamReportEntry", "RedTeamTask", "RedTeamEndTokenTask", "RedTeamClassifierTask", diff --git a/aisploit/redteam/report.py b/aisploit/redteam/report.py index 4be6124..b11c359 100644 --- a/aisploit/redteam/report.py +++ b/aisploit/redteam/report.py @@ -1,4 +1,4 @@ -from typing import List, Optional +from typing import Optional from dataclasses import dataclass from ..core import BaseReport, BasePromptValue, Score @@ -12,18 +12,54 @@ class RedTeamReportEntry: class RedTeamReport(BaseReport[RedTeamReportEntry]): + """ + A report class for storing red team evaluation entries. + """ + def __init__(self, *, run_id: str) -> None: + """ + Initialize the RedTeamReport instance. + + Args: + run_id (str): The ID of the run. + """ super().__init__(run_id=run_id) def add_entry(self, entry: RedTeamReportEntry): + """ + Add an entry to the report. + + Args: + entry (RedTeamReportEntry): The entry to add to the report. + """ self._entries.append(entry) @property def final_score(self) -> Optional[Score]: - last_entry = self._entries[-1] - if last_entry: - return last_entry.score - return None + """ + Get the final score of the report. + + Returns: + Optional[Score]: The final score of the report, or None if no entries exist. + """ + if len(self._entries) == 0: + return None + return self._entries[-1].score + + @property + def final_response(self) -> Optional[str]: + """ + Get the final response of the report. + + Returns: + Optional[str]: The final response of the report, or None if no entries exist. + """ + if len(self._entries) == 0: + return None + return self._entries[-1].response def _ipython_display_(self): + """ + Display the report in IPython environments. + """ print("TODO") diff --git a/tests/redteam/test_report.py b/tests/redteam/test_report.py new file mode 100644 index 0000000..a10b310 --- /dev/null +++ b/tests/redteam/test_report.py @@ -0,0 +1,45 @@ +import pytest + +from aisploit.core import BaseReport, Score +from aisploit.redteam import RedTeamReport, RedTeamReportEntry + + +@pytest.fixture +def red_team_report(): + return RedTeamReport(run_id="test_run") + + +@pytest.fixture +def red_team_report_entry(): + return RedTeamReportEntry( + attempt=1, + prompt="Test prompt", + response="Test response", + score=Score(flagged=True, value=0.8), + ) + + +def test_red_team_report_init(red_team_report): + assert isinstance(red_team_report, BaseReport) + assert red_team_report.run_id == "test_run" + assert len(red_team_report._entries) == 0 + + +def test_red_team_report_add_entry(red_team_report, red_team_report_entry): + red_team_report.add_entry(red_team_report_entry) + assert len(red_team_report._entries) == 1 + assert red_team_report._entries[0] == red_team_report_entry + + +def test_red_team_report_final_score(red_team_report, red_team_report_entry): + assert red_team_report.final_score is None + + red_team_report.add_entry(red_team_report_entry) + assert red_team_report.final_score == red_team_report_entry.score + + +def test_red_team_report_final_response(red_team_report, red_team_report_entry): + assert red_team_report.final_response is None + + red_team_report.add_entry(red_team_report_entry) + assert red_team_report.final_response == red_team_report_entry.response