mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
fix(ollama): fetch remote images asynchronously on ollama_chat acompletion
Override async_transform_request so acompletion inlines remote image_url parts with async_inline_remote_media, the same path the ollama completion config uses, instead of the blocking convert_url_to_base64. The sync path is kept for completion(). Both paths keep honouring MAX_IMAGE_URL_DOWNLOAD_SIZE_MB, where 0 turns remote downloads off.
This commit is contained in:
parent
4fd719659e
commit
c75b596b84
2 changed files with 87 additions and 0 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue