diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 81132fa89a5..24f3b8bca7f 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 Callable, Mapping +from dataclasses import dataclass +from types import MappingProxyType from typing import Final from httpx import Response @@ -11,9 +15,11 @@ 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 +MAX_CONCURRENT_REMOTE_MEDIA_FETCHES: Final = 20 in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY) @@ -72,6 +78,14 @@ def _process_image_response(response: Response, url: str) -> str: return result +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) + return litellm.ImageFetchError( + "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}" + ) + + async def async_convert_url_to_base64(url: str) -> str: if url.startswith("data:") and ";base64," in url: return url @@ -93,6 +107,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 _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}") @@ -119,8 +135,192 @@ def convert_url_to_base64(url: str) -> str: return _process_image_response(response, url) except litellm.ImageFetchError: raise + except SSRFError as e: + raise _rejected_image_fetch(url, e) from e except Exception as e: verbose_logger.exception(e) 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 + + +_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 + 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, source, url) if source is not None and 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 + 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"}) + + +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 _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, ...]: + 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_part(part: object, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]) -> object: + remote: Final = _parse_remote_part(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], 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, 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 + + +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 _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] + 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 should_inline(_remote_media(remote)) + ) + ) + if not remote_urls: + return messages + data_urls: Final = await _fetch_data_urls(remote_urls) + inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True))) + 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/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/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/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/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 62b631a7671..ddab4f54d57 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, object]] = {} - - # 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 552b1549db8..2f561809940 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, @@ -485,7 +486,7 @@ class BaseLLMHTTPHandler: def completion( self, model: str, - messages: list, + messages: list[AllMessageValues], api_base: str | None, custom_llm_provider: str, model_response: ModelResponse, @@ -504,7 +505,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) @@ -519,14 +520,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( @@ -538,93 +542,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(*await asyncio.to_thread(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( @@ -638,7 +666,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, @@ -648,7 +676,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, @@ -681,7 +709,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, @@ -6754,7 +6782,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/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/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..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,6 +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 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 ( @@ -68,6 +70,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", } @@ -556,7 +559,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 @@ -1307,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], @@ -1348,13 +1368,17 @@ 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, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + ) + + 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 +1387,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 d4da18e8f6f..75b7f7f10a5 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..a4f32df46ae 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,43 @@ 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 factory, image_handling + from litellm.llms.gemini.chat import transformation as gemini_chat_transformation + + 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), + ) + + 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 893472d63ae..8fa4bd6c14d 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,7 @@ +import asyncio +import copy +import time +import uuid from unittest.mock import patch import pytest @@ -7,9 +11,13 @@ 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, + 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) @@ -107,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 @@ -207,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: @@ -268,3 +272,259 @@ 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"}}, + {"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"}}, + ], + }, + ] + snapshot = copy.deepcopy(messages) + + 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?"}, + {"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}}, + { + "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 + + +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": 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, should_inline=inline_unhinted_outside_files_api) + + assert inlined[0]["content"] == [ + {"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 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 + + +_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." + ), + 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): + 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): + attempts = [] + messages = [] + url = f"http://internal.example/{uuid.uuid4()}.png" + + for verdict in _SSRF_VERDICTS: + + async def block(client, fetched_url, verdict=verdict, **kwargs): + attempts.append(fetched_url) + raise 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_verdict_free_messages(messages, url) + + +def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): + attempts = [] + messages = [] + url = f"http://internal.example/{uuid.uuid4()}.png" + + for verdict in _SSRF_VERDICTS: + + def block(client, fetched_url, verdict=verdict, **kwargs): + attempts.append(fetched_url) + raise 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_verdict_free_messages(messages, url) + + +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"}]}, + { + "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/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/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..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 @@ -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,95 @@ 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"] + + +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 new file mode 100644 index 00000000000..a8448f5fa7a --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -0,0 +1,99 @@ +import json +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"] + + +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"] 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 5c6cc4a52da..e16855da8cb 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 @@ -17,7 +18,8 @@ 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.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, @@ -30,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 TranscriptionResponse +from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" @@ -3184,3 +3186,288 @@ 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 = [] + self.sign_threads = [] + + @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 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, + 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 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 = {} + + 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=logging_obj if logging_obj is not None else 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_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) + + 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" + + +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" 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..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,13 @@ +import json +import uuid +from unittest.mock import Mock + +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 +345,133 @@ 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"] + + +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}, + ] + + +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])