fix(vertexai): mime type parsing from image URLs with params (#1377)

Author: lexi-k

libs/vertexai/langchain_google_vertexai/_anthropic_utils.py

@@ -1,5 +1,6 @@
 import base64
 import re
+import urllib
 import warnings
 from collections.abc import Callable, Sequence
 from typing import (
@@ -90,7 +91,8 @@ def _format_image(image_url: str, project: str | None) -> dict:
     if validators.url(image_url):
         loader = ImageBytesLoader(project=project)
         image_bytes = loader.load_bytes(image_url)
-        raw_mime_type = image_url.split(".")[-1].lower()
+        path = urllib.parse.urlparse(image_url).path
+        raw_mime_type = path.split(".")[-1].lower()
         doc_type = "application" if raw_mime_type == "pdf" else "image"
         mime_type = (
             f"{doc_type}/jpeg"

libs/vertexai/tests/unit_tests/test_anthropic_utils.py

@@ -1,5 +1,7 @@
 """Unit tests for _anthropic_utils.py."""
 
+import base64
+from unittest.mock import patch
 import pytest
 from anthropic.types import (
     RawContentBlockDeltaEvent,
@@ -23,6 +25,8 @@
     _thinking_in_params,
 )
 
+from langchain_google_vertexai._anthropic_utils import _format_image
+
 
 def test_format_message_anthropic_with_cache_control_in_kwargs() -> None:
     """Test formatting a message with cache control in additional_kwargs."""
@@ -1309,3 +1313,34 @@ def test_tool_message_preserves_cache_control() -> None:
         "tool_use_id": "call_1",
         "cache_control": {"type": "ephemeral"},
     }
+
+
+@pytest.mark.parametrize(
+    ("image_url", "expected_media_type"),
+    [
+        ("https://example.com/image.png?token=123", "image/png"),
+        ("https://example.com/image.jpg", "image/jpeg"),
+        ("https://example.com/document.pdf", "application/pdf"),
+    ],
+)
+def test_format_image(image_url: str, expected_media_type: str) -> None:
+    """Test that _format_image correctly handles various URLs."""
+    project = "test-project"
+
+    with patch(
+        "langchain_google_vertexai._anthropic_utils.ImageBytesLoader"
+    ) as MockLoader:
+        mock_loader_instance = MockLoader.return_value
+        mock_loader_instance.load_bytes.return_value = b"fake_image_data"
+
+        result = _format_image(image_url, project)
+
+        expected_data = base64.b64encode(b"fake_image_data").decode("ascii")
+
+        assert result == {
+            "type": "base64",
+            "media_type": expected_media_type,
+            "data": expected_data,
+        }
+
+        mock_loader_instance.load_bytes.assert_called_once_with(image_url)