From 2ec9f755c0b4e2f849d4670eac5fa80890a70f99 Mon Sep 17 00:00:00 2001 From: reachv Date: Wed, 22 Oct 2025 17:23:56 -0700 Subject: [PATCH] Fix cleanup_message keyerror --- tests/test_message_utils_image.py | 35 +++++++++++++++++++++++++++++++ verifiers/utils/message_utils.py | 4 +--- 2 files changed, 36 insertions(+), 3 deletions(-) create mode 100644 tests/test_message_utils_image.py diff --git a/tests/test_message_utils_image.py b/tests/test_message_utils_image.py new file mode 100644 index 00000000..42c75bbc --- /dev/null +++ b/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)) diff --git a/verifiers/utils/message_utils.py b/verifiers/utils/message_utils.py index d1001f10..c2cedee0 100644 --- a/verifiers/utils/message_utils.py +++ b/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)