From 2fe7f8f66c23756e400043d8a85970a665b810b1 Mon Sep 17 00:00:00 2001 From: Jerry-Xin Date: Sat, 7 Mar 2026 17:59:38 +0800 Subject: [PATCH] fix(vertex_ai): Python 3.9 compatibility and concurrent image downloads - Use Optional[str] instead of str | None for Python 3.9 compatibility - Use asyncio.gather to download multiple images concurrently instead of sequentially, improving performance for messages with multiple images Co-Authored-By: Claude Opus 4.5 --- .../prompt_templates/image_handling.py | 3 +- .../transformation.py | 93 ++++++++++++------- 2 files changed, 63 insertions(+), 33 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 47a6db99403..29db1df5832 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -3,6 +3,7 @@ Helper functions to handle images passed in messages """ import base64 +from typing import Optional from httpx import Response @@ -49,7 +50,7 @@ def _infer_media_type_from_url(url: str) -> str: return media_type -def _get_valid_media_type(content_type: str | None, url: str) -> str: +def _get_valid_media_type(content_type: Optional[str], url: str) -> str: """ Get a valid media type from Content-Type header, falling back to URL extension. diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 34167a3b80f..8ad9841ac14 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -1,3 +1,4 @@ +import asyncio from typing import Any, Dict, List, Optional, Tuple from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -208,9 +209,47 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert {"type": "image", "source": {"type": "url", "url": "https://..."}} to: {"type": "image", "source": {"type": "base64", "media_type": "...", "data": "..."}} + + Uses asyncio.gather to download multiple images concurrently. """ + # First pass: collect all image URLs and their positions + # Position is (message_idx, block_idx, url, original_block) + image_tasks: List[Tuple[int, int, str, Dict]] = [] + + for msg_idx, message in enumerate(messages): + if not isinstance(message, dict): + continue + content = message.get("content") + if not isinstance(content, list): + continue + for block_idx, block in enumerate(content): + if not isinstance(block, dict): + continue + if block.get("type") == "image": + source = block.get("source", {}) + if isinstance(source, dict) and source.get("type") == "url": + url = source.get("url") + if url: + image_tasks.append((msg_idx, block_idx, url, block)) + + # Download all images concurrently + if image_tasks: + download_coros = [ + async_convert_url_to_base64(task[2]) for task in image_tasks + ] + data_uris = await asyncio.gather(*download_coros) + + # Build a lookup: (msg_idx, block_idx) -> (data_uri, original_block) + image_results: Dict[Tuple[int, int], Tuple[str, Dict]] = {} + for i, task in enumerate(image_tasks): + msg_idx, block_idx, _, original_block = task + image_results[(msg_idx, block_idx)] = (data_uris[i], original_block) + else: + image_results = {} + + # Second pass: build the result converted_messages = [] - for message in messages: + for msg_idx, message in enumerate(messages): if not isinstance(message, dict): converted_messages.append(message) continue @@ -221,38 +260,28 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert continue new_content = [] - for block in content: - if not isinstance(block, dict): + for block_idx, block in enumerate(content): + key = (msg_idx, block_idx) + if key in image_results: + data_uri, original_block = image_results[key] + # Parse the data URI to extract media_type and data + # Format: "data:image/jpeg;base64,/9j/..." + media_type_part, base64_data = data_uri.split("data:")[1].split( + ";base64," + ) + # Preserve all original block fields (e.g., cache_control) + # while replacing the source + new_block = { + **original_block, + "source": { + "type": "base64", + "media_type": media_type_part, + "data": base64_data, + }, + } + new_content.append(new_block) + else: new_content.append(block) - continue - - # Check if this is an image block with URL source - if block.get("type") == "image": - source = block.get("source", {}) - if isinstance(source, dict) and source.get("type") == "url": - url = source.get("url") - if url: - # Convert URL to base64 using async utility - data_uri = await async_convert_url_to_base64(url) - # Parse the data URI to extract media_type and data - # Format: "data:image/jpeg;base64,/9j/..." - media_type_part, base64_data = data_uri.split( - "data:" - )[1].split(";base64,") - # Preserve all original block fields (e.g., cache_control) - # while replacing the source - new_block = { - **block, - "source": { - "type": "base64", - "media_type": media_type_part, - "data": base64_data, - }, - } - new_content.append(new_block) - continue - - new_content.append(block) new_message = {**message, "content": new_content} converted_messages.append(new_message)