diff --git a/pycodeloop/providers/_shapes.py b/pycodeloop/providers/_shapes.py index 7f800ad..c4e2e1e 100644 --- a/pycodeloop/providers/_shapes.py +++ b/pycodeloop/providers/_shapes.py @@ -5,12 +5,26 @@ from __future__ import annotations import json +import re from collections.abc import Callable from pycodeloop.core.session import Message RequestBuilder = Callable[[str, "list[Message]", "list[dict]", str], dict] +_DATA_URL = re.compile(r"^data:(?P[\w./+-]+);base64,(?P[^\n]+)$") + + +def _split_image(image: str) -> tuple[str, str]: + """Returns `(mime_type, base64_data)`. Accepts a full `data:; + base64,` URL (mime taken from it) or a bare base64 string + (assumed `image/png`, the historical/backward-compatible behavior + for a caller passing raw base64 straight to `Session.add_user`).""" + match = _DATA_URL.match(image) + if match: + return match.group("mime"), match.group("data") + return "image/png", image + def openai_tool_schema(tools: list[dict]) -> list[dict]: return [ @@ -38,9 +52,9 @@ def to_openai_messages( content = [ { "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{image}"}, + "image_url": {"url": f"data:{mime};base64,{data}"}, } - for image in msg.images + for mime, data in (_split_image(i) for i in msg.images) ] if msg.content: content.append({"type": "text", "text": msg.content}) @@ -108,11 +122,11 @@ def to_anthropic_messages(messages: list[Message]) -> list[dict]: "type": "image", "source": { "type": "base64", - "media_type": "image/png", - "data": image, + "media_type": mime, + "data": data, }, } - for image in msg.images + for mime, data in (_split_image(i) for i in msg.images) ] if msg.content: content.append({"type": "text", "text": msg.content}) diff --git a/tests/providers/test_shapes.py b/tests/providers/test_shapes.py index ad939e9..17364c7 100644 --- a/tests/providers/test_shapes.py +++ b/tests/providers/test_shapes.py @@ -4,6 +4,7 @@ from pycodeloop.core.session import Message from pycodeloop.providers._shapes import ( + _split_image, anthropic_tool_schema, request_builder_from_config, to_anthropic_messages, @@ -11,6 +12,19 @@ ) +class TestSplitImage(unittest.TestCase): + def test_data_url_rejects_embedded_newlines_in_the_base64_payload(self): + """A multiline/PEM-style base64 payload must not silently match + as a clean data URL — falls back to treating the whole string as + opaque `image/png` data instead of a partial, corrupted match.""" + malformed = "data:image/png;base64,abc\ndef" + + mime, data = _split_image(malformed) + + self.assertEqual(mime, "image/png") + self.assertEqual(data, malformed) + + class TestAnthropicMessageBuilding(unittest.TestCase): def test_plain_text_user_message_stays_a_string(self): out = to_anthropic_messages([Message(role="user", content="hi")]) @@ -50,6 +64,22 @@ def test_image_only_message_omits_empty_text_block(self): self.assertEqual(len(out[0]["content"]), 1) self.assertEqual(out[0]["content"][0]["type"], "image") + def test_data_url_image_preserves_its_real_mime_type(self): + out = to_anthropic_messages( + [ + Message( + role="user", + content="", + images=["data:image/jpeg;base64,b64data"], + ) + ] + ) + + self.assertEqual( + out[0]["content"][0]["source"], + {"type": "base64", "media_type": "image/jpeg", "data": "b64data"}, + ) + class TestOpenAIMessageBuilding(unittest.TestCase): def test_plain_text_user_message_stays_a_string(self): @@ -81,6 +111,23 @@ def test_user_message_with_images_becomes_content_blocks(self): }, ) + def test_data_url_image_preserves_its_real_mime_type(self): + out = to_openai_messages( + "sys", + [ + Message( + role="user", + content="", + images=["data:image/webp;base64,b64data"], + ) + ], + ) + + self.assertEqual( + out[-1]["content"][0]["image_url"]["url"], + "data:image/webp;base64,b64data", + ) + def test_tool_call_extra_fields_round_trip_back_to_the_wire(self): """A ToolCall's vendor-specific `extra` (e.g. Gemini's extra_content.google.thought_signature) must be re-emitted as a