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])