Fix cleanup_message keyerror

Author: reachv

tests/test_message_utils_image.py

@@ -0,0 +1,35 @@
+# tests/test_message_utils_image.py
+import copy
+
+from verifiers.utils.message_utils import cleanup_message
+
+
+def test_cleanup_message_image_url():
+    msg = {
+        "role": "user",
+        "content": [
+            {
+                "text": "t",
+                "type": "image_url",
+                "image_url": {"url": "https://example.com/image.jpg"},
+            },
+        ],
+    }
+    cleaned = cleanup_message(copy.deepcopy(msg))
+    assert cleaned["role"] == "user"
+    assert len(cleaned["content"]) == 1
+    assert cleaned["content"][0]["type"] == "image_url"
+    assert "text" not in cleaned["content"][0]
+
+
+def test_cleanup_message_no_pop():
+    msg = {
+        "role": "user",
+        "content": [
+            {
+                "type": "image_url",
+                "image_url": {"url": "https://example.com/image.jpg"},
+            },
+        ],
+    }
+    cleanup_message(copy.deepcopy(msg))

verifiers/utils/message_utils.py

@@ -71,9 +71,7 @@ def cleanup_message(message: ChatMessage) -> ChatMessage:
                 new_c.pop("image_url")
                 new_message["content"].append(new_c)
             elif (
-                "image_url" in c_dict
-                and "type" in c_dict
-                and c_dict["type"] == "image_url"
+                "text" in c_dict and "type" in c_dict and c_dict["type"] == "image_url"
             ):
                 new_c.pop("text")
                 new_message["content"].append(new_c)