fix(image_handling): inline url-sourced Anthropic document and image blocks in the async walker

litellm/litellm_core_utils/prompt_templates/image_handling.py

@@ -157,18 +157,41 @@ def _remote_url(candidate: object) -> str | None:
     return candidate if isinstance(candidate, str) and candidate.startswith(_REMOTE_URL_PREFIXES) else None
 
 
-def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | None:
+_ANTHROPIC_MEDIA_BLOCK_TYPES: Final = frozenset({"document", "image"})
+
+
+@dataclass(frozen=True, slots=True)
+class _RemoteSource:
+    part: Mapping[str, object]
+    url: str
+
+
+def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None:
+    if fields.get("type") != "image_url":
+        return None
+    image_url: Final = fields.get("image_url")
+    image_url_fields: Final = _as_mapping(image_url)
+    url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url)
+    return _RemoteImage(fields, image_url_fields, url) if url is not None else None
+
+
+def _parse_remote_file(fields: Mapping[str, object]) -> _RemoteFile | None:
+    file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None
+    url: Final = _remote_url(file.get("file_id")) if file is not None else None
+    return _RemoteFile(fields, file, url) if file is not None and url is not None else None
+
+
+def _parse_remote_source(fields: Mapping[str, object]) -> _RemoteSource | None:
+    source: Final = _as_mapping(fields.get("source")) if fields.get("type") in _ANTHROPIC_MEDIA_BLOCK_TYPES else None
+    url: Final = _remote_url(source.get("url")) if source is not None and source.get("type") == "url" else None
+    return _RemoteSource(fields, url) if url is not None else None
+
+
+def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSource | None:
     fields: Final = _as_mapping(part)
     if fields is None:
         return None
-    if fields.get("type") == "image_url":
-        image_url: Final = fields.get("image_url")
-        image_url_fields: Final = _as_mapping(image_url)
-        url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url)
-        return _RemoteImage(fields, image_url_fields, url) if url is not None else None
-    file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None
-    file_url: Final = _remote_url(file.get("file_id")) if file is not None else None
-    return _RemoteFile(fields, file, file_url) if file is not None and file_url is not None else None
+    return _parse_remote_image(fields) or _parse_remote_file(fields) or _parse_remote_source(fields)
 
 
 _PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"})
@@ -187,12 +210,20 @@ def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mappin
     return {**kept, **_inferred_format(file, url), "file_data": data_url}  # mutable-ok: json-serialized part
 
 
-def _inline(remote: _RemoteImage | _RemoteFile, data_url: str) -> Mapping[str, object]:
+def _base64_source(url: str, data_url: str) -> Mapping[str, str]:
+    fetched_media_type, data = data_url.removeprefix("data:").split(";base64,", 1)
+    media_type: Final = "application/pdf" if url.lower().endswith(".pdf") else fetched_media_type
+    return {"type": "base64", "media_type": media_type, "data": data}  # mutable-ok: json-serialized message part
+
+
+def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) -> Mapping[str, object]:
     match remote:
         case _RemoteImage(part, image_url, _):
             return {**part, "image_url": _inlined_image_url(image_url, data_url)}  # mutable-ok: json-serialized part
         case _RemoteFile(part, file, url):
             return {**part, "file": _inlined_file(file, url, data_url)}  # mutable-ok: json-serialized message part
+        case _RemoteSource(part, url):
+            return {**part, "source": _base64_source(url, data_url)}  # mutable-ok: json-serialized message part
 
 
 def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]:

tests/test_litellm/litellm_core_utils/test_image_handling.py

@@ -289,6 +289,9 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o
                 {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
                 {"type": "file", "file": {"file_id": pdf_url}},
                 {"type": "file", "file": {"file_id": image_url, "format": "image/png"}},
+                {"type": "document", "source": {"type": "url", "url": pdf_url}, "title": "the doc"},
+                {"type": "image", "source": {"type": "url", "url": image_url}},
+                {"type": "document", "source": {"type": "file", "file_id": "file_abc"}},
             ],
         },
     ]
@@ -297,6 +300,7 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o
     inlined = await async_inline_remote_media(messages)
 
     data_url = async_only_image_fetch.data_url
+    base64_png = async_only_image_fetch.base64_png
     assert inlined[0] == {"role": "system", "content": "be terse"}
     assert inlined[1]["content"] == [
         {"type": "text", "text": "what is this?"},
@@ -305,6 +309,13 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o
         {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
         {"type": "file", "file": {"format": "application/pdf", "file_data": data_url}},
         {"type": "file", "file": {"format": "image/png", "file_data": data_url}},
+        {
+            "type": "document",
+            "source": {"type": "base64", "media_type": "application/pdf", "data": base64_png},
+            "title": "the doc",
+        },
+        {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": base64_png}},
+        {"type": "document", "source": {"type": "file", "file_id": "file_abc"}},
     ]
     assert sorted(async_only_image_fetch.fetched) == sorted([image_url, pdf_url])
     assert messages == snapshot

tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py

@@ -763,3 +763,50 @@ async def test_bedrock_invoke_claude_async_completion_inlines_remote_images_off_
     assert async_only_image_fetch.fetched == [image_url]
     assert image_url not in captured["body"]
     assert async_only_image_fetch.base64_png in captured["body"]
+
+
+async def test_bedrock_invoke_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch):
+    pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf"
+    captured = {}
+
+    def handle(request):
+        captured["body"] = json.loads(request.content)
+        return httpx.Response(
+            200,
+            json={
+                "id": "msg_1",
+                "type": "message",
+                "role": "assistant",
+                "model": "us.anthropic.claude-sonnet-5",
+                "content": [{"type": "text", "text": "A lease"}],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 1, "output_tokens": 1},
+            },
+        )
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+    response = await litellm.acompletion(
+        model="bedrock/invoke/us.anthropic.claude-sonnet-5",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What is this document?"},
+                    {"type": "document", "source": {"type": "url", "url": pdf_url}},
+                ],
+            }
+        ],
+        aws_access_key_id="AKIAEXAMPLE",
+        aws_secret_access_key="fake-secret",
+        aws_region_name="us-east-1",
+        client=client,
+    )
+
+    assert response.choices[0].message.content == "A lease"
+    assert async_only_image_fetch.fetched == [pdf_url]
+    assert {
+        "type": "document",
+        "source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png},
+    } in captured["body"]["messages"][0]["content"]

tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py

@@ -1,3 +1,4 @@
+import json
 import uuid
 
 import httpx
@@ -49,3 +50,50 @@ async def test_bedrock_mantle_claude_async_completion_inlines_remote_images_off_
     assert async_only_image_fetch.fetched == [image_url]
     assert image_url not in captured["body"]
     assert async_only_image_fetch.base64_png in captured["body"]
+
+
+async def test_bedrock_mantle_claude_async_completion_inlines_document_url_sources_off_the_event_loop(async_only_image_fetch):
+    pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf"
+    captured = {}
+
+    def handle(request):
+        captured["body"] = json.loads(request.content)
+        return httpx.Response(
+            200,
+            json={
+                "id": "msg_1",
+                "type": "message",
+                "role": "assistant",
+                "model": "us.anthropic.claude-sonnet-5",
+                "content": [{"type": "text", "text": "A lease"}],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 1, "output_tokens": 1},
+            },
+        )
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+    response = await litellm.acompletion(
+        model="bedrock/mantle/us.anthropic.claude-sonnet-5",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What is this document?"},
+                    {"type": "document", "source": {"type": "url", "url": pdf_url}},
+                ],
+            }
+        ],
+        aws_access_key_id="AKIAEXAMPLE",
+        aws_secret_access_key="fake-secret",
+        aws_region_name="us-east-1",
+        client=client,
+    )
+
+    assert response.choices[0].message.content == "A lease"
+    assert async_only_image_fetch.fetched == [pdf_url]
+    assert {
+        "type": "document",
+        "source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png},
+    } in captured["body"]["messages"][0]["content"]