diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 3f151628b6d..a4e36c40b77 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -14,6 +14,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, extract_images_from_message, ) +from litellm.litellm_core_utils.prompt_templates.image_handling import ( + async_inline_remote_media, + inline_remote_image_urls, +) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.ollama.common_utils import prepare_ollama_images @@ -228,6 +232,26 @@ class OllamaChatConfig(BaseConfig): return url + @property + def uses_async_transform_request(self) -> bool: + return True + + async def async_transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: BaseConfig signature + optional_params: dict[str, object], # mutable-ok: BaseConfig signature + litellm_params: dict[str, object], # mutable-ok: BaseConfig signature + headers: dict[str, object], # mutable-ok: BaseConfig signature + ) -> dict[str, object]: # mutable-ok: BaseConfig signature + return self.transform_request( + model=model, + messages=await async_inline_remote_media(messages, should_inline=inline_remote_image_urls), + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + def transform_request( self, model: str, diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 8bf26e45160..afce1642af6 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -1011,3 +1011,66 @@ def test_prepare_ollama_images_downloads_urls_and_passes_through_base64(monkeypa ) assert result == ["alreadybase64data", "downloaded"] + + +_REMOTE_IMAGE_MESSAGES = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}}, + ], + } +] + + +@pytest.mark.asyncio +async def test_async_transform_request_downloads_remote_images_without_blocking(monkeypatch): + """acompletion fetches remote images with async_convert_url_to_base64 + instead of the blocking convert_url_to_base64.""" + + async def fake_async_convert_url_to_base64(url): + return "data:image/png;base64,downloaded" + + def blocking_fetch(url): + raise AssertionError("the async path must not use the blocking fetch") + + monkeypatch.setattr( + "litellm.litellm_core_utils.prompt_templates.image_handling.async_convert_url_to_base64", + fake_async_convert_url_to_base64, + ) + monkeypatch.setattr("litellm.llms.ollama.common_utils.convert_url_to_base64", blocking_fetch) + config = OllamaChatConfig() + + result = await config.async_transform_request( + model="llava", + messages=_REMOTE_IMAGE_MESSAGES, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert config.uses_async_transform_request is True + assert result["messages"][0]["images"] == ["downloaded"] + + +@pytest.mark.asyncio +async def test_remote_image_download_respects_max_image_url_download_size_gate(monkeypatch): + """MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0 turns remote image downloads off on both + the sync and async ollama_chat paths, before any request goes out.""" + monkeypatch.setattr( + "litellm.litellm_core_utils.prompt_templates.image_handling.MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", 0 + ) + config = OllamaChatConfig() + request = { + "model": "llava", + "messages": _REMOTE_IMAGE_MESSAGES, + "optional_params": {}, + "litellm_params": {}, + "headers": {}, + } + + with pytest.raises(litellm.ImageFetchError, match="download is disabled"): + config.transform_request(**request) + with pytest.raises(litellm.ImageFetchError, match="download is disabled"): + await config.async_transform_request(**request)