From 17403967aba7edf07d6f6eadc42450c413e274e9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 18 Dec 2025 13:22:43 +0530 Subject: [PATCH] [Bug fix] Vertex Multimodal embeddings - Support text + base64 image combinations (#18172) * TestVertexMultimodalEmbedding * fix _try_merge_text_with_media * ruff fix --- .../chat/guardrail_translation/handler.py | 1 - .../multimodal_embeddings/transformation.py | 122 +++++++++++------- .../guardrail_hooks/grayswan/grayswan.py | 13 +- .../tag_management_endpoints.py | 1 - ..._ai_multimodal_embedding_transformation.py | 59 +++++++++ 5 files changed, 141 insertions(+), 55 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 6f53bf65714..9d50cc4d92d 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -43,7 +43,6 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, - AnthropicResponseTextBlock, ) diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py index 5bf02ad765f..2cb2ac9ed8f 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py @@ -58,36 +58,81 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): headers.update(default_headers) return headers + def _is_gcs_uri(self, input_str: str) -> bool: + """Check if the input string is a GCS URI.""" + return "gs://" in input_str + + def _is_video(self, input_str: str) -> bool: + """Check if the input string represents a video (mp4).""" + return "mp4" in input_str + + def _is_media_input(self, input_str: str) -> bool: + """Check if the input string is a media element (GCS URI or base64 image).""" + return self._is_gcs_uri(input_str) or is_base64_encoded(s=input_str) + + def _create_image_instance(self, input_str: str) -> InstanceImage: + """Create an InstanceImage from a GCS URI or base64 string.""" + if self._is_gcs_uri(input_str): + return InstanceImage(gcsUri=input_str) + else: + return InstanceImage( + bytesBase64Encoded=( + input_str.split(",")[1] if "," in input_str else input_str + ) + ) + + def _create_video_instance(self, input_str: str) -> InstanceVideo: + """Create an InstanceVideo from a GCS URI.""" + return InstanceVideo(gcsUri=input_str) + def _process_input_element(self, input_element: str) -> Instance: """ - Process the input element for multimodal embedding requests. checks if the if the input is gcs uri, base64 encoded image or plain text. + Process a single input element for multimodal embedding requests. + Detects if the input is a GCS URI, base64 encoded image, or plain text. Args: input_element (str): The input element to process. Returns: - Dict[str, Any]: A dictionary representing the processed input element. + Instance: A dictionary representing the processed input element. """ if len(input_element) == 0: return Instance(text=input_element) - elif "gs://" in input_element: - if "mp4" in input_element: - return Instance(video=InstanceVideo(gcsUri=input_element)) + elif self._is_gcs_uri(input_element): + if self._is_video(input_element): + return Instance(video=self._create_video_instance(input_element)) else: - return Instance(image=InstanceImage(gcsUri=input_element)) + return Instance(image=self._create_image_instance(input_element)) elif is_base64_encoded(s=input_element): - return Instance( - image=InstanceImage( - bytesBase64Encoded=( - input_element.split(",")[1] - if "," in input_element - else input_element - ) - ) - ) + return Instance(image=self._create_image_instance(input_element)) else: return Instance(text=input_element) + def _try_merge_text_with_media( + self, text_str: str, next_elem: Optional[str] + ) -> tuple[Instance, bool]: + """ + Try to merge a text element with a following media element into a single instance. + + Args: + text_str: The text string to potentially merge. + next_elem: The next element in the input list (may be media). + + Returns: + A tuple of (Instance, consumed_next) where consumed_next indicates + if the next element was merged into this instance. + """ + instance_args: Instance = {"text": text_str} + + if next_elem and isinstance(next_elem, str) and self._is_media_input(next_elem): + if self._is_gcs_uri(next_elem) and self._is_video(next_elem): + instance_args["video"] = self._create_video_instance(next_elem) + else: + instance_args["image"] = self._create_image_instance(next_elem) + return instance_args, True + + return instance_args, False + def process_openai_embedding_input( self, _input: Union[list, str] ) -> List[Instance]: @@ -98,50 +143,33 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): _input (Union[list, str]): The input data to process. Returns: - Union[Instance, List[Instance]]: Either a single Instance or list of Instance objects. + List[Instance]: List of Instance objects for the embedding request. """ _input_list = [_input] if not isinstance(_input, list) else _input - processed_instances = [] + processed_instances: List[Instance] = [] i = 0 while i < len(_input_list): current = _input_list[i] - - # Look ahead for potential media elements next_elem = _input_list[i + 1] if i + 1 < len(_input_list) else None - # If current is a text and next is a GCS URI, or current is a GCS URI if isinstance(current, str): - instance_args: Instance = {} - - # Process current element - if "gs://" not in current: - instance_args["text"] = current - elif "mp4" in current: - instance_args["video"] = InstanceVideo(gcsUri=current) + if self._is_media_input(current): + # Current element is media - process it standalone + processed_instances.append(self._process_input_element(current)) + i += 1 else: - instance_args["image"] = InstanceImage(gcsUri=current) - - # Check next element if it's a GCS URI - if next_elem and isinstance(next_elem, str) and "gs://" in next_elem: - if "mp4" in next_elem: - instance_args["video"] = InstanceVideo(gcsUri=next_elem) - else: - instance_args["image"] = InstanceImage(gcsUri=next_elem) - i += 2 # Skip next element since we processed it - else: - i += 1 # Move to next element - - processed_instances.append(instance_args) - continue - - # Handle dict or other types - if isinstance(current, dict): - instance = Instance(**current) - processed_instances.append(instance) + # Current element is text - try to merge with next media element + instance, consumed_next = self._try_merge_text_with_media( + text_str=current, next_elem=next_elem + ) + processed_instances.append(instance) + i += 2 if consumed_next else 1 + elif isinstance(current, dict): + processed_instances.append(Instance(**current)) + i += 1 else: raise ValueError(f"Unsupported input type: {type(current)}") - i += 1 return processed_instances diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 38e75ebabbc..59e737f7d21 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -8,7 +8,6 @@ from fastapi import HTTPException from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, - ModifyResponseException, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import ( @@ -506,11 +505,13 @@ class GraySwanGuardrail(CustomGuardrail): # Handle legacy format where detection_info is a list if isinstance(detection_info, list) and len(detection_info) > 0: detection_info = detection_info[0] - - violation_score = detection_info.get("violation_score", 0.0) - violated_rules = detection_info.get("violated_rules", []) - mutation = detection_info.get("mutation", False) - ipi = detection_info.get("ipi", False) + + # Extract fields from detection_info dict + detection_dict: dict = detection_info if isinstance(detection_info, dict) else {} + violation_score = detection_dict.get("violation_score", 0.0) + violated_rules = detection_dict.get("violated_rules", []) + mutation = detection_dict.get("mutation", False) + ipi = detection_dict.get("ipi", False) violation_location = "the model response" if is_output else "input query" diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 5085fb2a5b8..95b7300992c 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -17,7 +17,6 @@ from typing import TYPE_CHECKING, Dict, List, Optional from fastapi import APIRouter, Depends, HTTPException from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_daily_activity import ( diff --git a/tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py b/tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py index 88a60cb7c0a..63677c0f5f1 100644 --- a/tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py @@ -76,3 +76,62 @@ class TestVertexMultimodalEmbedding: assert ( self.config.process_openai_embedding_input(input_data) == expected_output ), f"Expected {expected_output}, but got {self.config.process_openai_embedding_input(input_data)}" + + def test_process_text_and_base64_image_input(self): + """Test that text + base64 image combinations are correctly merged into a single instance.""" + base64_image = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+ip1sAAAAASUVORK5CYII=" + input_data = ["describe this image", base64_image] + expected_output = [ + Instance( + text="describe this image", + image=InstanceImage(bytesBase64Encoded=base64_image.split(",")[1]), + ), + ] + result = self.config.process_openai_embedding_input(input_data) + assert result == expected_output, f"Expected {expected_output}, but got {result}" + + def test_process_multiple_text_and_base64_image_pairs(self): + """Test multiple text + base64 image pairs in a single request.""" + base64_image = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+ip1sAAAAASUVORK5CYII=" + input_data = [ + "first description", + base64_image, + "second description", + base64_image, + ] + expected_output = [ + Instance( + text="first description", + image=InstanceImage(bytesBase64Encoded=base64_image.split(",")[1]), + ), + Instance( + text="second description", + image=InstanceImage(bytesBase64Encoded=base64_image.split(",")[1]), + ), + ] + result = self.config.process_openai_embedding_input(input_data) + assert result == expected_output, f"Expected {expected_output}, but got {result}" + + def test_process_base64_image_only_in_list(self): + """Test that standalone base64 images in a list are processed correctly.""" + base64_image = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+ip1sAAAAASUVORK5CYII=" + input_data = [base64_image, base64_image] + expected_output = [ + Instance(image=InstanceImage(bytesBase64Encoded=base64_image.split(",")[1])), + Instance(image=InstanceImage(bytesBase64Encoded=base64_image.split(",")[1])), + ] + result = self.config.process_openai_embedding_input(input_data) + assert result == expected_output, f"Expected {expected_output}, but got {result}" + + def test_process_text_and_gcs_image_input(self): + """Test that text + GCS image combinations are correctly merged.""" + gcs_uri = "gs://my-bucket/image.png" + input_data = ["describe this image", gcs_uri] + expected_output = [ + Instance( + text="describe this image", + image=InstanceImage(gcsUri=gcs_uri), + ), + ] + result = self.config.process_openai_embedding_input(input_data) + assert result == expected_output, f"Expected {expected_output}, but got {result}"