diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index d6d753635f6..b3c8833811f 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -77,7 +77,7 @@ from .common_utils import ( parse_tool_call_arguments, salvage_concatenated_tool_arguments, ) -from .image_handling import convert_url_to_base64 +from .image_handling import convert_url_to_base64, is_pdf_media_type def default_pt(messages): @@ -911,8 +911,17 @@ def convert_to_anthropic_image_obj(openai_image_url: str, format: str | None) -> ) +def _anthropic_https_url_block( + image_url: str, image_format: str | None +) -> AnthropicMessagesImageParam | AnthropicMessagesDocumentParam: + source: Final = AnthropicContentParamSourceUrl(type="url", url=image_url) + if image_format and is_pdf_media_type(image_format): + return AnthropicMessagesDocumentParam(type="document", source=source) + return AnthropicMessagesImageParam(type="image", source=source) + + def create_anthropic_image_param( - image_url_input: str | dict, + image_url_input: str | Mapping[str, object], format: str | None = None, is_bedrock_invoke: bool = False, ) -> AnthropicMessagesImageParam | AnthropicMessagesDocumentParam: @@ -921,13 +930,11 @@ def create_anthropic_image_param( Supports both URL references (for HTTP/HTTPS URLs) and base64 encoding. """ - # Extract URL and format from input - if isinstance(image_url_input, str): - image_url = image_url_input - else: - image_url = image_url_input.get("url", "") - if format is None: - format = image_url_input.get("format") + url_value: Final = image_url_input if isinstance(image_url_input, str) else image_url_input.get("url") + format_value: Final = None if isinstance(image_url_input, str) else image_url_input.get("format") + image_url: Final = url_value if isinstance(url_value, str) else "" + explicit_format: Final = format if format is not None else format_value + image_format: Final = explicit_format if isinstance(explicit_format, str) else None # Check if the image URL is an HTTP/HTTPS URL if image_url.startswith("http://") or image_url.startswith("https://"): @@ -935,21 +942,14 @@ def create_anthropic_image_param( # as these providers don't support URL sources for images if is_bedrock_invoke or image_url.startswith("http://"): base64_url: Final = convert_url_to_base64(url=image_url) - image_chunk = convert_to_anthropic_image_obj(openai_image_url=base64_url, format=format) - return _anthropic_content_element_factory(image_chunk) + image_chunk = convert_to_anthropic_image_obj(openai_image_url=base64_url, format=image_format) + return anthropic_content_element_factory(image_chunk) else: - # HTTPS URL - pass directly for regular Anthropic - return AnthropicMessagesImageParam( - type="image", - source=AnthropicContentParamSourceUrl( - type="url", - url=image_url, - ), - ) + return _anthropic_https_url_block(image_url, image_format) else: # Convert to base64 for data URIs or other formats - image_chunk = convert_to_anthropic_image_obj(openai_image_url=image_url, format=format) - return _anthropic_content_element_factory(image_chunk) + image_chunk = convert_to_anthropic_image_obj(openai_image_url=image_url, format=image_format) + return anthropic_content_element_factory(image_chunk) # The following XML functions will be deprecated once JSON schema support is available on Bedrock and Vertex @@ -2255,16 +2255,16 @@ def add_cache_control_to_content( return anthropic_content_element -def _anthropic_content_element_factory( +def anthropic_content_element_factory( image_chunk: GenericImageParsingChunk, ) -> AnthropicMessagesImageParam | AnthropicMessagesDocumentParam: - if image_chunk["media_type"] == "application/pdf": + if is_pdf_media_type(image_chunk["media_type"]): _anthropic_content_element: AnthropicMessagesDocumentParam | AnthropicMessagesImageParam = ( AnthropicMessagesDocumentParam( type="document", source=AnthropicContentParamSource( type="base64", - media_type=image_chunk["media_type"], + media_type="application/pdf", data=image_chunk["data"], ), ) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 3287515c81d..ffc0e82e410 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -5,28 +5,33 @@ Helper functions to handle images passed in messages import asyncio import base64 import time -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -from itertools import chain +from functools import partial +from itertools import accumulate, chain, takewhile, tee from types import MappingProxyType -from typing import Final +from typing import Final, TypeVar +from urllib.parse import urlparse import httpx -from httpx import Response +from httpx import USE_CLIENT_DEFAULT, Response 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.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB, NON_LLM_CONNECTION_TIMEOUT from litellm.litellm_core_utils.prompt_templates.common_utils import infer_content_type_from_url_and_content 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 +_Fetched = TypeVar("_Fetched") _FETCH_ATTEMPTS: Final = 3 in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY) +pdf_probe_verdicts: Final = InMemoryCache() def _header_value(response: Response, name: str) -> str | None: @@ -209,6 +214,173 @@ def inline_remote_image_urls(media: RemoteMedia) -> bool: return media.part_type == "image_url" +_PDF_MEDIA_TYPES: Final = frozenset({"application/pdf", "application/x-pdf"}) +_GENERIC_MEDIA_TYPES: Final = frozenset({"application/octet-stream", "binary/octet-stream"}) +_PDF_SIGNATURE: Final = b"%PDF-" +_PDF_SIGNATURE_WINDOW: Final = 1024 +_PDF_PROBE_HEADERS: Final = MappingProxyType({"Range": f"bytes=0-{_PDF_SIGNATURE_WINDOW - 1}"}) +_PDF_PROBE_STATUSES: Final = frozenset({200, 206}) +_PDF_PROBE_FAILURES: Final = (httpx.HTTPError, httpx.InvalidURL, ValueError) +_ASYNC_PDF_PROBE_FAILURES: Final = (*_PDF_PROBE_FAILURES, asyncio.TimeoutError) +_PDF_PROBE_VERDICT_TTL: Final = 600 +_PDF_PROBE_FAILURE_TTL: Final = 60 +_DECODED_BODY_HEADERS: Final = frozenset({"content-encoding", "content-length", "transfer-encoding"}) + + +def _bare_media_type(media_type: str) -> str: + return media_type.split(";", 1)[0].strip().lower() + + +def is_pdf_media_type(media_type: str) -> bool: + return _bare_media_type(media_type) in _PDF_MEDIA_TYPES + + +def is_unformatted_pdf_link(url: str, image_format: str | None) -> bool: + lowered: Final = url.lower() + return not image_format and (urlparse(lowered).path.endswith(".pdf") or lowered.endswith(".pdf")) + + +def is_unformatted_pdf_link_image(media: RemoteMedia) -> bool: + image_format: Final = media.fields.get("format") + return inline_remote_image_urls(media) and is_unformatted_pdf_link( + media.url, image_format if isinstance(image_format, str) else None + ) + + +def _serves_pdf(response: Response) -> bool: + if response.status_code not in _PDF_PROBE_STATUSES: + return False + served_type: Final = _bare_media_type(_header_value(response, "Content-Type") or "") + if served_type and served_type not in _GENERIC_MEDIA_TYPES: + return served_type in _PDF_MEDIA_TYPES + return _PDF_SIGNATURE in response.content[:_PDF_SIGNATURE_WINDOW] + + +def _signature_window_chunks(chunks: Iterable[bytes], deadline: float) -> Iterator[bytes]: + reading, counting = tee(chunks) + read_so_far: Final = chain((0,), accumulate(map(len, counting))) + still_reading: Final = takewhile( + lambda total: total < _PDF_SIGNATURE_WINDOW and time.monotonic() < deadline, read_so_far + ) + return (chunk for _, chunk in zip(still_reading, reading)) + + +def _with_body_prefix(response: Response, prefix: bytes) -> Response: + return Response( + response.status_code, + headers=[(name, value) for name, value in response.headers.multi_items() if name not in _DECODED_BODY_HEADERS], + content=prefix, + request=response.request, + ) + + +@dataclass(frozen=True, slots=True) +class _BodyPrefixReader: + client: httpx.Client + deadline: float + + def get( + self, + url: str, + *, + headers: Mapping[str, str] | None = None, + follow_redirects: bool = False, + timeout: float | httpx.Timeout | None = None, + ) -> Response: + with self.client.stream( + "GET", + url, + headers=headers, + follow_redirects=follow_redirects, + timeout=USE_CLIENT_DEFAULT if timeout is None else timeout, + ) as response: + prefix: Final = b"".join(_signature_window_chunks(response.iter_bytes(), self.deadline)) + return _with_body_prefix(response, prefix[:_PDF_SIGNATURE_WINDOW]) + + +@dataclass(frozen=True, slots=True) +class _AsyncBodyPrefixReader: + client: httpx.AsyncClient + + async def get( + self, + url: str, + *, + headers: Mapping[str, str] | None = None, + follow_redirects: bool = False, + timeout: float | httpx.Timeout | None = None, + max_response_bytes: int | None = None, + ) -> Response: + async with self.client.stream( + "GET", + url, + headers=headers, + follow_redirects=follow_redirects, + timeout=USE_CLIENT_DEFAULT if timeout is None else timeout, + ) as response: + return _with_body_prefix(response, await anext(response.aiter_bytes(max_response_bytes), b"")) + + +def _cached_pdf_verdict(url: str) -> bool | None: + cached: Final = pdf_probe_verdicts.get_cache(url) + return cached if isinstance(cached, bool) else None + + +def _remember_pdf_verdict(url: str, verdict: bool, ttl: int) -> bool: + pdf_probe_verdicts.set_cache(url, verdict, ttl=ttl) + return verdict + + +def _remember_probe_response(url: str, response: Response) -> bool: + ttl: Final = _PDF_PROBE_VERDICT_TTL if response.status_code in _PDF_PROBE_STATUSES else _PDF_PROBE_FAILURE_TTL + return _remember_pdf_verdict(url, _serves_pdf(response), ttl) + + +def _failed_pdf_probe(url: str, error: Exception) -> bool: + verbose_logger.debug("PDF probe of %s failed, sending it as an image link: %s", url, error) + return _remember_pdf_verdict(url, False, _PDF_PROBE_FAILURE_TTL) + + +def serves_pdf(url: str, timeout: float = NON_LLM_CONNECTION_TIMEOUT) -> bool: + if MAX_IMAGE_URL_DOWNLOAD_SIZE_MB == 0: + return False + cached: Final = _cached_pdf_verdict(url) + if cached is not None: + return cached + try: + response: Final = safe_get( + _BodyPrefixReader(litellm.module_level_client.client, deadline=time.monotonic() + timeout), + url, + headers=dict(_PDF_PROBE_HEADERS), + timeout=timeout, + ) + except _PDF_PROBE_FAILURES as e: + return _failed_pdf_probe(url, e) + return _remember_probe_response(url, response) + + +async def async_serves_pdf(url: str, timeout: float = NON_LLM_CONNECTION_TIMEOUT) -> bool: + if MAX_IMAGE_URL_DOWNLOAD_SIZE_MB == 0: + return False + cached: Final = _cached_pdf_verdict(url) + if cached is not None: + return cached + try: + response: Final = await asyncio.wait_for( + async_safe_get( + _AsyncBodyPrefixReader(litellm.module_level_aclient.client), + url, + headers=dict(_PDF_PROBE_HEADERS), + timeout=timeout, + max_response_bytes=_PDF_SIGNATURE_WINDOW, + ), + timeout, + ) + except _ASYNC_PDF_PROBE_FAILURES as e: + return _failed_pdf_probe(url, e) + return _remember_probe_response(url, response) + + def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None: if fields.get("type") != "image_url": return None @@ -292,15 +464,40 @@ def _inline_part(part: object, data_urls: Mapping[str, str], should_inline: Call 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: +def _with_parts(message: AllMessageValues, rewrite: Callable[[object], object]) -> AllMessageValues: parts: Final = _content_parts(message) if not parts: return message - inlined_parts: Final = [_inline_part(part, data_urls, should_inline) for part in parts] - inlined_message: Final = {**message, "content": inlined_parts} - return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined + rewritten: Final = {**message, "content": [rewrite(part) for part in parts]} + return rewritten # pyright: ignore[reportReturnType] # the same message with some of its parts rewritten + + +def _inline_message( + message: AllMessageValues, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool] +) -> AllMessageValues: + return _with_parts(message, partial(_inline_part, data_urls=data_urls, should_inline=should_inline)) + + +def _pdf_formatted_part(part: object, pdf_links: frozenset[str], should_probe: Callable[[RemoteMedia], bool]) -> object: + remote: Final = _parse_remote_part(part) + match remote: + case _RemoteImage(fields, image_url, url) if url in pdf_links and should_probe(_remote_media(remote)): + image_url_fields: Final = image_url if image_url is not None else _NO_FIELDS + return {**fields, "image_url": {**image_url_fields, "url": url, "format": "application/pdf"}} + case _: + return part + + +def _with_pdf_formats( + messages: Iterable[AllMessageValues], + probed: Iterable[tuple[str, bool]], + should_probe: Callable[[RemoteMedia], bool], +) -> tuple[AllMessageValues, ...]: + pdf_links: Final = frozenset(url for url, is_pdf in probed if is_pdf) + if not pdf_links: + return tuple(messages) + rewrite: Final = partial(_pdf_formatted_part, pdf_links=pdf_links, should_probe=should_probe) + return tuple(_with_parts(message, rewrite) for message in messages) async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore, timeout: float | httpx.Timeout | None) -> str: @@ -308,11 +505,12 @@ async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore, timeout: float return await async_convert_url_to_base64(url, timeout=timeout) -async def async_convert_urls_to_base64( - remote_urls: tuple[str, ...], timeout: float | httpx.Timeout | None = None -) -> tuple[str, ...]: - in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) - fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight, timeout)) for url in remote_urls) +async def _serves_pdf_within(url: str, in_flight: asyncio.Semaphore) -> bool: + async with in_flight: + return await async_serves_pdf(url) + + +async def _gather_or_cancel_all(fetches: tuple["asyncio.Task[_Fetched]", ...]) -> tuple[_Fetched, ...]: try: return tuple(await asyncio.gather(*fetches)) except BaseException: @@ -322,6 +520,15 @@ async def async_convert_urls_to_base64( raise +async def async_convert_urls_to_base64( + remote_urls: tuple[str, ...], timeout: float | httpx.Timeout | None = None +) -> tuple[str, ...]: + in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) + return await _gather_or_cancel_all( + tuple(asyncio.create_task(_fetch_data_url(url, in_flight, timeout)) for url in remote_urls) + ) + + def _remote_urls_to_inline( messages: Iterable[AllMessageValues], should_inline: Callable[[RemoteMedia], bool] ) -> tuple[str, ...]: @@ -351,3 +558,27 @@ async def async_inline_remote_media( data_urls: Final = await async_convert_urls_to_base64(remote_urls) inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True))) return [_inline_message(message, inlined, should_inline) for message in messages] + + +def format_pdf_links( + messages: Sequence[AllMessageValues], should_probe: Callable[[RemoteMedia], bool] +) -> tuple[AllMessageValues, ...]: + links: Final = _remote_urls_to_inline(messages, should_probe) + if not links: + return tuple(messages) + with ThreadPoolExecutor(max_workers=min(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES, len(links))) as probes: + verdicts: Final = tuple(probes.map(serves_pdf, links)) + return _with_pdf_formats(messages, zip(links, verdicts, strict=True), should_probe) + + +async def async_format_pdf_links( + messages: Sequence[AllMessageValues], should_probe: Callable[[RemoteMedia], bool] +) -> tuple[AllMessageValues, ...]: + links: Final = _remote_urls_to_inline(messages, should_probe) + if not links: + return tuple(messages) + in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES) + verdicts: Final = await _gather_or_cancel_all( + tuple(asyncio.create_task(_serves_pdf_within(url, in_flight)) for url in links) + ) + return _with_pdf_formats(messages, zip(links, verdicts, strict=True), should_probe) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 7f688ae733b..7e3fda06b1a 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -34,6 +34,7 @@ from litellm.llms.anthropic.messages.adapters.transformation import ( ) from litellm.llms.anthropic.messages.convertible_blocks import ( keep_non_string_tool_result_texts, + keep_url_documents, scannable_tool_result_request, ) from litellm.llms.anthropic.messages.native.utils import filtered_safeguard_results, redacted_stop_details @@ -1187,7 +1188,10 @@ class AnthropicMessagesHandler(BaseTranslation): cast(Sequence[object], converted), # cast-ok: Anthropic message payload original_tool_names, ) - data["messages"] = keep_non_string_tool_result_texts(converted, data.get("messages")) + original_messages: Final = data.get("messages") + data["messages"] = keep_url_documents( + keep_non_string_tool_result_texts(converted, original_messages), original_messages + ) @staticmethod def _extract_midturn_system_text( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 2e908e66626..1977e3e2a7d 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -27,6 +27,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, RESPONSE_FORMAT_TOOL_NAME, ) +from litellm.litellm_core_utils.asyncify import is_event_loop_running from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, map_finish_reason from litellm.litellm_core_utils.prompt_templates.common_utils import ( DEFS_MAX_INLINED_BYTES, @@ -47,8 +48,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( RemoteMedia, + async_format_pdf_links, async_inline_remote_media, + format_pdf_links, inline_remote_image_urls, + is_unformatted_pdf_link_image, ) from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( message_field, @@ -2333,6 +2337,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def inlines_remote_media(self, media: RemoteMedia) -> bool: return inline_remote_image_urls(media) and media.url.startswith("http://") + def probes_pdf_link(self, media: RemoteMedia) -> bool: + return is_unformatted_pdf_link_image(media) and not self.inlines_remote_media(media) + + def _with_pdf_links_formatted(self, messages: Sequence[AllMessageValues]) -> Sequence[AllMessageValues]: + return messages if is_event_loop_running() else format_pdf_links(messages, should_probe=self.probes_pdf_link) + async def async_transform_request( self, model: str, @@ -2341,9 +2351,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): 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: Final = await async_inline_remote_media(messages, should_inline=self.inlines_remote_media) return self.transform_request( model=model, - messages=await async_inline_remote_media(messages, should_inline=self.inlines_remote_media), + messages=list(await async_format_pdf_links(inlined, should_probe=self.probes_pdf_link)), optional_params=optional_params, litellm_params=litellm_params, headers=headers, @@ -2477,7 +2488,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): try: anthropic_messages = anthropic_messages_pt( model=model, - messages=list(conversation), + messages=list(self._with_pdf_links_formatted(conversation)), llm_provider=self._resolved_provider, drop_params=drop_params, ) diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 1c8f0c206cd..3034927c612 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -49,7 +49,10 @@ _PLAIN_MESSAGE_KEYS: Final = frozenset({"role", "content"}) def _count_image(block: JsonValue) -> JsonValue: if not isinstance(block, dict) or block.get("type") != "image_url": return block - from litellm.litellm_core_utils.prompt_templates.factory import convert_to_anthropic_image_obj + from litellm.litellm_core_utils.prompt_templates.factory import ( + anthropic_content_element_factory, + convert_to_anthropic_image_obj, + ) image_block: Final = _IMAGE_BLOCK.validate_python(block) image_url: Final = image_block["image_url"] @@ -57,7 +60,7 @@ def _count_image(block: JsonValue) -> JsonValue: openai_image_url=image_url if isinstance(image_url, str) else image_url["url"], format=image_url.get("format") if isinstance(image_url, dict) else None, ) - image: Final = _COUNT_REQUEST.validate_python({"type": "image", "source": source}) + image: Final = _COUNT_REQUEST.validate_python(anthropic_content_element_factory(source)) return {**{key: value for key, value in block.items() if key not in {"type", "image_url"}}, **image} diff --git a/litellm/llms/anthropic/messages/convertible_blocks.py b/litellm/llms/anthropic/messages/convertible_blocks.py index 09d662da543..b98b4cf3f75 100644 --- a/litellm/llms/anthropic/messages/convertible_blocks.py +++ b/litellm/llms/anthropic/messages/convertible_blocks.py @@ -314,6 +314,76 @@ def keep_non_string_tool_result_texts(rewritten: Sequence[object], original: obj return [_with_user_blocks(message, partial(_restored_block, kept=kept_in_order)) for message in rewritten] +def _url_source_url(block: object) -> str | None: + match as_str_mapping(block): + case {"type": "image" | "document", "source": {"type": "url", "url": str(url)}}: + return url + case _: + return None + + +def _is_document(block: object) -> bool: + match as_str_mapping(block): + case {"type": "document"}: + return True + case _: + return False + + +@dataclass(frozen=True, slots=True, order=True) +class _UrlBlockPlace: + in_tool_result: bool + url: str + + +def _url_sourced_blocks(block: object, in_tool_result: bool = False) -> tuple[tuple[_UrlBlockPlace, object], ...]: + url: Final = _url_source_url(block) + if url is not None: + return ((_UrlBlockPlace(in_tool_result, url), block),) + match as_str_mapping(block): + case {"type": "tool_result", "content": list() as content}: + in_this_tool_result: Final = partial(_url_sourced_blocks, in_tool_result=True) + return tuple(chain.from_iterable(map(in_this_tool_result, _list_items(content)))) + case _: + return () + + +def _entry_place(entry: tuple[_UrlBlockPlace, object]) -> _UrlBlockPlace: + return entry[0] + + +def _entry_block(entry: tuple[_UrlBlockPlace, object]) -> object: + return entry[1] + + +def _restored_url_document( + block: object, originals: Mapping[_UrlBlockPlace, Iterator[object]], in_tool_result: bool = False +) -> object: + url: Final = _url_source_url(block) + if url is not None: + remaining: Final = originals.get(_UrlBlockPlace(in_tool_result, url)) + original: Final = None if remaining is None else next(remaining, None) + return original if _is_document(original) else block + match as_str_mapping(block): + case {"type": "tool_result", "content": list() as content} as tool_result: + restore: Final = partial(_restored_url_document, originals=originals, in_tool_result=True) + return {**tool_result, "content": list(map(restore, _list_items(content)))} + case _: + return block + + +def keep_url_documents(rewritten: Sequence[object], original: object) -> Sequence[object]: + blocks: Final = chain.from_iterable(map(_user_content_blocks, _list_items(original))) + by_place: Final = tuple(sorted(chain.from_iterable(map(_url_sourced_blocks, blocks)), key=_entry_place)) + if not any(_is_document(block) for _, block in by_place): + return rewritten + originals_in_order: Final = MappingProxyType( + {place: iter(tuple(map(_entry_block, group))) for place, group in groupby(by_place, key=_entry_place)} + ) + restore: Final = partial(_restored_url_document, originals=originals_in_order) + return [_with_user_blocks(message, restore) for message in rewritten] + + def _without_non_string_text(block: object) -> object: match as_str_mapping(block): case {"type": "tool_result", "content": list()} as tool_result: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py index e2f12b123de..37acb5423ff 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py @@ -13,8 +13,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( + RemoteMedia, async_inline_remote_media, convert_url_to_base64, + inline_every_remote_url, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import ( @@ -117,6 +119,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): def should_strip_billing_metadata(self) -> bool: return True + def inlines_remote_media(self, media: RemoteMedia) -> bool: + return inline_every_remote_url(media) + def get_supported_openai_params(self, model: str) -> list[str]: return AnthropicConfig.get_supported_openai_params(self, model) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index ed7dd5ec146..914eb1f2508 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -216,6 +216,7 @@ def _reset_aws_auth_caches() -> None: def _flush_client_caches() -> None: litellm.in_memory_llm_clients_cache.flush_cache() image_handling_module.in_memory_cache.flush_cache() + image_handling_module.pdf_probe_verdicts.flush_cache() _reset_aws_auth_caches() diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 8a755d47817..d65f094971a 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1298,6 +1298,71 @@ def test_anthropic_messages_pt_uses_pdf_format_for_generic_mime_data_uri(llm_pro } +@pytest.mark.parametrize( + ("image_url", "expected_type"), + ( + ({"url": "https://example.com/download?id=7", "format": "application/pdf"}, "document"), + ({"url": "https://example.com/download?id=7", "format": "Application/PDF; charset=binary"}, "document"), + ({"url": "https://example.com/files/report.pdf", "format": "image/png"}, "image"), + ({"url": "https://example.com/pdf/photo.png"}, "image"), + ({"url": "https://example.com/pdf/photo.png", "format": 5}, "image"), + ({"url": "https://example.com/pdf/photo.png", "format": ["application/pdf"]}, "image"), + ), +) +def test_anthropic_messages_pt_picks_block_type_for_https_image_url( + image_url: dict[str, object], expected_type: str +) -> None: + result: Final = anthropic_messages_pt( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize"}, + {"type": "image_url", "image_url": image_url}, + ], + } + ], + model="claude-opus-5-5", + llm_provider="anthropic", + ) + + assert result[0]["content"][1] == { + "type": expected_type, + "source": {"type": "url", "url": image_url["url"]}, + } + + +@pytest.mark.parametrize( + "image_url", + ( + {"url": "data:Application/PDF;base64,JVBERi0xLjQK"}, + {"url": "data:application/x-pdf;base64,JVBERi0xLjQK"}, + {"url": "data:application/octet-stream;base64,JVBERi0xLjQK", "format": "application/pdf; charset=binary"}, + ), +) +def test_anthropic_messages_pt_normalizes_a_pdf_media_type_with_case_or_parameters( + image_url: dict[str, str], +) -> None: + result: Final = anthropic_messages_pt( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize"}, + {"type": "image_url", "image_url": image_url}, + ], + } + ], + model="claude-opus-5-5", + llm_provider="anthropic", + ) + + assert result[0]["content"][1] == { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + } + + def test_anthropic_messages_pt_keeps_png_image_url_as_image() -> None: image_data = "iVBORw0KGgo=" result: Final = anthropic_messages_pt( diff --git a/tests/unit/litellm_core_utils/test_image_handling.py b/tests/unit/litellm_core_utils/test_image_handling.py index d1bebde9917..2d4cd4b9a6d 100644 --- a/tests/unit/litellm_core_utils/test_image_handling.py +++ b/tests/unit/litellm_core_utils/test_image_handling.py @@ -1,23 +1,36 @@ import asyncio import base64 import copy +import threading import time import uuid +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final +from unittest.mock import MagicMock import pytest from httpx import Request, Response import litellm +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.litellm_core_utils import url_utils 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_format_pdf_links, async_inline_remote_media, + async_serves_pdf, convert_url_to_base64, + format_pdf_links, inline_remote_media, + is_unformatted_pdf_link_image, + serves_pdf, ) from litellm.litellm_core_utils.url_utils import SSRFError +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @pytest.fixture(autouse=True) @@ -695,3 +708,295 @@ async def test_async_inline_remote_media_raises_image_fetch_error_when_the_fetch with pytest.raises(litellm.ImageFetchError, match="Status code: 404"): await async_inline_remote_media([{"role": "user", "content": [{"type": "image_url", "image_url": url}]}]) + + +PDF_BODY: Final = b"%PDF-1.7\n" + b"0" * 4096 +PNG_BODY: Final = b"\x89PNG\r\n\x1a\n" + b"0" * 64 + + +@dataclass(frozen=True, slots=True) +class _Served: + status: int = 200 + content_type: str | None = "application/octet-stream" + body: bytes = PDF_BODY + honors_range: bool = True + barrier: threading.Barrier | None = None + declared_length: int | None = None + hold_open_until: threading.Event | None = None + respond_after: threading.Event | None = None + drip_every: float | None = None + + +@dataclass(frozen=True, slots=True) +class _ProbeServer: + base_url: str + routes: dict[str, _Served] + seen: list[tuple[str, str | None]] + + def serve(self, path: str, served: _Served) -> str: + self.routes[path] = served + return f"{self.base_url}{path}" + + +@pytest.fixture +def probe_server(monkeypatch): + monkeypatch.setattr(litellm, "user_url_validation", False) + monkeypatch.setattr(image_handling, "safe_get", url_utils.safe_get) + monkeypatch.setattr(litellm, "module_level_client", HTTPHandler()) + monkeypatch.setattr(litellm, "module_level_aclient", AsyncHTTPHandler()) + routes: Final[dict[str, _Served]] = {} + seen: Final[list[tuple[str, str | None]]] = [] + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + served = routes[self.path] + requested_range = self.headers.get("Range") + seen.append((self.path, requested_range)) + if served.barrier is not None: + served.barrier.wait() + if served.respond_after is not None: + served.respond_after.wait(timeout=30) + ranged = served.honors_range and requested_range == "bytes=0-1023" + body = served.body[:1024] if ranged else served.body + self.send_response(206 if ranged and served.status == 200 else served.status) + if served.content_type is not None: + self.send_header("Content-Type", served.content_type) + self.send_header("Content-Length", str(served.declared_length or len(body))) + self.end_headers() + if served.drip_every is None: + self.wfile.write(body) + else: + for byte in body: + self.wfile.write(bytes([byte])) + self.wfile.flush() + time.sleep(served.drip_every) + if served.hold_open_until is not None: + served.hold_open_until.wait(timeout=30) + + def log_message(self, *args): + pass + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield _ProbeServer(f"http://127.0.0.1:{server.server_port}", routes, seen) + server.shutdown() + server.server_close() + + +async def _serves_pdf(mode: str, url: str, **kwargs) -> bool: + return serves_pdf(url, **kwargs) if mode == "sync" else await async_serves_pdf(url, **kwargs) + + +async def _format_pdf_links(mode, messages): + if mode == "sync": + return format_pdf_links(messages, should_probe=is_unformatted_pdf_link_image) + return await async_format_pdf_links(messages, should_probe=is_unformatted_pdf_link_image) + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_asks_for_only_the_leading_kilobyte(probe_server, mode): + url = probe_server.serve("/files/report.pdf", _Served()) + + assert await _serves_pdf(mode, url) is True + assert probe_server.seen == [("/files/report.pdf", "bytes=0-1023")] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize( + ("served", "expected"), + ( + (_Served(content_type="application/pdf", body=PNG_BODY), True), + (_Served(content_type="Application/X-PDF; charset=binary", body=b""), True), + (_Served(content_type="binary/octet-stream", body=PDF_BODY), True), + (_Served(content_type=None, body=PDF_BODY), True), + (_Served(content_type="application/octet-stream", body=b"\n" * 1000 + PDF_BODY), True), + (_Served(content_type="application/octet-stream", body=b"\n" * 1024 + PDF_BODY), False), + (_Served(content_type="binary/octet-stream", body=PNG_BODY), False), + (_Served(content_type="image/png", body=PDF_BODY), False), + (_Served(content_type="text/html", body=b"%PDF-"), False), + (_Served(status=404, content_type="application/pdf"), False), + (_Served(honors_range=False), True), + ), + ids=( + "served-pdf", + "served-x-pdf", + "generic-type-pdf-bytes", + "no-type-pdf-bytes", + "signature-inside-the-first-kilobyte", + "signature-past-the-first-kilobyte", + "generic-type-png-bytes", + "served-png", + "served-html", + "not-found", + "range-ignored", + ), +) +async def test_serves_pdf_trusts_a_specific_served_type_and_sniffs_a_generic_one(probe_server, mode, served, expected): + url = probe_server.serve(f"/{uuid.uuid4()}/report.pdf", served) + + assert await _serves_pdf(mode, url) is expected + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_makes_no_request_when_url_downloads_are_disabled(probe_server, monkeypatch, mode): + monkeypatch.setattr(image_handling, "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", 0) + url = probe_server.serve("/report.pdf", _Served(content_type="application/pdf")) + + assert await _serves_pdf(mode, url) is False + assert probe_server.seen == [] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_is_false_when_the_host_refuses_the_connection(probe_server, mode): + assert await _serves_pdf(mode, "http://127.0.0.1:1/report.pdf") is False + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_is_false_without_a_request_when_the_ssrf_check_blocks_the_link(probe_server, monkeypatch, mode): + monkeypatch.setattr(litellm, "user_url_validation", True) + url = probe_server.serve("/report.pdf", _Served(content_type="application/pdf")) + + assert await _serves_pdf(mode, url) is False + assert probe_server.seen == [] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_stops_reading_after_the_leading_kilobyte_when_the_host_ignores_range(probe_server, mode): + rest_of_body_sent = threading.Event() + url = probe_server.serve( + "/report.pdf", _Served(honors_range=False, declared_length=1024**3, hold_open_until=rest_of_body_sent) + ) + + try: + assert await _serves_pdf(mode, url) is True + finally: + rest_of_body_sent.set() + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_gives_up_within_its_timeout_when_the_host_sends_one_byte_at_a_time(probe_server, mode): + url = probe_server.serve("/report.pdf", _Served(body=b"0" * 1024, honors_range=False, drip_every=0.02)) + started = time.monotonic() + + assert await _serves_pdf(mode, url, timeout=0.5) is False + assert time.monotonic() - started < 3 + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_reuses_a_link_verdict_instead_of_probing_the_link_again(probe_server, mode): + pdf_link = probe_server.serve("/a/report.pdf", _Served()) + png_link = probe_server.serve("/a/page.pdf", _Served(content_type="image/png", body=PNG_BODY)) + + verdicts = [await _serves_pdf(mode, url) for url in (pdf_link, png_link, pdf_link, png_link)] + + assert verdicts == [True, False, True, False] + assert [path for path, _ in probe_server.seen] == ["/a/report.pdf", "/a/page.pdf"] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_serves_pdf_remembers_a_timed_out_probe_so_the_next_turn_does_not_wait_again(probe_server, mode): + host_answers = threading.Event() + url = probe_server.serve("/report.pdf", _Served(respond_after=host_answers)) + + try: + first = await _serves_pdf(mode, url, timeout=0.3) + started = time.monotonic() + second = await _serves_pdf(mode, url, timeout=0.3) + second_took = time.monotonic() - started + finally: + host_answers.set() + + assert (first, second) == (False, False) + assert second_took < 0.1 + assert probe_server.seen == [("/report.pdf", "bytes=0-1023")] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +@pytest.mark.parametrize("status", (404, 503)) +async def test_serves_pdf_probes_a_link_again_a_minute_after_its_host_answered_an_error( + probe_server, monkeypatch, mode, status +): + clock: Final = MagicMock(return_value=1_000.0) + monkeypatch.setattr(image_handling, "pdf_probe_verdicts", InMemoryCache(clock=clock)) + url = probe_server.serve("/report.pdf", _Served(status=status)) + + while_down = await _serves_pdf(mode, url) + probe_server.serve("/report.pdf", _Served()) + clock.return_value = 1_030.0 + within_the_minute = await _serves_pdf(mode, url) + clock.return_value = 1_061.0 + after_the_minute = await _serves_pdf(mode, url) + + assert (while_down, within_the_minute, after_the_minute) == (False, False, True) + assert len(probe_server.seen) == 2 + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_format_pdf_links_marks_only_links_that_serve_a_pdf(probe_server, mode): + pdf_link = probe_server.serve("/a/report.pdf", _Served()) + query_link = probe_server.serve("/a/download?file=summary.pdf", _Served(content_type="application/pdf")) + png_at_pdf_link = probe_server.serve("/a/page.pdf?fm=png", _Served(content_type="image/png", body=PNG_BODY)) + png_link = f"{probe_server.base_url}/a/photo.png" + messages = [ + {"role": "system", "content": "read the files"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "compare"}, + {"type": "image_url", "image_url": {"url": pdf_link, "detail": "high"}, "cache_control": {"type": "ephemeral"}}, + {"type": "image_url", "image_url": query_link}, + {"type": "image_url", "image_url": {"url": png_at_pdf_link}}, + {"type": "image_url", "image_url": {"url": png_link}}, + {"type": "image_url", "image_url": {"url": pdf_link, "format": "image/png"}}, + ], + }, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": pdf_link}}]}, + ] + original = copy.deepcopy(messages) + + formatted = await _format_pdf_links(mode, messages) + + assert messages == original + assert formatted == ( + messages[0], + { + "role": "user", + "content": [ + {"type": "text", "text": "compare"}, + { + "type": "image_url", + "image_url": {"url": pdf_link, "detail": "high", "format": "application/pdf"}, + "cache_control": {"type": "ephemeral"}, + }, + {"type": "image_url", "image_url": {"url": query_link, "format": "application/pdf"}}, + {"type": "image_url", "image_url": {"url": png_at_pdf_link}}, + {"type": "image_url", "image_url": {"url": png_link}}, + {"type": "image_url", "image_url": {"url": pdf_link, "format": "image/png"}}, + ], + }, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": pdf_link, "format": "application/pdf"}}]}, + ) + assert sorted(path for path, _ in probe_server.seen) == ["/a/download?file=summary.pdf", "/a/page.pdf?fm=png", "/a/report.pdf"] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_format_pdf_links_probes_every_link_at_once(probe_server, mode): + both_in_flight = threading.Barrier(2, timeout=5) + links = [probe_server.serve(f"/{name}.pdf", _Served(barrier=both_in_flight)) for name in ("first", "second")] + messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": link}} for link in links]}] + + formatted = await _format_pdf_links(mode, messages) + + assert [part["image_url"]["format"] for part in formatted[0]["content"]] == ["application/pdf", "application/pdf"] + + +@pytest.mark.parametrize("mode", ("sync", "async")) +async def test_format_pdf_links_leaves_messages_without_pdf_links_alone(probe_server, mode): + messages = [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": f"{probe_server.base_url}/a.png"}}]}, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:application/pdf;base64,JVBERi0="}}]}, + ] + + assert await _format_pdf_links(mode, messages) == tuple(messages) + assert probe_server.seen == [] diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index ce14adab4a0..7e2786327dc 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -2,6 +2,7 @@ import json import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import chain +from collections.abc import Callable from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -13,7 +14,8 @@ from openai.types.chat import ChatCompletionChunk import litellm from litellm._uuid import uuid -from litellm.constants import RESPONSE_FORMAT_TOOL_NAME +from litellm.constants import NON_LLM_CONNECTION_TIMEOUT, RESPONSE_FORMAT_TOOL_NAME +from litellm.litellm_core_utils.prompt_templates import image_handling from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call, make_sync_call from litellm.llms.base_llm.base_utils import response_format_schema, unwrap_response_format_envelope @@ -253,6 +255,7 @@ def test_anthropic_tool_call_error_is_mapped_from_the_provider_response(): async def test_anthropic_async_completion_inlines_http_images_off_the_event_loop(async_only_image_fetch): http_image_url = f"http://img.example/{uuid.uuid4()}.png" https_image_url = f"https://img.example/{uuid.uuid4()}.png" + https_pdf_link = f"https://img.example/{uuid.uuid4()}/page.pdf?w=600" captured = {} def handle(request): @@ -282,6 +285,7 @@ async def test_anthropic_async_completion_inlines_http_images_off_the_event_loop {"type": "text", "text": "What colour is this?"}, {"type": "image_url", "image_url": {"url": http_image_url}}, {"type": "image_url", "image_url": {"url": https_image_url}}, + {"type": "image_url", "image_url": {"url": https_pdf_link}}, ], } ], @@ -290,11 +294,193 @@ async def test_anthropic_async_completion_inlines_http_images_off_the_event_loop ) assert response.choices[0].message.content == "Green" - assert async_only_image_fetch.fetched == [http_image_url] + assert async_only_image_fetch.fetched == [http_image_url, https_pdf_link] sources = [part["source"] for part in captured["body"]["messages"][0]["content"] if part["type"] == "image"] assert sources == [ {"type": "base64", "media_type": "image/png", "data": async_only_image_fetch.base64_png}, {"type": "url", "url": https_image_url}, + {"type": "url", "url": https_pdf_link}, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("download_limit_mb", (0, 50), ids=("downloads-disabled", "link-unreachable")) +async def test_anthropic_async_completion_sends_a_pdf_path_link_it_cannot_download_as_a_link( + monkeypatch: pytest.MonkeyPatch, download_limit_mb: int +) -> None: + pdf_link: Final = f"https://img.example/{uuid.uuid4()}/page.pdf?fm=png" + captured: Final[dict[str, object]] = {} + + async def unreachable(client: object, url: str, **kwargs: object) -> httpx.Response: + return httpx.Response(404, request=httpx.Request("GET", url)) + + def handle(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "A chart"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + monkeypatch.setattr(image_handling, "async_safe_get", unreachable) + monkeypatch.setattr(image_handling, "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", download_limit_mb) + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response: Final = await litellm.acompletion( + model="anthropic/claude-haiku-5-5", + messages=[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": pdf_link}}]}], + client=client, + api_key="sk-ant-unit-test", + ) + + assert response.choices[0].message.content == "A chart" + assert captured["body"]["messages"][0]["content"] == [{"type": "image", "source": {"type": "url", "url": pdf_link}}] + + +@pytest.mark.asyncio +async def test_anthropic_async_completion_gives_up_on_a_pdf_link_within_the_non_llm_connection_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pdf_link: Final = f"https://unreachable.example/{uuid.uuid4()}/report.pdf" + fetch_timeouts: Final[list[object]] = [] + captured: Final[dict[str, object]] = {} + + async def unreachable(client: object, url: str, **kwargs: object) -> httpx.Response: + fetch_timeouts.append(kwargs.get("timeout")) + raise httpx.ConnectTimeout("no route to host") + + def handle(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "A chart"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + monkeypatch.setattr(image_handling, "async_safe_get", unreachable) + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + await litellm.acompletion( + model="anthropic/claude-haiku-5-5", + messages=[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": pdf_link}}]}], + client=client, + api_key="sk-ant-unit-test", + ) + + assert captured["body"]["messages"][0]["content"] == [{"type": "image", "source": {"type": "url", "url": pdf_link}}] + assert fetch_timeouts == [NON_LLM_CONNECTION_TIMEOUT] + + +def _anthropic_reply(captured: dict[str, object]) -> Callable[[httpx.Request], httpx.Response]: + def handle(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "A report"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + return handle + + +def _served_pdf_slice(url: str) -> httpx.Response: + return httpx.Response( + 206, content=b"%PDF-1.7\n", headers={"content-type": "binary/octet-stream"}, request=httpx.Request("GET", url) + ) + + +@pytest.mark.asyncio +async def test_anthropic_async_completion_sends_a_link_that_serves_a_pdf_as_a_url_document( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pdf_link: Final = f"https://files.example/{uuid.uuid4()}/download?file=report.pdf" + probes: Final[list[object]] = [] + captured: Final[dict[str, object]] = {} + + async def serve_pdf_slice(client: object, url: str, **kwargs: object) -> httpx.Response: + probes.append(kwargs.get("headers")) + return _served_pdf_slice(url) + + monkeypatch.setattr(image_handling, "async_safe_get", serve_pdf_slice) + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(_anthropic_reply(captured))) + + await litellm.acompletion( + model="anthropic/claude-haiku-5-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": pdf_link}, "cache_control": {"type": "ephemeral"}} + ], + } + ], + client=client, + api_key="sk-ant-unit-test", + ) + + assert captured["body"]["messages"][0]["content"] == [ + {"type": "document", "source": {"type": "url", "url": pdf_link}, "cache_control": {"type": "ephemeral"}} + ] + assert probes == [{"Range": "bytes=0-1023"}] + + +def test_anthropic_sync_completion_sends_a_link_that_serves_a_pdf_as_a_url_document( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pdf_link: Final = f"https://files.example/{uuid.uuid4()}/report.pdf" + png_at_pdf_link: Final = f"https://files.example/{uuid.uuid4()}/page.pdf?fm=png" + captured: Final[dict[str, object]] = {} + + def serve(client: object, url: str, **kwargs: object) -> httpx.Response: + if url == pdf_link: + return _served_pdf_slice(url) + return httpx.Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=httpx.Request("GET", url)) + + monkeypatch.setattr(image_handling, "safe_get", serve) + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_anthropic_reply(captured)))) + + litellm.completion( + model="anthropic/claude-haiku-5-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": pdf_link}}, + {"type": "image_url", "image_url": {"url": png_at_pdf_link}}, + ], + } + ], + client=client, + api_key="sk-ant-unit-test", + ) + + assert captured["body"]["messages"][0]["content"] == [ + {"type": "document", "source": {"type": "url", "url": pdf_link}}, + {"type": "image", "source": {"type": "url", "url": png_at_pdf_link}}, ] diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index aa7fa2727b4..13c17ebd2bf 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -50,20 +50,30 @@ COMPACTED_HISTORY = [ "config_type", (AnthropicCountTokensConfig, AzureAIAnthropicCountTokensConfig) ) @pytest.mark.parametrize( - ("image_url", "source"), + ("image_url", "block"), ( - ("data:image/png;base64,aW1hZ2U=", {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}), - ({"url": "data:image/png;base64,aW1hZ2U="}, {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}), + ( + "data:image/png;base64,aW1hZ2U=", + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + ), + ( + {"url": "data:image/png;base64,aW1hZ2U="}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + ), ( {"url": "data:image/png;base64,aW1hZ2U=", "format": "image/jpeg", "detail": "high"}, - {"type": "base64", "media_type": "image/jpeg", "data": "aW1hZ2U="}, + {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "aW1hZ2U="}}, + ), + ( + "data:application/pdf;base64,JVBERi0=", + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0="}}, ), ), ) def test_count_translates_openai_images_without_mutating_input( config_type: type[AnthropicCountTokensConfig], image_url: str | dict[str, JsonValue], - source: dict[str, JsonValue], + block: dict[str, JsonValue], ) -> None: cache_control: Final[dict[str, JsonValue]] = {"type": "ephemeral"} messages: Final[list[dict[str, JsonValue]]] = [{ @@ -81,7 +91,7 @@ def test_count_translates_openai_images_without_mutating_input( "model": "claude-opus-5-5", "messages": [{ "role": "user", "content": [ {"type": "text", "text": "Count this image"}, - {"type": "image", "source": source, "cache_control": cache_control}, + {**block, "cache_control": cache_control}, ], }], } diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py index 05d755b1215..34a4cf5225b 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py @@ -18,7 +18,7 @@ from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.unit.conftest import NativeStructuredOutputKey from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE @@ -1052,6 +1052,64 @@ async def test_bedrock_invoke_claude_async_completion_inlines_document_url_sourc } in captured["body"]["messages"][0]["content"] +def test_bedrock_invoke_claude_sync_completion_downloads_a_pdf_link_once_and_sends_it_as_a_base64_document(monkeypatch): + from litellm.litellm_core_utils.prompt_templates import image_handling + + pdf_url = f"https://docs.example/{uuid.uuid4()}/report.pdf" + pdf_bytes = b"%PDF-1.7\n" + b"0" * 64 + fetched = [] + captured = {} + + def serve_pdf(client, url, **kwargs): + fetched.append(url) + return httpx.Response( + 200, content=pdf_bytes, headers={"content-type": "application/pdf"}, request=httpx.Request("GET", url) + ) + + 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}, + }, + ) + + monkeypatch.setattr(image_handling, "safe_get", serve_pdf) + client = HTTPHandler() + client.client = httpx.Client(transport=httpx.MockTransport(handle)) + + response = litellm.completion( + model="bedrock/invoke/us.anthropic.claude-sonnet-5", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this document?"}, + {"type": "image_url", "image_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 fetched == [pdf_url] + assert captured["body"]["messages"][0]["content"][1] == { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": base64.b64encode(pdf_bytes).decode()}, + } + + @pytest.mark.parametrize( "model, expected_betas", [ diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py index d2e5b407e30..1baf9504208 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py @@ -200,3 +200,78 @@ async def test_write_back_keeps_real_tool_results_under_modify_params(): assert "skipped" not in serialized assert "Please continue" not in serialized assert [m["role"] for m in written] == ["user", "assistant", "user", "user"] + + +@pytest.mark.asyncio +async def test_write_back_sends_the_client_url_blocks_as_the_client_sent_them(async_only_image_fetch): + image_link = "https://img.example/presskit/page.pdf?fm=png" + document_link = "https://files.example/reports/q3.pdf" + messages = [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu_1", "name": "fetch_report", "input": {}}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": [ + {"type": "document", "source": {"type": "url", "url": document_link}, "title": "Q3 attachment"} + ], + }, + {"type": "text", "text": "compare these"}, + {"type": "image", "source": {"type": "url", "url": image_link}}, + { + "type": "document", + "source": {"type": "url", "url": document_link}, + "title": "Q3 report", + "context": "Board copy", + "citations": {"enabled": True}, + "cache_control": {"type": "ephemeral"}, + }, + {"type": "document", "source": {"type": "url", "url": document_link}, "title": "Q3 report, again"}, + ], + }, + ] + + written = await _write_back_identity([dict(m) for m in messages]) + + assert async_only_image_fetch.fetched == [] + assert [m["content"] for m in written[1:]] == [messages[1]["content"][:1], messages[1]["content"][1:]] + + +@pytest.mark.asyncio +async def test_write_back_keeps_each_url_document_fields_when_a_tool_result_after_it_has_the_same_link( + async_only_image_fetch, +): + document_link = "https://files.example/reports/q3.pdf" + messages = [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu_1", "name": "fetch_report", "input": {}}], + }, + { + "role": "user", + "content": [ + { + "type": "document", + "source": {"type": "url", "url": document_link}, + "title": "Q3 report", + "context": "Board copy", + }, + { + "type": "tool_result", + "tool_use_id": "tu_1", + "content": [ + {"type": "document", "source": {"type": "url", "url": document_link}, "title": "Q3 attachment"} + ], + }, + ], + }, + ] + + written = await _write_back_identity([dict(m) for m in messages]) + + assert [m["content"] for m in written[1:]] == [messages[1]["content"][1:], messages[1]["content"][:1]]