From 8b71baad75ef497b6f5782ffd771425b787f10bd Mon Sep 17 00:00:00 2001
From: Leon Derczynski <leonderczynski@gmail.com>
Date: Wed, 15 Jan 2025 13:49:47 +0100
Subject: [PATCH 1/2] 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(
-                "</?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
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)

From 70f8d9b7b6d98adba06a7a3020067c11f79d7590 Mon Sep 17 00:00:00 2001
From: Leon Derczynski <leonderczynski@gmail.com>
Date: Wed, 15 Jan 2025 13:55:27 +0100
Subject: [PATCH 2/2] credit for reporter & suggestions

---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

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"