guard against Nones in message history when rewriting outputs for lea…

…kreplay / attempt history
NVIDIA · Jan 15, 2025 · 8b71baa · 8b71baa
1 parent 0315908
commit 8b71baa
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git 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(
-                "</?name>", "", thread[-1]["content"]
-            )
+            if thread[-1]["content"] is not None:
+                attempt.messages[idx][-1]["content"] = re.sub(
+                    "</?name>", "", thread[-1]["content"]
+                )
         return attempt
 
 

diff --git 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)