diff --git a/garak/probes/leakreplay.py b/garak/probes/leakreplay.py index 46b57ffc5..3658c933e 100644 --- a/garak/probes/leakreplay.py +++ b/garak/probes/leakreplay.py @@ -64,9 +64,10 @@ def _attempt_prestore_hook(self, attempt: Attempt, seq: int) -> Attempt: def _postprocess_hook(self, attempt: Attempt) -> Attempt: for idx, thread in enumerate(attempt.messages): - attempt.messages[idx][-1]["content"] = re.sub( - "", "", thread[-1]["content"] - ) + if thread[-1]["content"] is not None: + attempt.messages[idx][-1]["content"] = re.sub( + "", "", thread[-1]["content"] + ) return attempt diff --git a/pyproject.toml b/pyproject.toml index e5cd675b0..caa6f2bbf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ authors = [ { name = "Shine-afk" }, { name = "Rafael Sandroni" }, { name = "Eric Hacker" }, + { name = "Blessed Uyo" }, ] license = { file = "LICENSE" } description = "LLM vulnerability scanner" diff --git a/tests/probes/test_probes_leakreplay.py b/tests/probes/test_probes_leakreplay.py index 6bc77d0d9..76523d00c 100644 --- a/tests/probes/test_probes_leakreplay.py +++ b/tests/probes/test_probes_leakreplay.py @@ -8,6 +8,7 @@ import garak._plugins import garak.attempt import garak.cli +import garak.probes.leakreplay def test_leakreplay_hitlog(): @@ -29,3 +30,10 @@ def test_leakreplay_output_count(): p.generator = g results = p._execute_all([a]) assert len(a.all_outputs) == generations + + +def test_leakreplay_handle_incomplete_attempt(): + p = garak.probes.leakreplay.LiteratureCloze80() + a = garak.attempt.Attempt(prompt="IS THIS BROKEN") + a.outputs = ["", None] + p._postprocess_hook(a)