[Bug fix] Vertex Multimodal embeddings - Support text + base64 image combinations (#18172)

* TestVertexMultimodalEmbedding

* fix _try_merge_text_with_media

* ruff fix

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

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
 

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"
 

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 (

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}"