diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 5ae8d224556..4beaabc8b24 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/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 diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index bcf4c7cb6ff..ae9016f2f9e 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/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"}]},