Fix redteam report

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",

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")

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