From 43413f52836c6f24932b77096285a1135d8eaff2 Mon Sep 17 00:00:00 2001
From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com>
Date: Tue, 20 Oct 2026 09:36:42 +0530
Subject: [PATCH] FIX: Handle missing prompt keys gracefully in write_txt

Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com>
---
 pyrit/common/text_helper.py           |  2 +-
 tests/unit/common/test_text_helper.py | 41 ++++++++++++++++++++++++++-
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/pyrit/common/text_helper.py b/pyrit/common/text_helper.py
index d776d6aab4..bf927cb7a0 100644
--- a/pyrit/common/text_helper.py
+++ b/pyrit/common/text_helper.py
@@ -38,7 +38,7 @@ def write_txt(file: IO[Any], examples: list[dict[str, str]]) -> None:
         file: A file-like object opened for writing TXT data.
         examples (list[dict[str, str]]): List of dictionaries to write as TXT.
     """
-    file.write("\n".join([ex["prompt"] for ex in examples]))
+    file.write("\n".join([str(ex.get("prompt") or "") for ex in examples]))
 
 
 def escape_control_characters(text: str) -> str:
diff --git a/tests/unit/common/test_text_helper.py b/tests/unit/common/test_text_helper.py
index 876d8062d2..6d2d4814c9 100644
--- a/tests/unit/common/test_text_helper.py
+++ b/tests/unit/common/test_text_helper.py
@@ -5,7 +5,12 @@
 
 import pytest
 
-from pyrit.common.text_helper import escape_control_characters, is_non_empty_string, read_txt
+from pyrit.common.text_helper import (
+    escape_control_characters,
+    is_non_empty_string,
+    read_txt,
+    write_txt,
+)
 
 
 def test_read_txt_ignores_blank_lines():
@@ -48,3 +53,37 @@ def test_escape_control_characters_leaves_other_text_unchanged(text: str):
 )
 def test_is_non_empty_string_preserves_runtime_validation(value, expected):
     assert is_non_empty_string(value) is expected
+
+
+def test_write_txt_produces_expected_output():
+    file = StringIO()
+    examples = [{"prompt": "first line"}, {"prompt": "second line"}]
+    write_txt(file, examples)
+
+    assert file.getvalue() == "first line\nsecond line"
+
+
+def test_write_txt_empty_examples_writes_nothing():
+    file = StringIO()
+    write_txt(file, [])
+
+    assert file.getvalue() == ""
+
+
+def test_write_then_read_roundtrip():
+    examples = [{"prompt": "first prompt"}, {"prompt": "second prompt"}]
+    file = StringIO()
+    write_txt(file, examples)
+    file.seek(0)
+
+    assert read_txt(file) == examples
+
+
+def test_write_txt_handles_missing_prompt_key_and_none():
+    file = StringIO()
+    examples = [{"prompt": "first"}, {}, {"prompt": None}, {"prompt": "second"}]
+    write_txt(file, examples)
+
+    assert file.getvalue() == "first\n\n\nsecond"
+    file.seek(0)
+    assert read_txt(file) == [{"prompt": "first"}, {"prompt": "second"}]