From 8b71baad75ef497b6f5782ffd771425b787f10bd Mon Sep 17 00:00:00 2001 From: Leon Derczynski Date: Wed, 15 Jan 2025 13:49:47 +0100 Subject: [PATCH] guard against Nones in message history when rewriting outputs for leakreplay / attempt history --- garak/probes/leakreplay.py | 7 ++++--- tests/probes/test_probes_leakreplay.py | 8 ++++++++ 2 files changed, 12 insertions(+), 3 deletions(-) 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/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)