From 199b44a475719500915dbd9059c0ea5dfeeae782 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:48:56 -0700 Subject: [PATCH 1/8] fix(async): move remote image fetches off the event loop for Snowflake, Bedrock invoke Claude, Mantle and Gemini --- .../prompt_templates/image_handling.py | 103 ++++++++++ litellm/llms/base_llm/chat/transformation.py | 4 + .../anthropic_claude3_transformation.py | 61 +----- .../bedrock/chat/mantle/transformation.py | 15 +- litellm/llms/custom_httpx/llm_http_handler.py | 194 ++++++++++-------- litellm/llms/snowflake/chat/transformation.py | 16 ++ .../llms/vertex_ai/gemini/transformation.py | 9 +- litellm/main.py | 8 +- tests/test_litellm/conftest.py | 35 ++++ .../litellm_core_utils/test_image_handling.py | 62 ++++++ ...ations_anthropic_claude3_transformation.py | 49 +++++ ...test_bedrock_chat_mantle_transformation.py | 51 +++++ .../custom_httpx/test_llm_http_handler.py | 106 +++++++++- .../test_snowflake_chat_transformation.py | 46 +++++ .../vertex_ai/gemini/test_transformation.py | 42 ++++ 15 files changed, 644 insertions(+), 157 deletions(-) create mode 100644 tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 81132fa89a5..1890d4eb682 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -2,7 +2,11 @@ Helper functions to handle images passed in messages """ +import asyncio import base64 +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType from typing import Final from httpx import Response @@ -12,6 +16,7 @@ from litellm import verbose_logger from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB 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 @@ -124,3 +129,101 @@ def convert_url_to_base64(url: str) -> str: raise litellm.ImageFetchError( f"Error: Unable to fetch image from URL after 3 attempts. url={url}", ) + + +_REMOTE_URL_PREFIXES: Final = ("http://", "https://") + + +@dataclass(frozen=True, slots=True) +class _RemoteImage: + part: Mapping[str, object] + image_url: Mapping[str, object] | None + url: str + + +@dataclass(frozen=True, slots=True) +class _RemoteFile: + part: Mapping[str, object] + file: Mapping[str, object] + url: str + + +def _as_mapping(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, Mapping) else None # pyright: ignore[reportUnknownVariableType] # fields are parsed one by one + + +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: + 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 + + +_PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"}) + + +def _inferred_format(file: Mapping[str, object], url: str) -> Mapping[str, str]: + return _PDF_FORMAT if "format" not in file and url.lower().endswith(".pdf") else MappingProxyType({}) + + +def _inlined_image_url(image_url: Mapping[str, object] | None, data_url: str) -> Mapping[str, object] | str: + return {**image_url, "url": data_url} if image_url is not None else data_url # mutable-ok: json-serialized part + + +def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mapping[str, object]: + kept: Final = {k: v for k, v in file.items() if k != "file_id"} # mutable-ok: json-serialized message part + 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]: + 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 + + +def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]: + content: Final = message.get("content") + return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one + + +def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> AllMessageValues: + parts: Final = _content_parts(message) + if not parts: + return message + inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks + _inline(remote, data_urls[remote.url]) if (remote := _parse_remote_part(part)) is not None else part + for part in parts + ] + inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message + return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined + + +async def async_inline_remote_media( + messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues] +) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues] + remote_urls: Final = tuple( + dict.fromkeys( + remote.url + for message in messages + for part in _content_parts(message) + if (remote := _parse_remote_part(part)) is not None + ) + ) + if not remote_urls: + return messages + data_urls: Final = await asyncio.gather(*(async_convert_url_to_base64(url) 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/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index bbe1cc85df1..7bfc87a30d6 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -411,6 +411,10 @@ class BaseConfig(ABC): def has_custom_stream_wrapper(self) -> bool: return False + @property + def uses_async_transform_request(self) -> bool: + return False + @property def supports_stream_param_in_request_body(self) -> bool: """ diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 67720451c00..38f280eef03 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -8,7 +8,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( - async_convert_url_to_base64, + async_inline_remote_media, convert_url_to_base64, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -172,6 +172,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): return _anthropic_request + @property + def uses_async_transform_request(self) -> bool: + return True + async def async_transform_request( self, model: str, @@ -180,26 +184,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): litellm_params: dict, headers: dict, ) -> dict: - _anthropic_request: Final = self._build_bedrock_anthropic_request_base( + return self.transform_request( model=model, - messages=messages, + messages=await async_inline_remote_media(messages), optional_params=optional_params, litellm_params=litellm_params, headers=headers, ) - await self._async_convert_document_url_sources_to_base64(_anthropic_request) - beta_list: Final = self._compute_bedrock_invoke_beta_headers( - model=model, - messages=messages, - optional_params=optional_params, - headers=headers, - ) - if beta_list: - _anthropic_request["anthropic_beta"] = beta_list - - return _anthropic_request - def _build_bedrock_anthropic_request_base( self, model: str, @@ -321,45 +313,6 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): "data": image_chunk["data"], } - async def _async_convert_document_url_sources_to_base64(self, anthropic_request: dict) -> None: - """ - Async version of document URL conversion for async completion paths. - """ - messages: Final = anthropic_request.get("messages") - if not isinstance(messages, list): - return - - for message in messages: - if not isinstance(message, dict): - continue - content = message.get("content") - if not isinstance(content, list): - continue - - for block in content: - if not isinstance(block, dict) or block.get("type") != "document": - continue - source = block.get("source") - if not isinstance(source, dict) or source.get("type") != "url": - continue - source_url = source.get("url") - if not isinstance(source_url, str): - continue - - inferred_format: str | None = None - if source_url.lower().endswith(".pdf"): - inferred_format = "application/pdf" - base64_url = await async_convert_url_to_base64(url=source_url) - image_chunk = convert_to_anthropic_image_obj( - openai_image_url=base64_url, - format=inferred_format, - ) - block["source"] = { - "type": "base64", - "media_type": image_chunk["media_type"], - "data": image_chunk["data"], - } - def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict: """ Convert tool search entries to the format supported by the Bedrock Invoke API. diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index 31fc079c0c9..7e2037c33f1 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -10,6 +10,7 @@ at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth. from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Any, Final +from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -110,21 +111,13 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): litellm_params: dict, headers: dict, ) -> dict: - model_id: Final = model.replace("mantle/", "", 1) - - request: Final = self._build_bedrock_anthropic_request_base( - model=model_id, - messages=messages, + return self.transform_request( + model=model, + messages=await async_inline_remote_media(messages), optional_params=optional_params, litellm_params=litellm_params, headers=headers, ) - await self._async_convert_document_url_sources_to_base64(request) - return self._restore_mantle_body_fields( - request=request, - model_id=model_id, - optional_params=optional_params, - ) @staticmethod def _restore_mantle_body_fields(request: dict, model_id: str, optional_params: dict) -> dict: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f281c249c72..54c21f8d5b3 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -116,6 +116,7 @@ from litellm.types.llms.anthropic_skills import ( Skill, ) from litellm.types.llms.openai import ( + AllMessageValues, CreateBatchRequest, CreateFileRequest, FileContentRequest, @@ -488,7 +489,7 @@ class BaseLLMHTTPHandler: def completion( self, model: str, - messages: list, + messages: list[AllMessageValues], api_base: str | None, custom_llm_provider: str, model_response: ModelResponse, @@ -507,7 +508,7 @@ class BaseLLMHTTPHandler: shared_session: Optional["ClientSession"] = None, ): json_mode: Final[bool] = optional_params.pop("json_mode", False) - extra_body: Final[dict | None] = optional_params.pop("extra_body", None) + extra_body: Final[Mapping[str, object] | None] = optional_params.pop("extra_body", None) provider_config = provider_config or ProviderConfigManager.get_provider_chat_config( model=model, provider=litellm.LlmProviders(custom_llm_provider) @@ -522,14 +523,17 @@ class BaseLLMHTTPHandler: ) # get config from model, custom llm provider - headers = provider_config.validate_environment( - api_key=api_key, - headers=headers or {}, - model=model, - messages=messages, - optional_params=optional_params, - api_base=api_base, - litellm_params=litellm_params, + request_headers: Final = cast( # cast-ok: validate_environment is declared as a bare dict + "dict[str, object]", + provider_config.validate_environment( + api_key=api_key, + headers=headers or {}, + model=model, + messages=messages, + optional_params=optional_params, + api_base=api_base, + litellm_params=litellm_params, + ), ) api_base = provider_config.get_complete_url( @@ -541,93 +545,117 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, ) - data: dict[str, object] = provider_config.transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - headers=headers, - ) - - if extra_body is not None: - data = {**data, **extra_body} - - headers, signed_json_body = provider_config.sign_request( - headers=headers, - optional_params={ - **optional_params, - **_aws_signing_overrides(optional_params, litellm_params), - }, - request_data=data, - api_base=api_base, - api_key=api_key, - stream=stream, - fake_stream=fake_stream, - model=model, - ) - - ## LOGGING - logging_obj.pre_call( - input=messages, - api_key=api_key, - additional_args={ - "complete_input_dict": data, - "api_base": api_base, - "headers": headers, - }, - ) - - # Check if stream was converted for WebSearch interception - # This is set by the async_pre_request_hook in WebSearchInterceptionLogger - if litellm_params.get("_websearch_interception_converted_stream", False): - logging_obj.model_call_details["websearch_interception_converted_stream"] = True - - if acompletion is True: - if stream is True: - data = self._add_stream_param_to_request_body( - data=data, - provider_config=provider_config, + def sign_and_log( + transformed: dict[str, object], # mutable-ok: async_completion takes dict + ) -> tuple[dict[str, object], dict[str, object], bytes | None]: # mutable-ok: async_completion takes dict + data: Final = {**transformed, **extra_body} if extra_body is not None else transformed + signed: Final = cast( # cast-ok: sign_request is declared as a bare dict + "tuple[dict[str, object], bytes | None]", + provider_config.sign_request( + headers=request_headers, + optional_params={ + **optional_params, + **_aws_signing_overrides(optional_params, litellm_params), + }, + request_data=data, + api_base=api_base, + api_key=api_key, + stream=stream, fake_stream=fake_stream, - ) + model=model, + ), + ) + logging_obj.pre_call( + input=messages, + api_key=api_key, + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": signed[0], + }, + ) + if litellm_params.get("_websearch_interception_converted_stream", False): + logging_obj.model_call_details["websearch_interception_converted_stream"] = True + return data, signed[0], signed[1] + + def dispatch_async( + data: dict[str, object], # mutable-ok: async_completion takes dict + signed_headers: dict[str, object], # mutable-ok: async_completion takes dict + signed_json_body: bytes | None, + ): + async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None + if stream is True: return self.acompletion_stream_function( model=model, messages=messages, api_base=api_base, - headers=headers, + headers=signed_headers, custom_llm_provider=custom_llm_provider, provider_config=provider_config, timeout=timeout, logging_obj=logging_obj, - data=data, + data=self._add_stream_param_to_request_body( + data=data, + provider_config=provider_config, + fake_stream=fake_stream, + ), fake_stream=fake_stream, - client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), + client=async_client, litellm_params=litellm_params, json_mode=json_mode, optional_params=optional_params, signed_json_body=signed_json_body, ) + return self.async_completion( + custom_llm_provider=custom_llm_provider, + provider_config=provider_config, + api_base=api_base, + headers=signed_headers, + data=data, + timeout=timeout, + model=model, + model_response=model_response, + logging_obj=logging_obj, + api_key=api_key, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + client=async_client, + json_mode=json_mode, + signed_json_body=signed_json_body, + shared_session=shared_session, + ) - else: - return self.async_completion( - custom_llm_provider=custom_llm_provider, - provider_config=provider_config, - api_base=api_base, - headers=headers, - data=data, - timeout=timeout, - model=model, - model_response=model_response, - logging_obj=logging_obj, - api_key=api_key, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - encoding=encoding, - client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), - json_mode=json_mode, - signed_json_body=signed_json_body, - shared_session=shared_session, + if acompletion is True and provider_config.uses_async_transform_request: + + async def transform_then_dispatch(): + transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict + "dict[str, object]", + await provider_config.async_transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=request_headers, + ), ) + return await dispatch_async(*sign_and_log(transformed)) + + return transform_then_dispatch() + + data, signed_headers, signed_json_body = sign_and_log( + provider_config.transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=request_headers, + ) + ) + + if acompletion is True: + return dispatch_async(data, signed_headers, signed_json_body) if stream is True: data = self._add_stream_param_to_request_body( @@ -641,7 +669,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, messages=messages, @@ -651,7 +679,7 @@ class BaseLLMHTTPHandler: completion_stream, headers = self.make_sync_call( provider_config=provider_config, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, original_data=data, @@ -684,7 +712,7 @@ class BaseLLMHTTPHandler: sync_httpx_client=sync_httpx_client, provider_config=provider_config, api_base=api_base, - headers=headers, + headers=signed_headers, data=data, signed_json_body=signed_json_body, timeout=timeout, diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index c64fc583edc..f65b0876202 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( create_anthropic_image_param, select_anthropic_content_block_type_for_file, ) +from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload @@ -421,6 +422,21 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body) return self._transform_request_openai(model, messages, optional_params, stream, extra_body) + @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 + inlined_messages: Final = await async_inline_remote_media(messages) if _is_claude_model(model) else messages + return self.transform_request(model, inlined_messages, optional_params, litellm_params, headers) + def _transform_request_openai( self, model: str, diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index e2d62be6a69..c9480f07150 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -27,6 +27,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, response_schema_prompt, ) +from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.files import ( @@ -1348,13 +1349,15 @@ async def async_transform_request_body( vertex_auth_header=vertex_auth_header, ) - if _openai_messages_may_need_sync_gcs_metadata_fetch(messages): + inlined_messages: Final = await async_inline_remote_media(messages) if custom_llm_provider == "gemini" else messages + + if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): # _transform_request_body may issue a sync httpx.get (up to 5s timeout) # via _get_gcs_object_content_type to fetch GCS object metadata. Run the # whole sync transformation on a worker thread so it does not block the # async event loop. return await asyncify(_transform_request_body)( - messages=messages, + messages=inlined_messages, model=model, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, @@ -1363,7 +1366,7 @@ async def async_transform_request_body( ) return _transform_request_body( - messages=messages, + messages=inlined_messages, model=model, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, diff --git a/litellm/main.py b/litellm/main.py index 2929790f2bd..9970c203e03 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4474,7 +4474,7 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client = _dispatch_client_http(ctx) + injected_client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4486,11 +4486,11 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR shared_session: Final = ctx.shared_session stream: Final = ctx.stream timeout: Final = ctx.timeout + client: Final = ( + injected_client if injected_client is not None else (HTTPHandler(timeout=timeout) if stream is False else None) + ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible try: - client = ( - HTTPHandler(timeout=timeout) if stream is False else None - ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible response: Final = base_llm_http_handler.completion( model=model, messages=messages, diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index 62c95cb100b..aa8ba168ecc 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -7,9 +7,12 @@ # 4. Added proper cleanup in fixtures # 5. Added worker-specific isolation for parallel execution +import base64 import importlib import os from pathlib import Path +from types import SimpleNamespace +import httpx import pytest import asyncio @@ -595,3 +598,35 @@ def pytest_sessionfinish(session, exitstatus): _close_handler_if_needed(getattr(litellm, "aclient", None)) _close_handler_if_needed(getattr(litellm, "client", None)) _run_coroutine_if_needed(close_litellm_async_clients()) + + +ONE_PIXEL_PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) + + +@pytest.fixture +def async_only_image_fetch(monkeypatch): + from litellm.litellm_core_utils.prompt_templates import image_handling + + fetch = SimpleNamespace( + fetched=[], + base64_png=base64.b64encode(ONE_PIXEL_PNG).decode(), + data_url="data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode(), + ) + + def forbid_sync_fetch(client, url, **kwargs): + raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}") + + async def serve_png(client, url, **kwargs): + fetch.fetched.append(url) + return httpx.Response( + 200, + content=ONE_PIXEL_PNG, + headers={"content-type": "image/png"}, + request=httpx.Request("GET", url), + ) + + monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(image_handling, "async_safe_get", serve_png) + return fetch 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 893472d63ae..5b6d403fa21 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,5 @@ +import copy +import uuid from unittest.mock import patch import pytest @@ -8,6 +10,7 @@ from litellm import constants from litellm.litellm_core_utils.prompt_templates import image_handling from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, + async_inline_remote_media, convert_url_to_base64, ) @@ -268,3 +271,62 @@ def test_image_size_limit_disabled(monkeypatch): assert "Image URL download is disabled" in str(excinfo.value) assert "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0" in str(excinfo.value) + + +async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf" + messages = [ + {"role": "system", "content": "be terse"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this?"}, + {"type": "image_url", "image_url": {"url": image_url, "detail": "low"}}, + {"type": "image_url", "image_url": image_url}, + {"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"}}, + ], + }, + ] + snapshot = copy.deepcopy(messages) + + inlined = await async_inline_remote_media(messages) + + data_url = async_only_image_fetch.data_url + assert inlined[0] == {"role": "system", "content": "be terse"} + assert inlined[1]["content"] == [ + {"type": "text", "text": "what is this?"}, + {"type": "image_url", "image_url": {"url": data_url, "detail": "low"}}, + {"type": "image_url", "image_url": data_url}, + {"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}}, + ] + assert sorted(async_only_image_fetch.fetched) == sorted([image_url, pdf_url]) + assert messages == snapshot + + +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"}]}, + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}], + }, + ] + + assert await async_inline_remote_media(messages) is messages + assert async_only_image_fetch.fetched == [] + + +async def test_async_inline_remote_media_raises_image_fetch_error_when_the_fetch_fails(monkeypatch): + async def serve_404(client, url, **kwargs): + return Response(404, request=Request("GET", url)) + + monkeypatch.setattr(image_handling, "async_safe_get", serve_404) + url = f"http://img.example/{uuid.uuid4()}.png" + + with pytest.raises(litellm.ImageFetchError, match="Status code: 404"): + await async_inline_remote_media([{"role": "user", "content": [{"type": "image_url", "image_url": url}]}]) diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 0d7573a2536..e808a087e3d 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1,15 +1,19 @@ import asyncio import json +import uuid from unittest.mock import patch +import httpx import pytest # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. +import litellm from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeConfig, ) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler def test_get_supported_params_thinking(): @@ -714,3 +718,48 @@ def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking assert result["thinking"] == {"type": "adaptive"} assert result["output_config"] == {"effort": "high"} + + +async def test_bedrock_invoke_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "Green"}], + "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 colour is this?"}, + {"type": "image_url", "image_url": {"url": image_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 == "Green" + 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"] diff --git a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py new file mode 100644 index 00000000000..9b491d305c7 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -0,0 +1,51 @@ +import uuid + +import httpx + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + +async def test_bedrock_mantle_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "us.anthropic.claude-sonnet-5", + "content": [{"type": "text", "text": "Green"}], + "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 colour is this?"}, + {"type": "image_url", "image_url": {"url": image_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 == "Green" + 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"] diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 023e2d8843f..e86e671b939 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -17,7 +17,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, ) -from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, @@ -30,7 +30,7 @@ from litellm.llms.azure.videos.transformation import AzureVideoConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import TranscriptionResponse +from litellm.types.utils import ModelResponse, TranscriptionResponse _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" @@ -3186,3 +3186,105 @@ async def test_async_container_list_handler_transforms_success_response(): assert [container.id for container in response.data] == ["cntr_a"] assert response.has_more is True + + +class _TransformRecordingConfig(BaseConfig): + def __init__(self, transform_async: bool): + self.transform_async = transform_async + self.transform_calls = [] + + @property + def uses_async_transform_request(self) -> bool: + return self.transform_async + + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, non_default_params, optional_params, model, drop_params): + return optional_params + + def validate_environment( + self, headers, model, messages, optional_params, litellm_params, api_key=None, api_base=None + ): + return {} + + def transform_request(self, model, messages, optional_params, litellm_params, headers): + self.transform_calls.append("sync") + return {"transformed_by": "sync"} + + async def async_transform_request(self, model, messages, optional_params, litellm_params, headers): + self.transform_calls.append("async") + return {"transformed_by": "async"} + + def transform_response( + self, + model, + raw_response, + model_response, + logging_obj, + request_data, + messages, + optional_params, + litellm_params, + encoding, + api_key=None, + json_mode=None, + ): + model_response.choices[0].message.content = raw_response.json()["transformed_by"] + return model_response + + def get_error_class(self, error_message, status_code, headers): + return BaseLLMException(status_code=status_code, message=error_message, headers=headers) + + +def _start_async_completion(config): + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=captured["body"]) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + pending = BaseLLMHTTPHandler().completion( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + api_base="https://llm.example/v1/chat", + custom_llm_provider="openai", + model_response=ModelResponse(), + encoding=None, + logging_obj=Mock(dynamic_success_callbacks=None, model_call_details={}), + optional_params={}, + timeout=10.0, + litellm_params={}, + acompletion=True, + client=client, + provider_config=config, + ) + return pending, captured + + +async def test_completion_awaits_async_transform_request_when_config_opts_in(): + config = _TransformRecordingConfig(transform_async=True) + + pending, captured = _start_async_completion(config) + assert config.transform_calls == [] + + response = await pending + + assert config.transform_calls == ["async"] + assert captured["body"] == {"transformed_by": "async"} + assert response.choices[0].message.content == "async" + + +async def test_completion_keeps_sync_transform_request_before_returning_by_default(): + config = _TransformRecordingConfig(transform_async=False) + + pending, captured = _start_async_completion(config) + assert config.transform_calls == ["sync"] + + response = await pending + + assert config.transform_calls == ["sync"] + assert captured["body"] == {"transformed_by": "sync"} + assert response.choices[0].message.content == "sync" diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py index 25a961c3413..5687a319f06 100644 --- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -6,6 +6,7 @@ Tests tool calling request/response transformations and chat completions import asyncio import os import copy +import uuid import json from typing import Any, Dict, List @@ -945,3 +946,48 @@ class TestSnowflakeChatCompletion: assert len(chunks_received) > 0 content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content) + + +async def test_snowflake_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "Green"}], + "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="snowflake/claude-sonnet-4-6", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + api_key="fake-jwt", + account_id="FAKE-ACCOUNT", + api_base=FAKE_API_BASE, + client=client, + ) + + assert response.choices[0].message.content == "Green" + 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"] diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py index fad310fc5c0..ec445342523 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py @@ -1,6 +1,10 @@ +import uuid +import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.vertex_ai.gemini import transformation from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, @@ -338,3 +342,41 @@ def test_map_function_enterprise_web_search_snake_case(): assert len(result) == 1 assert "enterpriseWebSearch" in result[0] + + +async def test_gemini_ai_studio_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): + image_url = f"http://img.example/{uuid.uuid4()}.png" + captured = {} + + def handle(request): + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "Green"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="gemini/gemini-3.8-flash", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What colour is this?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + api_key="fake-gemini-key", + client=client, + ) + + assert response.choices[0].message.content == "Green" + 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"] From e355203014326462c3e70e6b6c8d37fcc0abe656 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:11:41 -0700 Subject: [PATCH 2/8] fix(gemini): keep Files API refs out of the async remote media walker Gemini AI Studio file URIs under generativelanguage.googleapis.com/v1beta/files/ answer 403 when fetched and must pass through as file_data.file_uri, which the sync transform already did. Give the shared walker a skip_url_prefixes parameter, pass that prefix from the Gemini body builder, and skip it in the AI Studio message transform for both image_url and file parts --- .../prompt_templates/image_handling.py | 12 +++-- litellm/llms/gemini/chat/transformation.py | 14 ++++-- .../llms/vertex_ai/gemini/transformation.py | 9 +++- .../litellm_core_utils/test_image_handling.py | 28 ++++++++++++ .../vertex_ai/gemini/test_transformation.py | 44 +++++++++++++++++++ 5 files changed, 99 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 1890d4eb682..5ae8d224556 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -199,13 +199,18 @@ def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]: return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one +def _inline_part(part: object, data_urls: Mapping[str, str]) -> object: + remote: Final = _parse_remote_part(part) + data_url: Final = data_urls.get(remote.url) if remote is not None else None + return _inline(remote, data_url) if remote is not None and data_url is not None else part + + def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> AllMessageValues: parts: Final = _content_parts(message) if not parts: return message inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks - _inline(remote, data_urls[remote.url]) if (remote := _parse_remote_part(part)) is not None else part - for part in parts + _inline_part(part, data_urls) for part in parts ] inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined @@ -213,13 +218,14 @@ def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> async def async_inline_remote_media( messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues] + skip_url_prefixes: tuple[str, ...] = (), ) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues] remote_urls: Final = tuple( dict.fromkeys( remote.url for message in messages for part in _content_parts(message) - if (remote := _parse_remote_part(part)) is not None + if (remote := _parse_remote_part(part)) is not None and not remote.url.startswith(skip_url_prefixes) ) ) if not remote_urls: diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 1a67b33665b..42c9ef13730 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -12,7 +12,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionFileObject from litellm.types.llms.vertex_ai import ContentType, PartType from litellm.utils import supports_reasoning -from ...vertex_ai.gemini.transformation import _gemini_convert_messages_with_history +from ...vertex_ai.gemini.transformation import GEMINI_FILES_API_URI_PREFIX, _gemini_convert_messages_with_history from ...vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig @@ -127,7 +127,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): if element.get("type") == "image_url": img_element = cast(ChatCompletionImageObject, element) # cast-ok: runtime type tag checked _image_url, format, detail = _image_url_fields(img_element) - if _image_url and "https://" in _image_url: + if ( + _image_url + and "https://" in _image_url + and not _image_url.startswith(GEMINI_FILES_API_URI_PREFIX) + ): image_obj = convert_to_anthropic_image_obj(_image_url, format=format) converted_image_url = convert_generic_image_chunk_to_openai_image_obj(image_obj) if detail is not None: @@ -147,7 +151,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): llm_provider="gemini", ) file_id = _file_field.get("file_id") - if file_id and ("http://" in file_id or "https://" in file_id): + if ( + file_id + and ("http://" in file_id or "https://" in file_id) + and not file_id.startswith(GEMINI_FILES_API_URI_PREFIX) + ): # Convert HTTP/HTTPS file URL to base64 data try: base64_data = convert_url_to_base64(file_id) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index c9480f07150..6d100143e52 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -69,6 +69,7 @@ _GCS_METADATA_VERTEX_BASE: Any | None = None # Shared sync client for GCS JSON API metadata reads so proxy/SSL settings # from litellm's HTTP stack apply (see Greptile review on PR #27278). _GCS_METADATA_HTTP_HANDLER: HTTPHandler | None = None +GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/" _GEMINI_MIME_TYPE_ALIASES: Final[dict[str, str]] = { "image/jpg": "image/jpeg", } @@ -557,7 +558,7 @@ def _process_gemini_media( file_data = FileDataType(mime_type=mime_type, file_uri=image_url) part: PartType = {"file_data": file_data} return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata) - elif image_url.startswith("https://generativelanguage.googleapis.com/v1beta/files/"): + elif image_url.startswith(GEMINI_FILES_API_URI_PREFIX): # Gemini Files API URIs — the file is already uploaded to Google's # servers; pass the URI through as file_data without fetching it. # These URLs return 403 when accessed directly, so we must not try @@ -1349,7 +1350,11 @@ async def async_transform_request_body( vertex_auth_header=vertex_auth_header, ) - inlined_messages: Final = await async_inline_remote_media(messages) if custom_llm_provider == "gemini" else messages + inlined_messages: Final = ( + await async_inline_remote_media(messages, skip_url_prefixes=(GEMINI_FILES_API_URI_PREFIX,)) + if custom_llm_provider == "gemini" + else messages + ) if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): # _transform_request_body may issue a sync httpx.get (up to 5s timeout) 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 5b6d403fa21..bcf4c7cb6ff 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -308,6 +308,34 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o assert messages == snapshot +async def test_async_inline_remote_media_leaves_skipped_url_prefixes_untouched(async_only_image_fetch): + skipped_prefix = "https://generativelanguage.googleapis.com/v1beta/files/" + skipped_file = f"{skipped_prefix}{uuid.uuid4().hex}" + skipped_image = f"{skipped_prefix}{uuid.uuid4().hex}" + fetched_image = f"https://img.example/{uuid.uuid4()}.png" + messages = [ + { + "role": "user", + "content": [ + {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}}, + {"type": "image_url", "image_url": {"url": skipped_image}}, + {"type": "image_url", "image_url": fetched_image}, + ], + } + ] + snapshot = copy.deepcopy(messages) + + inlined = await async_inline_remote_media(messages, skip_url_prefixes=(skipped_prefix,)) + + assert inlined[0]["content"] == [ + {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}}, + {"type": "image_url", "image_url": {"url": skipped_image}}, + {"type": "image_url", "image_url": async_only_image_fetch.data_url}, + ] + assert async_only_image_fetch.fetched == [fetched_image] + assert messages == snapshot + + 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"}]}, diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py index ec445342523..d31254746d4 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py @@ -1,4 +1,5 @@ +import json import uuid import httpx import pytest @@ -380,3 +381,46 @@ async def test_gemini_ai_studio_async_completion_inlines_remote_images_off_the_e 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_gemini_ai_studio_async_completion_passes_files_api_uris_through_unfetched(async_only_image_fetch): + files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + files_api_image = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + captured = {} + + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "A report"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.acompletion( + model="gemini/gemini-3.8-flash", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize these"}, + {"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}}, + {"type": "image_url", "image_url": {"url": files_api_image, "format": "image/png"}}, + ], + } + ], + api_key="fake-gemini-key", + client=client, + ) + + assert response.choices[0].message.content == "A report" + assert async_only_image_fetch.fetched == [] + file_parts = [part["file_data"] for part in captured["body"]["contents"][0]["parts"] if "file_data" in part] + assert file_parts == [ + {"mime_type": "application/pdf", "file_uri": files_api_pdf}, + {"mime_type": "image/png", "file_uri": files_api_image}, + ] From 08bb7de868f0d740c937585835333ac226faccad Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:28:17 -0700 Subject: [PATCH 3/8] fix(image_handling): cap in-flight remote media fetches per request --- .../prompt_templates/image_handling.py | 9 +++++++- .../litellm_core_utils/test_image_handling.py | 22 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) 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"}]}, From 567915aeee82953c4127456654f3a63ec9862970 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:51:14 -0700 Subject: [PATCH 4/8] fix(image_handling): inline url-sourced Anthropic document and image blocks in the async walker --- .../prompt_templates/image_handling.py | 51 +++++++++++++++---- .../litellm_core_utils/test_image_handling.py | 11 ++++ ...ations_anthropic_claude3_transformation.py | 47 +++++++++++++++++ ...test_bedrock_chat_mantle_transformation.py | 48 +++++++++++++++++ 4 files changed, 147 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 4beaabc8b24..55c01b3b1cb 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/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, ...]: 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 ae9016f2f9e..106bf90213b 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/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 diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index e808a087e3d..cbf160c451f 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/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"] diff --git a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py index 9b491d305c7..a8448f5fa7a 100644 --- a/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py +++ b/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"] From f09992b73082e37d7337402954387def48fc7420 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:54:40 -0700 Subject: [PATCH 5/8] fix(image_handling): keep DNS, signing, and Vertex Gemini fetches off the event loop The SSRF check in async_safe_get resolved DNS on the event loop and a blocked address was retried three times; validate_url now runs in a thread and an SSRFError fails the fetch on the first attempt in both fetchers. The shared HTTP handler signed the request and ran pre_call logging on the loop after an async transform; both now run in a thread. Vertex AI Gemini still fetched http:// images and https images without an inferrable mime type with the sync converter inside its async body builder; the walker takes a should_inline predicate and Vertex AI inlines exactly those URLs, leaving https images with a known mime type and Files API refs to Google. When one download fails the other in-flight downloads for that request are now cancelled instead of finishing in the background --- .../prompt_templates/image_handling.py | 73 ++++++++-- litellm/litellm_core_utils/url_utils.py | 3 +- litellm/llms/custom_httpx/llm_http_handler.py | 2 +- .../llms/vertex_ai/gemini/transformation.py | 26 +++- .../litellm_core_utils/test_image_handling.py | 135 ++++++++++++++++-- .../litellm_core_utils/test_url_utils.py | 35 +++++ .../custom_httpx/test_llm_http_handler.py | 28 +++- .../vertex_ai/gemini/test_transformation.py | 51 +++++++ 8 files changed, 318 insertions(+), 35 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 55c01b3b1cb..dcad7776b58 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -4,7 +4,7 @@ Helper functions to handle images passed in messages import asyncio import base64 -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from types import MappingProxyType from typing import Final @@ -15,7 +15,7 @@ import litellm from litellm import verbose_logger from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB -from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 @@ -99,6 +99,8 @@ async def async_convert_url_to_base64(url: str) -> str: return _process_image_response(response, url) except litellm.ImageFetchError: raise + except SSRFError as e: + raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e except Exception: pass raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}") @@ -125,6 +127,8 @@ def convert_url_to_base64(url: str) -> str: return _process_image_response(response, url) except litellm.ImageFetchError: raise + except SSRFError as e: + raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e except Exception as e: verbose_logger.exception(e) raise litellm.ImageFetchError( @@ -163,9 +167,23 @@ _ANTHROPIC_MEDIA_BLOCK_TYPES: Final = frozenset({"document", "image"}) @dataclass(frozen=True, slots=True) class _RemoteSource: part: Mapping[str, object] + source: Mapping[str, object] url: str +@dataclass(frozen=True, slots=True) +class RemoteMedia: + url: str + fields: Mapping[str, object] + + +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def inline_every_remote_url(_media: RemoteMedia) -> bool: + return True + + def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None: if fields.get("type") != "image_url": return None @@ -184,7 +202,7 @@ def _parse_remote_file(fields: Mapping[str, object]) -> _RemoteFile | 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 + return _RemoteSource(fields, source, url) if source is not None and url is not None else None def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSource | None: @@ -194,6 +212,16 @@ def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSour return _parse_remote_image(fields) or _parse_remote_file(fields) or _parse_remote_source(fields) +def _remote_media(remote: _RemoteImage | _RemoteFile | _RemoteSource) -> RemoteMedia: + match remote: + case _RemoteImage(_, image_url, url): + return RemoteMedia(url, image_url if image_url is not None else _NO_FIELDS) + case _RemoteFile(_, file, url): + return RemoteMedia(url, file) + case _RemoteSource(_, source, url): + return RemoteMedia(url, source) + + _PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"}) @@ -222,7 +250,7 @@ def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) - 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): + case _RemoteSource(part, _, url): return {**part, "source": _base64_source(url, data_url)} # mutable-ok: json-serialized message part @@ -231,18 +259,22 @@ def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]: return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one -def _inline_part(part: object, data_urls: Mapping[str, str]) -> object: +def _inline_part(part: object, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]) -> object: remote: Final = _parse_remote_part(part) - data_url: Final = data_urls.get(remote.url) if remote is not None else None - return _inline(remote, data_url) if remote is not None and data_url is not None else part + if remote is None or not should_inline(_remote_media(remote)): + return part + data_url: Final = data_urls.get(remote.url) + return _inline(remote, data_url) if data_url is not None else part -def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> AllMessageValues: +def _inline_message( + message: AllMessageValues, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool] +) -> AllMessageValues: parts: Final = _content_parts(message) if not parts: return message inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks - _inline_part(part, data_urls) for part in parts + _inline_part(part, data_urls, should_inline) for part in parts ] inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined @@ -253,21 +285,34 @@ async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str: return await async_convert_url_to_base64(url) +async def _fetch_data_urls(remote_urls: tuple[str, ...]) -> tuple[str, ...]: + in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) + fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight)) for url in remote_urls) + try: + return tuple(await asyncio.gather(*fetches)) + except BaseException: + for fetch in fetches: + fetch.cancel() + await asyncio.gather(*fetches, return_exceptions=True) + raise + + async def async_inline_remote_media( messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues] - skip_url_prefixes: tuple[str, ...] = (), + should_inline: Callable[[RemoteMedia], bool] = inline_every_remote_url, ) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues] remote_urls: Final = tuple( dict.fromkeys( remote.url for message in messages for part in _content_parts(message) - if (remote := _parse_remote_part(part)) is not None and not remote.url.startswith(skip_url_prefixes) + if (remote := _parse_remote_part(part)) is not None and should_inline(_remote_media(remote)) ) ) if not remote_urls: return messages - 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)) + data_urls: Final = await _fetch_data_urls(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 + return [ # mutable-ok: transform_request takes a list + _inline_message(message, inlined, should_inline) for message in messages + ] diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 1e43117933d..fa070a648f5 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -19,6 +19,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): check but still resolve DNS and still rewrite HTTP to the resolved IP. """ +import asyncio import socket from ipaddress import ip_address, ip_network from typing import Any, Final, Protocol @@ -471,7 +472,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response kwargs.pop("follow_redirects", None) headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})} for _ in range(_MAX_REDIRECTS): - validated_url, original_host = validate_url(url) + validated_url, original_host = await asyncio.to_thread(validate_url, url) response = await fetcher.get( validated_url, headers={**headers_view["headers"], "Host": original_host}, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 54c21f8d5b3..d57aee025d5 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -640,7 +640,7 @@ class BaseLLMHTTPHandler: headers=request_headers, ), ) - return await dispatch_async(*sign_and_log(transformed)) + return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed)) return transform_then_dispatch() diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 6d100143e52..13e2238fdf6 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -7,6 +7,7 @@ Why separate file? Make it easy to see how transformation works import json import os import re +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote @@ -27,7 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, response_schema_prompt, ) -from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media +from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, async_inline_remote_media from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.files import ( @@ -1309,6 +1310,23 @@ def sync_transform_request_body( ) +def _explicit_mime_type(fields: Mapping[str, object]) -> str | None: + hint: Final = fields.get("format") or fields.get("mime_type") or fields.get("content_type") + return hint if isinstance(hint, str) else None + + +def _ai_studio_inlines(media: RemoteMedia) -> bool: + return not media.url.startswith(GEMINI_FILES_API_URI_PREFIX) + + +def _vertex_inlines(media: RemoteMedia) -> bool: + if media.url.startswith(GEMINI_FILES_API_URI_PREFIX): + return False + return media.url.startswith("http://") or ( + _explicit_mime_type(media.fields) is None and _get_image_mime_type_from_url(media.url) is None + ) + + async def async_transform_request_body( gemini_api_key: str | None, messages: list[AllMessageValues], @@ -1350,10 +1368,8 @@ async def async_transform_request_body( vertex_auth_header=vertex_auth_header, ) - inlined_messages: Final = ( - await async_inline_remote_media(messages, skip_url_prefixes=(GEMINI_FILES_API_URI_PREFIX,)) - if custom_llm_provider == "gemini" - else messages + inlined_messages: Final = await async_inline_remote_media( + messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines ) if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): 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 106bf90213b..e7b28d0d3a2 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -1,5 +1,6 @@ import asyncio import copy +import time import uuid from unittest.mock import patch @@ -11,10 +12,12 @@ 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, + RemoteMedia, async_convert_url_to_base64, async_inline_remote_media, convert_url_to_base64, ) +from litellm.litellm_core_utils.url_utils import SSRFError @pytest.fixture(autouse=True) @@ -321,34 +324,142 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o assert messages == snapshot -async def test_async_inline_remote_media_leaves_skipped_url_prefixes_untouched(async_only_image_fetch): - skipped_prefix = "https://generativelanguage.googleapis.com/v1beta/files/" - skipped_file = f"{skipped_prefix}{uuid.uuid4().hex}" - skipped_image = f"{skipped_prefix}{uuid.uuid4().hex}" - fetched_image = f"https://img.example/{uuid.uuid4()}.png" +async def test_async_inline_remote_media_inlines_only_the_parts_the_predicate_accepts(async_only_image_fetch): + files_api_prefix = "https://generativelanguage.googleapis.com/v1beta/files/" + files_api_pdf = f"{files_api_prefix}{uuid.uuid4().hex}" + hinted_image = f"https://img.example/{uuid.uuid4()}.png" + plain_image = f"https://img.example/{uuid.uuid4()}.png" + hinted_document = f"https://docs.example/{uuid.uuid4()}.pdf" + seen = [] + + def inline_unhinted_outside_files_api(media: RemoteMedia) -> bool: + seen.append(media) + return not media.url.startswith(files_api_prefix) and "format" not in media.fields + messages = [ { "role": "user", "content": [ - {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}}, - {"type": "image_url", "image_url": {"url": skipped_image}}, - {"type": "image_url", "image_url": fetched_image}, + {"type": "file", "file": {"file_id": files_api_pdf}}, + {"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": plain_image}}, + {"type": "image_url", "image_url": plain_image}, + {"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}}, ], } ] snapshot = copy.deepcopy(messages) - inlined = await async_inline_remote_media(messages, skip_url_prefixes=(skipped_prefix,)) + inlined = await async_inline_remote_media(messages, should_inline=inline_unhinted_outside_files_api) assert inlined[0]["content"] == [ - {"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}}, - {"type": "image_url", "image_url": {"url": skipped_image}}, + {"type": "file", "file": {"file_id": files_api_pdf}}, + {"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}}, {"type": "image_url", "image_url": async_only_image_fetch.data_url}, + {"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}}, + ] + assert async_only_image_fetch.fetched == [plain_image] + assert [(media.url, dict(media.fields)) for media in seen[:5]] == [ + (files_api_pdf, {"file_id": files_api_pdf}), + (hinted_image, {"url": hinted_image, "format": "image/png"}), + (plain_image, {"url": plain_image}), + (plain_image, {}), + (hinted_document, {"type": "url", "url": hinted_document, "format": "application/pdf"}), ] - assert async_only_image_fetch.fetched == [fetched_image] assert messages == snapshot +async def test_async_inline_remote_media_inlines_a_shared_url_only_where_the_predicate_accepts_it( + async_only_image_fetch, +): + shared = f"https://img.example/{uuid.uuid4()}.png" + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": shared, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": shared}}, + ], + } + ] + + inlined = await async_inline_remote_media(messages, should_inline=lambda media: "format" not in media.fields) + + assert inlined[0]["content"] == [ + {"type": "image_url", "image_url": {"url": shared, "format": "image/png"}}, + {"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}}, + ] + assert async_only_image_fetch.fetched == [shared] + + +async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fails(monkeypatch): + missing = f"http://img.example/{uuid.uuid4()}-missing.png" + slow = f"http://img.example/{uuid.uuid4()}-slow.png" + slow_fetch_outcomes = [] + + async def serve(client, url, **kwargs): + if url == missing: + return Response(404, request=Request("GET", url)) + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + slow_fetch_outcomes.append("cancelled") + raise + slow_fetch_outcomes.append("finished") + return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url)) + + monkeypatch.setattr(image_handling, "async_safe_get", serve) + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": missing}}, + {"type": "image_url", "image_url": {"url": slow}}, + ], + } + ] + started = time.perf_counter() + + with pytest.raises(litellm.ImageFetchError, match="Status code: 404"): + await async_inline_remote_media(messages) + + assert slow_fetch_outcomes == ["cancelled"] + assert time.perf_counter() - started < 1 + + +async def test_async_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch): + attempts = [] + + async def block(client, url, **kwargs): + attempts.append(url) + raise SSRFError("URL targets a blocked address (10.0.0.8)") + + monkeypatch.setattr(image_handling, "async_safe_get", block) + url = f"http://internal.example/{uuid.uuid4()}.png" + + with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"): + await async_convert_url_to_base64(url) + + assert attempts == [url] + + +def test_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch): + attempts = [] + + def block(client, url, **kwargs): + attempts.append(url) + raise SSRFError("URL targets a blocked address (10.0.0.8)") + + monkeypatch.setattr(image_handling, "safe_get", block) + url = f"http://internal.example/{uuid.uuid4()}.png" + + with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"): + convert_url_to_base64(url) + + assert attempts == [url] + + async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch): in_flight = {"now": 0, "peak": 0} diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index aaaa43a0dc4..fccdc1a2a0a 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -1,5 +1,9 @@ +import asyncio import socket +import threading +import time +import httpx import pytest import litellm @@ -535,3 +539,34 @@ def test_assert_same_origin_error_message_does_not_leak_hostnames(): detail = str(exc.value) assert "attacker.example.com" not in detail assert "api.internal-corp.example" not in detail + + +async def test_async_safe_get_resolves_dns_off_the_event_loop(monkeypatch): + loop_thread = threading.current_thread() + resolver_threads = [] + + def slow_getaddrinfo(host, port, *args, **kwargs): + resolver_threads.append(threading.current_thread()) + time.sleep(0.4) + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port or 443))] + + monkeypatch.setattr(url_utils.socket, "getaddrinfo", slow_getaddrinfo) + + class FakeClient: + async def get(self, url, **kwargs): + return httpx.Response(200, request=httpx.Request("GET", url)) + + ticks = [time.perf_counter()] + + async def heartbeat(): + while True: + await asyncio.sleep(0.01) + ticks.append(time.perf_counter()) + + beating = asyncio.create_task(heartbeat()) + response = await url_utils.async_safe_get(FakeClient(), "https://img.example/a.png") + beating.cancel() + + assert response.status_code == 200 + assert resolver_threads and all(thread is not loop_thread for thread in resolver_threads) + assert max(b - a for a, b in zip(ticks, ticks[1:])) < 0.2 diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index e86e671b939..f1614654ffb 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import threading import time from unittest.mock import AsyncMock, Mock, patch @@ -3192,6 +3193,7 @@ class _TransformRecordingConfig(BaseConfig): def __init__(self, transform_async: bool): self.transform_async = transform_async self.transform_calls = [] + self.sign_threads = [] @property def uses_async_transform_request(self) -> bool: @@ -3216,6 +3218,12 @@ class _TransformRecordingConfig(BaseConfig): self.transform_calls.append("async") return {"transformed_by": "async"} + def sign_request( + self, headers, optional_params, request_data, api_base, api_key=None, model=None, stream=None, fake_stream=None + ): + self.sign_threads.append(threading.current_thread()) + return headers, None + def transform_response( self, model, @@ -3237,7 +3245,7 @@ class _TransformRecordingConfig(BaseConfig): return BaseLLMException(status_code=status_code, message=error_message, headers=headers) -def _start_async_completion(config): +def _start_async_completion(config, logging_obj=None): captured = {} def handle(request): @@ -3253,7 +3261,7 @@ def _start_async_completion(config): custom_llm_provider="openai", model_response=ModelResponse(), encoding=None, - logging_obj=Mock(dynamic_success_callbacks=None, model_call_details={}), + logging_obj=logging_obj if logging_obj is not None else Mock(dynamic_success_callbacks=None, model_call_details={}), optional_params={}, timeout=10.0, litellm_params={}, @@ -3277,6 +3285,22 @@ async def test_completion_awaits_async_transform_request_when_config_opts_in(): assert response.choices[0].message.content == "async" +async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_transform(): + config = _TransformRecordingConfig(transform_async=True) + loop_thread = threading.current_thread() + pre_call_threads = [] + logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={}) + logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread()) + + pending, captured = _start_async_completion(config, logging_obj) + response = await pending + + assert response.choices[0].message.content == "async" + assert captured["body"] == {"transformed_by": "async"} + assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) + assert pre_call_threads and all(thread is not loop_thread for thread in pre_call_threads) + + async def test_completion_keeps_sync_transform_request_before_returning_by_default(): config = _TransformRecordingConfig(transform_async=False) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py index d31254746d4..f135acd094f 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_transformation.py @@ -1,6 +1,8 @@ import json import uuid +from unittest.mock import Mock + import httpx import pytest @@ -424,3 +426,52 @@ async def test_gemini_ai_studio_async_completion_passes_files_api_uris_through_u {"mime_type": "application/pdf", "file_uri": files_api_pdf}, {"mime_type": "image/png", "file_uri": files_api_image}, ] + + +async def test_vertex_ai_async_transform_inlines_only_the_urls_gemini_cannot_fetch_itself(async_only_image_fetch): + plain_http_png = f"http://img.example/{uuid.uuid4()}.png" + extensionless_https = f"https://cdn.example/files/{uuid.uuid4().hex}" + https_png = f"https://img.example/{uuid.uuid4()}.png" + hinted_extensionless = f"https://cdn.example/files/{uuid.uuid4().hex}" + files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe these"}, + {"type": "image_url", "image_url": {"url": plain_http_png}}, + {"type": "image_url", "image_url": {"url": extensionless_https}}, + {"type": "image_url", "image_url": {"url": https_png}}, + {"type": "image_url", "image_url": {"url": hinted_extensionless, "mime_type": "image/webp"}}, + {"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}}, + ], + } + ] + + body = await transformation.async_transform_request_body( + gemini_api_key=None, + messages=messages, + api_base=None, + model="gemini-3.8-flash", + client=None, + timeout=None, + extra_headers=None, + optional_params={}, + logging_obj=Mock(), + custom_llm_provider="vertex_ai", + litellm_params={}, + vertex_project="qa-project", + vertex_location="us-central1", + vertex_auth_header=None, + ) + + inlined = {"inline_data": {"mime_type": "image/png", "data": async_only_image_fetch.base64_png}} + assert body["contents"][0]["parts"] == [ + {"text": "Describe these"}, + inlined, + inlined, + {"file_data": {"mime_type": "image/png", "file_uri": https_png}}, + {"file_data": {"mime_type": "image/webp", "file_uri": hinted_extensionless}}, + {"file_data": {"mime_type": "application/pdf", "file_uri": files_api_pdf}}, + ] + assert sorted(async_only_image_fetch.fetched) == sorted([plain_http_png, extensionless_https]) From 0ee3bec046bf0f15eebdddc1d52610d389ca220f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 23:20:21 -0700 Subject: [PATCH 6/8] fix(image_edit): await an async transform hook and hide the URL policy verdict from callers The image edit handler now awaits BaseImageEditConfig.async_transform_image_edit_request, and Black Forest Labs overrides it so URL images and masks download through async_safe_get instead of the blocking safe_get on the event loop. Rejected image fetches raise a fixed policy message with the user_url_allowed_hosts hint rather than echoing the resolver's verdict (resolved IP, DNS failure) back to the caller. The test fixture also fails any request-path call of the sync convert_url_to_base64 so a regression cannot pass unnoticed. --- .../prompt_templates/image_handling.py | 12 +- .../base_llm/image_edit/transformation.py | 19 ++ .../image_edit/transformation.py | 102 ++++++----- litellm/llms/custom_httpx/llm_http_handler.py | 2 +- tests/test_litellm/conftest.py | 10 +- .../litellm_core_utils/test_image_handling.py | 65 ++++--- .../test_bfl_image_edit_transformation.py | 78 ++++++++- .../custom_httpx/test_llm_http_handler.py | 163 +++++++++++++++++- 8 files changed, 379 insertions(+), 72 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index dcad7776b58..4e7dfdb4017 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -78,6 +78,14 @@ def _process_image_response(response: Response, url: str) -> str: return result +def _url_policy_rejection(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": + verbose_logger.warning("Image fetch of %s rejected by the URL policy: %s", url, verdict) + return litellm.ImageFetchError( + "Error: Unable to fetch image from URL. The proxy's URL policy rejected this host; " + f"an admin can allow it with `user_url_allowed_hosts` in general_settings. url={url}" + ) + + async def async_convert_url_to_base64(url: str) -> str: if url.startswith("data:") and ";base64," in url: return url @@ -100,7 +108,7 @@ async def async_convert_url_to_base64(url: str) -> str: except litellm.ImageFetchError: raise except SSRFError as e: - raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e + raise _url_policy_rejection(url, e) from e except Exception: pass raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}") @@ -128,7 +136,7 @@ def convert_url_to_base64(url: str) -> str: except litellm.ImageFetchError: raise except SSRFError as e: - raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e + raise _url_policy_rejection(url, e) from e except Exception as e: verbose_logger.exception(e) raise litellm.ImageFetchError( diff --git a/litellm/llms/base_llm/image_edit/transformation.py b/litellm/llms/base_llm/image_edit/transformation.py index 9a25d3294e0..4faf0aaaf30 100644 --- a/litellm/llms/base_llm/image_edit/transformation.py +++ b/litellm/llms/base_llm/image_edit/transformation.py @@ -1,5 +1,6 @@ import types from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any import httpx @@ -102,6 +103,24 @@ class BaseImageEditConfig(ABC): ) -> tuple[dict, RequestFiles]: pass + async def async_transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: FileTypes | None, + image_edit_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[dict, RequestFiles]: + return self.transform_image_edit_request( + model=model, + prompt=prompt, + image=image, + image_edit_optional_request_params=dict(image_edit_optional_request_params), + litellm_params=litellm_params, + headers=dict(headers), + ) + def finalize_image_edit_request_data(self, data: dict, resolved_request_url: str) -> dict: """ Last pass on the request dict after ``transform_image_edit_request``, using the diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 013053e5bd5..c6de9f9ba65 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -9,6 +9,7 @@ API Reference: https://docs.bfl.ai/ import base64 import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -16,7 +17,7 @@ from httpx._types import RequestFiles import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH -from litellm.litellm_core_utils.url_utils import safe_get +from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -37,6 +38,22 @@ else: LiteLLMLoggingObj = Any +_BFL_REQUEST_PARAMS: Final = ( + "seed", + "output_format", + "safety_tolerance", + "prompt_upsampling", + "aspect_ratio", + "steps", + "guidance", + "grow_mask", + "top", + "bottom", + "left", + "right", +) + + class BlackForestLabsImageEditConfig(BaseImageEditConfig): """ Configuration for Black Forest Labs image editing. @@ -85,34 +102,10 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): BFL-specific params are passed through directly. """ optional_params: Final[dict[str, Any]] = {} - - # Pass through BFL-specific params - bfl_params: Final = [ - "seed", - "output_format", - "safety_tolerance", - "prompt_upsampling", - # Kontext-specific - "aspect_ratio", - # Fill/Inpaint-specific - "steps", - "guidance", - "grow_mask", - # Expand-specific - "top", - "bottom", - "left", - "right", - ] - - # Convert TypedDict to regular dict for access - params_dict: Final = dict(image_edit_optional_params) - - for param in bfl_params: - if param in params_dict: - value = params_dict[param] - if value is not None: - optional_params[param] = value + params: Final[Mapping[str, object]] = image_edit_optional_params + for param in _BFL_REQUEST_PARAMS: + if (value := params.get(param)) is not None: + optional_params[param] = value # Set default output format if "output_format" not in optional_params: @@ -251,23 +244,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): "input_image": b64_image, } - # Add optional params (only BFL-recognized parameters) - bfl_request_params: Final = [ - "seed", - "output_format", - "safety_tolerance", - "prompt_upsampling", - "aspect_ratio", - "steps", - "guidance", - "grow_mask", - "top", - "bottom", - "left", - "right", - ] for key, value in image_edit_optional_request_params.items(): - if key in bfl_request_params and value is not None: + if key in _BFL_REQUEST_PARAMS and value is not None: request_body[key] = value # Handle mask if provided (for inpainting) @@ -277,7 +255,39 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8") # BFL uses JSON, not multipart - return empty files - return request_body, [] + return request_body, () + + async def async_transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: FileTypes | None, + image_edit_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[dict, RequestFiles]: + downloaded_image: Final = await self._fetch_remote_image(image) + downloaded_mask: Final = await self._fetch_remote_image(image_edit_optional_request_params.get("mask")) + return self.transform_image_edit_request( + model=model, + prompt=prompt, + image=image if downloaded_image is None else downloaded_image, + image_edit_optional_request_params=( + dict(image_edit_optional_request_params) + if downloaded_mask is None + else {**image_edit_optional_request_params, "mask": downloaded_mask} + ), + litellm_params=litellm_params, + headers=dict(headers), + ) + + async def _fetch_remote_image(self, image: object) -> bytes | None: + candidate: Final = image[0] if isinstance(image, list) and image else image + if not isinstance(candidate, str) or not candidate.startswith(("http://", "https://")): + return None + response: Final = await async_safe_get(litellm.module_level_aclient, candidate, timeout=60.0) + response.raise_for_status() + return response.content def transform_image_edit_response( self, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d57aee025d5..6a17841faa2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6787,7 +6787,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - data, files = image_edit_provider_config.transform_image_edit_request( + data, files = await image_edit_provider_config.async_transform_image_edit_request( model=model, image=image, prompt=prompt, diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index aa8ba168ecc..a4f32df46ae 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -607,7 +607,8 @@ ONE_PIXEL_PNG = base64.b64decode( @pytest.fixture def async_only_image_fetch(monkeypatch): - from litellm.litellm_core_utils.prompt_templates import image_handling + from litellm.litellm_core_utils.prompt_templates import factory, image_handling + from litellm.llms.gemini.chat import transformation as gemini_chat_transformation fetch = SimpleNamespace( fetched=[], @@ -627,6 +628,13 @@ def async_only_image_fetch(monkeypatch): request=httpx.Request("GET", url), ) + def forbid_sync_convert(url, *args, **kwargs): + if url.startswith(("http://", "https://")): + raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}") + return url + monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch) monkeypatch.setattr(image_handling, "async_safe_get", serve_png) + for module in (image_handling, factory, gemini_chat_transformation): + monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert) return fetch 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 e7b28d0d3a2..d8ccc9251dd 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -428,36 +428,61 @@ async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fail assert time.perf_counter() - started < 1 -async def test_async_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch): +_SSRF_VERDICTS = ( + "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " + "add the host to `user_url_allowed_hosts` in general_settings.", + "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known", + "No addresses found for 'internal.example'", +) + + +def _assert_one_verdict_free_message(messages, url): + assert len(set(messages)) == 1 + assert "10.0.0.8" not in messages[0] + assert "DNS" not in messages[0] + assert "No addresses" not in messages[0] + assert "user_url_allowed_hosts" in messages[0] + assert url in messages[0] + + +async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): attempts = [] - - async def block(client, url, **kwargs): - attempts.append(url) - raise SSRFError("URL targets a blocked address (10.0.0.8)") - - monkeypatch.setattr(image_handling, "async_safe_get", block) + messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"): - await async_convert_url_to_base64(url) + for verdict in _SSRF_VERDICTS: - assert attempts == [url] + async def block(client, fetched_url, verdict=verdict, **kwargs): + attempts.append(fetched_url) + raise SSRFError(verdict) + + monkeypatch.setattr(image_handling, "async_safe_get", block) + with pytest.raises(litellm.ImageFetchError) as raised: + await async_convert_url_to_base64(url) + messages.append(raised.value.message) + + assert attempts == [url] * len(_SSRF_VERDICTS) + _assert_one_verdict_free_message(messages, url) -def test_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch): +def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): attempts = [] - - def block(client, url, **kwargs): - attempts.append(url) - raise SSRFError("URL targets a blocked address (10.0.0.8)") - - monkeypatch.setattr(image_handling, "safe_get", block) + messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"): - convert_url_to_base64(url) + for verdict in _SSRF_VERDICTS: - assert attempts == [url] + def block(client, fetched_url, verdict=verdict, **kwargs): + attempts.append(fetched_url) + raise SSRFError(verdict) + + monkeypatch.setattr(image_handling, "safe_get", block) + with pytest.raises(litellm.ImageFetchError) as raised: + convert_url_to_base64(url) + messages.append(raised.value.message) + + assert attempts == [url] * len(_SSRF_VERDICTS) + _assert_one_verdict_free_message(messages, url) async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch): diff --git a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py index ec243b7058d..d9d6e813d86 100644 --- a/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py +++ b/tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py @@ -16,6 +16,9 @@ import httpx import pytest +from litellm.llms.black_forest_labs.image_edit import ( + transformation as bfl_transformation, +) from litellm.llms.black_forest_labs.image_edit.transformation import ( BlackForestLabsImageEditConfig, ) @@ -186,7 +189,7 @@ class TestBlackForestLabsImageEditTransformation: assert data["output_format"] == "jpeg" # BFL uses JSON, not multipart - files should be empty - assert files == [] + assert files == () def test_transform_image_edit_request_with_mask(self): """Test request transformation with mask for inpainting.""" @@ -299,3 +302,76 @@ class TestBlackForestLabsImageEditTransformation: def test_use_multipart_form_data_returns_false(self): """Test that use_multipart_form_data returns False for BFL.""" assert self.config.use_multipart_form_data() is False + + +async def test_async_transform_image_edit_request_downloads_url_images_with_the_async_fetcher(monkeypatch): + served = b"png-bytes-from-cdn" + fetched = [] + + def forbid_sync_fetch(client, url, **kwargs): + raise AssertionError(f"sync image fetch ran on the event loop: {url}") + + async def serve(client, url, **kwargs): + fetched.append((url, kwargs.get("timeout"))) + return httpx.Response(200, content=served, request=httpx.Request("GET", url)) + + monkeypatch.setattr(bfl_transformation, "safe_get", forbid_sync_fetch) + monkeypatch.setattr(bfl_transformation, "async_safe_get", serve) + + data, files = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image="https://cdn.example/photo.png", + image_edit_optional_request_params={"mask": "https://cdn.example/mask.png", "seed": 7}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert base64.b64decode(data["input_image"]) == served + assert base64.b64decode(data["mask"]) == served + assert data["seed"] == 7 + assert files == () + assert fetched == [("https://cdn.example/photo.png", 60.0), ("https://cdn.example/mask.png", 60.0)] + + +async def test_async_transform_image_edit_request_never_fetches_for_local_images(monkeypatch): + def refuse(*args, **kwargs): + raise AssertionError("no network fetch expected for local image bytes") + + monkeypatch.setattr(bfl_transformation, "safe_get", refuse) + monkeypatch.setattr(bfl_transformation, "async_safe_get", refuse) + + data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image=[BytesIO(b"first"), BytesIO(b"other")], + image_edit_optional_request_params={"mask": b"mask-bytes"}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert base64.b64decode(data["input_image"]) == b"first" + assert base64.b64decode(data["mask"]) == b"mask-bytes" + + +async def test_async_transform_image_edit_request_downloads_only_the_first_url_of_a_list(monkeypatch): + fetched = [] + + async def serve(client, url, **kwargs): + fetched.append(url) + return httpx.Response(200, content=b"first-bytes", request=httpx.Request("GET", url)) + + monkeypatch.setattr(bfl_transformation, "safe_get", lambda *args, **kwargs: pytest.fail("sync fetch ran")) + monkeypatch.setattr(bfl_transformation, "async_safe_get", serve) + + data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request( + model="flux-kontext-pro", + prompt="Add a red hat", + image=["https://cdn.example/a.png", "https://cdn.example/b.png"], + image_edit_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert fetched == ["https://cdn.example/a.png"] + assert base64.b64decode(data["input_image"]) == b"first-bytes" diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index f1614654ffb..3afd57bbf34 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -19,6 +19,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( BaseAudioTranscriptionConfig, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, @@ -31,7 +32,7 @@ from litellm.llms.azure.videos.transformation import AzureVideoConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import ModelResponse, TranscriptionResponse +from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" @@ -3244,6 +3245,11 @@ class _TransformRecordingConfig(BaseConfig): def get_error_class(self, error_message, status_code, headers): return BaseLLMException(status_code=status_code, message=error_message, headers=headers) + def get_model_response_iterator(self, streaming_response, sync_stream, json_mode=False): + return litellm.OpenAIGPTConfig().get_model_response_iterator( + streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode + ) + def _start_async_completion(config, logging_obj=None): captured = {} @@ -3312,3 +3318,158 @@ async def test_completion_keeps_sync_transform_request_before_returning_by_defau assert config.transform_calls == ["sync"] assert captured["body"] == {"transformed_by": "sync"} assert response.choices[0].message.content == "sync" + + +def _sse_echoing_transformed_by(request): + transformed_by = json.loads(request.content)["transformed_by"] + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "stub-model", + "choices": [{"index": 0, "delta": {"content": transformed_by}, "finish_reason": None}], + } + return httpx.Response( + 200, + content=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode(), + headers={"content-type": "text/event-stream"}, + request=request, + ) + + +def _streaming_logging_obj(): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="async-transform-stream", + function_id="f", + ) + logging_obj.update_environment_variables( + model="stub-model", user="", optional_params={}, litellm_params={}, custom_llm_provider="openai" + ) + return logging_obj + + +async def test_completion_streams_after_the_async_transform_request(): + config = _TransformRecordingConfig(transform_async=True) + loop_thread = threading.current_thread() + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(_sse_echoing_transformed_by)) + + stream = await BaseLLMHTTPHandler().completion( + model="stub-model", + messages=[{"role": "user", "content": "hi"}], + api_base="https://llm.example/v1/chat", + custom_llm_provider="openai", + model_response=ModelResponse(), + encoding=None, + logging_obj=_streaming_logging_obj(), + optional_params={}, + timeout=10.0, + litellm_params={}, + acompletion=True, + stream=True, + client=client, + provider_config=config, + ) + collected = [chunk async for chunk in stream] + + assert config.transform_calls == ["async"] + assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) + assert "".join(chunk.choices[0].delta.content or "" for chunk in collected) == "async" + + +class _ImageEditRecordingConfig(BaseImageEditConfig): + def __init__(self): + self.transform_calls = [] + + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, image_edit_optional_params, model, drop_params): + return dict(image_edit_optional_params) + + def validate_environment(self, headers, model, api_key=None, litellm_params=None, api_base=None): + return {} + + def get_complete_url(self, model, api_base, litellm_params): + return "https://images.example/v1/edits" + + def use_multipart_form_data(self): + return False + + def transform_image_edit_request( + self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers + ): + self.transform_calls.append("sync") + return {"transformed_by": "sync"}, [] + + async def async_transform_image_edit_request( + self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers + ): + self.transform_calls.append("async") + return {"transformed_by": "async"}, [] + + def transform_image_edit_response(self, model, raw_response, logging_obj): + return ImageResponse(data=[ImageObject(b64_json=raw_response.json()["transformed_by"])]) + + +def _echo_json_transport(captured): + def handle(request): + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=captured["body"]) + + return httpx.MockTransport(handle) + + +async def test_async_image_edit_handler_awaits_the_async_transform(): + config = _ImageEditRecordingConfig() + captured = {} + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_echo_json_transport(captured)) + + response = await BaseLLMHTTPHandler().async_image_edit_handler( + model="edit-model", + image=b"raw-image", + prompt="add a hat", + image_edit_provider_config=config, + image_edit_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=Mock(), + timeout=10.0, + client=client, + ) + + assert config.transform_calls == ["async"] + assert captured["body"] == {"transformed_by": "async"} + assert response.data[0].b64_json == "async" + + +def test_image_edit_handler_keeps_the_sync_transform(): + config = _ImageEditRecordingConfig() + captured = {} + client = HTTPHandler() + client.client = httpx.Client(transport=_echo_json_transport(captured)) + + response = BaseLLMHTTPHandler().image_edit_handler( + model="edit-model", + image=b"raw-image", + prompt="add a hat", + image_edit_provider_config=config, + image_edit_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=Mock(), + timeout=10.0, + client=client, + ) + + assert config.transform_calls == ["sync"] + assert captured["body"] == {"transformed_by": "sync"} + assert response.data[0].b64_json == "sync" From 3920cf4dfe458b6d4982e2a64c162de21defa2a0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 23:45:37 -0700 Subject: [PATCH 7/8] fix(image_handling): tell callers when the image host did not resolve instead of blaming the URL policy validate_url raises HostResolutionError, a SSRFError subclass, for the two DNS outcomes (lookup failed, no addresses). The image fetch helper maps that to a "host could not be resolved" message and keeps the user_url_allowed_hosts hint for the policy verdicts it can actually fix. --- .../prompt_templates/image_handling.py | 14 +++-- litellm/litellm_core_utils/url_utils.py | 8 ++- .../litellm_core_utils/test_image_handling.py | 56 +++++++++++-------- .../litellm_core_utils/test_url_utils.py | 19 ++++++- 4 files changed, 64 insertions(+), 33 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 4e7dfdb4017..99efbd13320 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -15,7 +15,7 @@ import litellm from litellm import verbose_logger from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB -from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError, async_safe_get, safe_get from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 @@ -78,8 +78,12 @@ def _process_image_response(response: Response, url: str) -> str: return result -def _url_policy_rejection(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": - verbose_logger.warning("Image fetch of %s rejected by the URL policy: %s", url, verdict) +def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": + verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict) + if isinstance(verdict, HostResolutionError): + return litellm.ImageFetchError( + f"Error: Unable to fetch image from URL. The image host could not be resolved. url={url}" + ) return litellm.ImageFetchError( "Error: Unable to fetch image from URL. The proxy's URL policy rejected this host; " f"an admin can allow it with `user_url_allowed_hosts` in general_settings. url={url}" @@ -108,7 +112,7 @@ async def async_convert_url_to_base64(url: str) -> str: except litellm.ImageFetchError: raise except SSRFError as e: - raise _url_policy_rejection(url, e) from e + raise _rejected_image_fetch(url, e) from e except Exception: pass raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}") @@ -136,7 +140,7 @@ def convert_url_to_base64(url: str) -> str: except litellm.ImageFetchError: raise except SSRFError as e: - raise _url_policy_rejection(url, e) from e + raise _rejected_image_fetch(url, e) from e except Exception as e: verbose_logger.exception(e) raise litellm.ImageFetchError( diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index fa070a648f5..f4388d5fd12 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -93,6 +93,10 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" +class HostResolutionError(SSRFError): + pass + + def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. @@ -324,10 +328,10 @@ def validate_url(url: str) -> tuple[str, str]: try: addrinfo: Final = socket.getaddrinfo(hostname, effective_port, proto=socket.IPPROTO_TCP) except socket.gaierror as e: - raise SSRFError(f"DNS resolution failed for '{hostname}': {e}") + raise HostResolutionError(f"DNS resolution failed for '{hostname}': {e}") if not addrinfo: - raise SSRFError(f"No addresses found for '{hostname}'") + raise HostResolutionError(f"No addresses found for '{hostname}'") if not is_allowlisted: for addrinfo_entry in addrinfo: 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 d8ccc9251dd..3a909298aba 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -17,7 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_inline_remote_media, convert_url_to_base64, ) -from litellm.litellm_core_utils.url_utils import SSRFError +from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError @pytest.fixture(autouse=True) @@ -115,9 +115,7 @@ class StreamingLargeImageClient: request=Request("GET", url), ) # Mock the iter_bytes method to return our generator - response.iter_bytes = lambda chunk_size=8192: generate_chunks( - size_bytes, chunk_size - ) + response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size) return response @@ -215,9 +213,7 @@ def test_streaming_download_handles_petabyte_file(monkeypatch): """ # Simulate a 1 petabyte file (1,000,000 GB) # Without streaming protection, this would cause OOM or hang indefinitely - client = StreamingLargeImageClient( - size_mb=1_000_000_000, include_content_length=False - ) + client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False) monkeypatch.setattr(litellm, "module_level_client", client) with pytest.raises(litellm.ImageFetchError) as excinfo: @@ -429,20 +425,32 @@ async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fail _SSRF_VERDICTS = ( - "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " - "add the host to `user_url_allowed_hosts` in general_settings.", - "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known", - "No addresses found for 'internal.example'", + ( + SSRFError( + "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " + "add the host to `user_url_allowed_hosts` in general_settings." + ), + "The proxy's URL policy rejected this host; an admin can allow it with `user_url_allowed_hosts`", + ), + ( + HostResolutionError( + "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known" + ), + "The image host could not be resolved", + ), + (HostResolutionError("No addresses found for 'internal.example'"), "The image host could not be resolved"), ) -def _assert_one_verdict_free_message(messages, url): - assert len(set(messages)) == 1 - assert "10.0.0.8" not in messages[0] - assert "DNS" not in messages[0] - assert "No addresses" not in messages[0] - assert "user_url_allowed_hosts" in messages[0] - assert url in messages[0] +def _assert_verdict_free_messages(messages, url): + for message, (verdict, expected_guidance) in zip(messages, _SSRF_VERDICTS, strict=True): + assert expected_guidance in message + assert url in message + assert "10.0.0.8" not in message + assert "DNS" not in message + assert "No addresses" not in message + if isinstance(verdict, HostResolutionError): + assert "user_url_allowed_hosts" not in message async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): @@ -450,11 +458,11 @@ async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_r messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - for verdict in _SSRF_VERDICTS: + for verdict, _ in _SSRF_VERDICTS: async def block(client, fetched_url, verdict=verdict, **kwargs): attempts.append(fetched_url) - raise SSRFError(verdict) + raise verdict monkeypatch.setattr(image_handling, "async_safe_get", block) with pytest.raises(litellm.ImageFetchError) as raised: @@ -462,7 +470,7 @@ async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_r messages.append(raised.value.message) assert attempts == [url] * len(_SSRF_VERDICTS) - _assert_one_verdict_free_message(messages, url) + _assert_verdict_free_messages(messages, url) def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): @@ -470,11 +478,11 @@ def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeyp messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - for verdict in _SSRF_VERDICTS: + for verdict, _ in _SSRF_VERDICTS: def block(client, fetched_url, verdict=verdict, **kwargs): attempts.append(fetched_url) - raise SSRFError(verdict) + raise verdict monkeypatch.setattr(image_handling, "safe_get", block) with pytest.raises(litellm.ImageFetchError) as raised: @@ -482,7 +490,7 @@ def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeyp messages.append(raised.value.message) assert attempts == [url] * len(_SSRF_VERDICTS) - _assert_one_verdict_free_message(messages, url) + _assert_verdict_free_messages(messages, url) async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch): diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index fccdc1a2a0a..151c7ceb5bd 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm.litellm_core_utils import url_utils from litellm.litellm_core_utils.url_utils import ( + HostResolutionError, SSRFError, _is_blocked_ip, assert_same_origin, @@ -161,10 +162,24 @@ class TestValidateUrl: assert "/path" in rewritten assert "key=value" in rewritten - def test_dns_failure_raises(self, mock_dns_failure): - with pytest.raises(SSRFError, match="DNS resolution failed"): + def test_dns_failure_raises_a_host_resolution_error(self, mock_dns_failure): + with pytest.raises(HostResolutionError, match="DNS resolution failed"): validate_url("http://this-domain-does-not-exist-xyz123.invalid/test") + def test_empty_resolution_raises_a_host_resolution_error(self, monkeypatch): + monkeypatch.setattr(url_utils.socket, "getaddrinfo", lambda *args, **kwargs: []) + with pytest.raises(HostResolutionError, match="No addresses found"): + validate_url("http://this-domain-resolves-to-nothing.invalid/test") + + def test_blocked_address_is_not_a_host_resolution_error(self, monkeypatch): + def fake(host, port, *a, **kw): + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.8", port or 80))] + + monkeypatch.setattr(url_utils.socket, "getaddrinfo", fake) + with pytest.raises(SSRFError, match="blocked address") as raised: + validate_url("http://internal.example/test") + assert not isinstance(raised.value, HostResolutionError) + def test_blocks_localhost_hostname(self, monkeypatch): def fake(host, port, *a, **kw): return [ From 827554954d305b7c8c24f524c961cf7a27e3a029 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:16:14 -0700 Subject: [PATCH 8/8] fix(image_handling): answer every SSRF rejection with one message so error text cannot probe internal hostnames --- .../prompt_templates/image_handling.py | 10 ++--- litellm/litellm_core_utils/url_utils.py | 8 +--- .../litellm_core_utils/test_image_handling.py | 41 ++++++++----------- .../litellm_core_utils/test_url_utils.py | 19 +-------- 4 files changed, 24 insertions(+), 54 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 99efbd13320..24f3b8bca7f 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -15,7 +15,7 @@ import litellm from litellm import verbose_logger from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB -from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError, async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 @@ -80,13 +80,9 @@ def _process_image_response(response: Response, url: str) -> str: def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict) - if isinstance(verdict, HostResolutionError): - return litellm.ImageFetchError( - f"Error: Unable to fetch image from URL. The image host could not be resolved. url={url}" - ) return litellm.ImageFetchError( - "Error: Unable to fetch image from URL. The proxy's URL policy rejected this host; " - f"an admin can allow it with `user_url_allowed_hosts` in general_settings. url={url}" + "Error: Unable to fetch image from URL. The proxy could not resolve this host or its URL policy rejected it; " + f"an admin can check the proxy log and `user_url_allowed_hosts` in general_settings. url={url}" ) diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index f4388d5fd12..fa070a648f5 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -93,10 +93,6 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" -class HostResolutionError(SSRFError): - pass - - def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. @@ -328,10 +324,10 @@ def validate_url(url: str) -> tuple[str, str]: try: addrinfo: Final = socket.getaddrinfo(hostname, effective_port, proto=socket.IPPROTO_TCP) except socket.gaierror as e: - raise HostResolutionError(f"DNS resolution failed for '{hostname}': {e}") + raise SSRFError(f"DNS resolution failed for '{hostname}': {e}") if not addrinfo: - raise HostResolutionError(f"No addresses found for '{hostname}'") + raise SSRFError(f"No addresses found for '{hostname}'") if not is_allowlisted: for addrinfo_entry in addrinfo: 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 3a909298aba..8fa4bd6c14d 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -17,7 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_inline_remote_media, convert_url_to_base64, ) -from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError +from litellm.litellm_core_utils.url_utils import SSRFError @pytest.fixture(autouse=True) @@ -425,32 +425,25 @@ async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fail _SSRF_VERDICTS = ( - ( - SSRFError( - "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " - "add the host to `user_url_allowed_hosts` in general_settings." - ), - "The proxy's URL policy rejected this host; an admin can allow it with `user_url_allowed_hosts`", + SSRFError( + "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " + "add the host to `user_url_allowed_hosts` in general_settings." ), - ( - HostResolutionError( - "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known" - ), - "The image host could not be resolved", - ), - (HostResolutionError("No addresses found for 'internal.example'"), "The image host could not be resolved"), + SSRFError("DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known"), + SSRFError("No addresses found for 'internal.example'"), ) def _assert_verdict_free_messages(messages, url): - for message, (verdict, expected_guidance) in zip(messages, _SSRF_VERDICTS, strict=True): - assert expected_guidance in message - assert url in message - assert "10.0.0.8" not in message - assert "DNS" not in message - assert "No addresses" not in message - if isinstance(verdict, HostResolutionError): - assert "user_url_allowed_hosts" not in message + assert len(messages) == len(_SSRF_VERDICTS) + assert len(set(messages)) == 1, "a caller must not be able to tell a blocked host from one that does not resolve" + message = messages[0] + assert "The proxy could not resolve this host or its URL policy rejected it" in message + assert "user_url_allowed_hosts" in message + assert url in message + assert "10.0.0.8" not in message + assert "DNS" not in message + assert "No addresses" not in message async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): @@ -458,7 +451,7 @@ async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_r messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - for verdict, _ in _SSRF_VERDICTS: + for verdict in _SSRF_VERDICTS: async def block(client, fetched_url, verdict=verdict, **kwargs): attempts.append(fetched_url) @@ -478,7 +471,7 @@ def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeyp messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - for verdict, _ in _SSRF_VERDICTS: + for verdict in _SSRF_VERDICTS: def block(client, fetched_url, verdict=verdict, **kwargs): attempts.append(fetched_url) diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index 151c7ceb5bd..fccdc1a2a0a 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -9,7 +9,6 @@ import pytest import litellm from litellm.litellm_core_utils import url_utils from litellm.litellm_core_utils.url_utils import ( - HostResolutionError, SSRFError, _is_blocked_ip, assert_same_origin, @@ -162,24 +161,10 @@ class TestValidateUrl: assert "/path" in rewritten assert "key=value" in rewritten - def test_dns_failure_raises_a_host_resolution_error(self, mock_dns_failure): - with pytest.raises(HostResolutionError, match="DNS resolution failed"): + def test_dns_failure_raises(self, mock_dns_failure): + with pytest.raises(SSRFError, match="DNS resolution failed"): validate_url("http://this-domain-does-not-exist-xyz123.invalid/test") - def test_empty_resolution_raises_a_host_resolution_error(self, monkeypatch): - monkeypatch.setattr(url_utils.socket, "getaddrinfo", lambda *args, **kwargs: []) - with pytest.raises(HostResolutionError, match="No addresses found"): - validate_url("http://this-domain-resolves-to-nothing.invalid/test") - - def test_blocked_address_is_not_a_host_resolution_error(self, monkeypatch): - def fake(host, port, *a, **kw): - return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.8", port or 80))] - - monkeypatch.setattr(url_utils.socket, "getaddrinfo", fake) - with pytest.raises(SSRFError, match="blocked address") as raised: - validate_url("http://internal.example/test") - assert not isinstance(raised.value, HostResolutionError) - def test_blocks_localhost_hostname(self, monkeypatch): def fake(host, port, *a, **kw): return [