fix(image_handling): cap in-flight remote media fetches per request

litellm/litellm_core_utils/prompt_templates/image_handling.py

@@ -19,6 +19,7 @@ from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
 from litellm.types.llms.openai import AllMessageValues
 
 MAX_IMGS_IN_MEMORY: Final = 10
+MAX_CONCURRENT_REMOTE_MEDIA_FETCHES: Final = 20
 
 in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY)
 
@@ -216,6 +217,11 @@ def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) ->
     return inlined_message  # pyright: ignore[reportReturnType]  # the same message with its remote parts inlined
 
 
+async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str:
+    async with in_flight:
+        return await async_convert_url_to_base64(url)
+
+
 async def async_inline_remote_media(
     messages: list[AllMessageValues],  # mutable-ok: every transform_request takes list[AllMessageValues]
     skip_url_prefixes: tuple[str, ...] = (),
@@ -230,6 +236,7 @@ async def async_inline_remote_media(
     )
     if not remote_urls:
         return messages
-    data_urls: Final = await asyncio.gather(*(async_convert_url_to_base64(url) for url in remote_urls))
+    in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES)
+    data_urls: Final = await asyncio.gather(*(_fetch_data_url(url, in_flight) for url in remote_urls))
     inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True)))
     return [_inline_message(message, inlined) for message in messages]  # mutable-ok: transform_request takes a list

tests/test_litellm/litellm_core_utils/test_image_handling.py

@@ -1,3 +1,4 @@
+import asyncio
 import copy
 import uuid
 from unittest.mock import patch
@@ -9,6 +10,7 @@ import litellm
 from litellm import constants
 from litellm.litellm_core_utils.prompt_templates import image_handling
 from litellm.litellm_core_utils.prompt_templates.image_handling import (
+    MAX_CONCURRENT_REMOTE_MEDIA_FETCHES,
     async_convert_url_to_base64,
     async_inline_remote_media,
     convert_url_to_base64,
@@ -336,6 +338,26 @@ async def test_async_inline_remote_media_leaves_skipped_url_prefixes_untouched(a
     assert messages == snapshot
 
 
+async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch):
+    in_flight = {"now": 0, "peak": 0}
+
+    async def serve_png_slowly(client, url, **kwargs):
+        in_flight["now"] += 1
+        in_flight["peak"] = max(in_flight["peak"], in_flight["now"])
+        await asyncio.sleep(0.01)
+        in_flight["now"] -= 1
+        return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url))
+
+    monkeypatch.setattr(image_handling, "async_safe_get", serve_png_slowly)
+    urls = [f"https://img.example/{uuid.uuid4()}.png" for _ in range(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES + 5)]
+    messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}} for url in urls]}]
+
+    inlined = await async_inline_remote_media(messages)
+
+    assert in_flight["peak"] == MAX_CONCURRENT_REMOTE_MEDIA_FETCHES
+    assert all(part["image_url"]["url"].startswith("data:image/png;base64,") for part in inlined[0]["content"])
+
+
 async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch):
     messages = [
         {"role": "user", "content": [{"type": "text", "text": "hi"}]},