Fix matplotlib warnings and TclError in image generator

agentic_security/probe_data/image_generator.py

@@ -1,8 +1,12 @@
 import base64
 import io
+import re
 
 import httpx
-import matplotlib.pyplot as plt
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
 from cache_to_disk import cache_to_disk
 from tqdm import tqdm
 
@@ -49,6 +53,10 @@ def generate_image(prompt: str, variant: int = 0) -> bytes:
     Returns:
         bytes: The image data in JPG format.
     """
+    # Sanitize prompt: replace non-renderable whitespace characters (tabs, etc.)
+    # with spaces to avoid matplotlib UserWarning about missing glyphs.
+    prompt = re.sub(r"[\t\r\x0b\x0c]", " ", prompt)
+
     # Create a matplotlib figure
     fig, ax = plt.subplots(figsize=(6, 4))
 

agentic_security/probe_data/test_image_generator.py

@@ -39,3 +39,12 @@ def test_generate_image_dataset(mock_generate_image):
     assert isinstance(image_datasets[0], ImageProbeDataset)
     assert image_datasets[0].test_dataset.dataset_name == test_dataset_name
     assert image_datasets[0].image_prompts[0] == b"dummy_image_bytes"
+
+
+def test_generate_image_with_special_whitespace():
+    """Test that prompts with tab and other non-renderable whitespace don't raise warnings."""
+    prompt_with_tabs = "Hello\tWorld\tTest"
+    image_bytes = generate_image(prompt_with_tabs, 0)
+
+    assert isinstance(image_bytes, bytes)
+    assert len(image_bytes) > 0