fixed Anthropic url support (#839)

Author: lkuligin

libs/vertexai/langchain_google_vertexai/_anthropic_utils.py

@@ -1,3 +1,4 @@
+import base64
 import re
 import warnings
 from typing import (
@@ -31,7 +32,9 @@
 from langchain_core.utils.function_calling import convert_to_openai_tool
 from pydantic import BaseModel
 
-from langchain_google_vertexai._image_utils import image_bytes_to_b64_string
+from langchain_google_vertexai._image_utils import (
+    ImageBytesLoader,
+)
 from langchain_google_vertexai._utils import load_image_from_gcs
 
 if TYPE_CHECKING:
@@ -49,29 +52,43 @@
 
 def _format_image(image_url: str, project: Optional[str]) -> Dict:
     """Formats a message image to a dict for anthropic api."""
-    regex = r"^data:(?P<media_type>image/.+);base64,(?P<data>.+)$"
+    regex = r"^data:(?P<media_type>(?:image|application)/.+);base64,(?P<data>.+)$"
     match = re.match(regex, image_url)
-
     if match:
         return {
             "type": "base64",
             "media_type": match.group("media_type"),
             "data": match.group("data"),
         }
     elif validators.url(image_url):
+        loader = ImageBytesLoader(project=project)
+        image_bytes = loader.load_bytes(image_url)
+        raw_mime_type = image_url.split(".")[-1].lower()
+        doc_type = "application" if raw_mime_type == "pdf" else "image"
+        mime_type = (
+            f"{doc_type}/jpeg"
+            if raw_mime_type == "jpg"
+            else f"{doc_type}/{raw_mime_type}"
+        )
         return {
-            "type": "url",
-            "url": image_url,
+            "type": "base64",
+            "media_type": mime_type,
+            "data": base64.b64encode(image_bytes).decode("ascii"),
         }
     elif image_url.startswith("gs://"):
         # Gets image and encodes to base64.
-        image = load_image_from_gcs(image_url, project)
+        loader = ImageBytesLoader(project=project)
+        part = loader.load_part(image_url)
+        if part.file_data.mime_type:
+            mime_type = part.file_data.mime_type
+            image_data = load_image_from_gcs(image_url, project=project).data
+        else:
+            mime_type = part.inline_data.mime_type
+            image_data = part.inline_data.data
         return {
             "type": "base64",
-            "media_type": image._mime_type(),
-            "data": image_bytes_to_b64_string(
-                image.data(), "ascii", image._mime_type().split("/")[-1]
-            ),
+            "media_type": mime_type,
+            "data": base64.b64encode(image_data).decode("ascii"),
         }
     else:
         raise ValueError(
@@ -168,7 +185,11 @@ def _format_message_anthropic(
                 if block["type"] == "image_url":
                     # convert format
                     source = _format_image(block["image_url"]["url"], project)
-                    content.append({"type": "image", "source": source})
+                    if source["media_type"] == "application/pdf":
+                        doc_type = "document"
+                    else:
+                        doc_type = "image"
+                    content.append({"type": doc_type, "source": source})
                     continue
 
                 if block["type"] == "tool_use":

libs/vertexai/langchain_google_vertexai/_image_utils.py

@@ -252,8 +252,12 @@ def image_bytes_to_b64_string(
     Returns:
         B64 image encoded string.
     """
+    if image_format == "pdf":
+        image_type = "application"
+    else:
+        image_type = "image"
     encoded_bytes = base64.b64encode(image_bytes).decode(encoding)
-    return f"data:image/{image_format};base64,{encoded_bytes}"
+    return f"data:{image_type}/{image_format};base64,{encoded_bytes}"
 
 
 def create_text_content_part(message_str: str) -> Dict:

libs/vertexai/tests/integration_tests/test_anthropic_files.py

@@ -0,0 +1,78 @@
+"""Test ChatGoogleVertexAI chat model."""
+import os
+
+from langchain_google_vertexai._image_utils import image_bytes_to_b64_string
+from langchain_google_vertexai._utils import load_image_from_gcs
+from langchain_google_vertexai.model_garden import ChatAnthropicVertex
+
+
+def test_pdf_gcs_uri():
+    gcs_uri = "gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf"
+    llm = ChatAnthropicVertex(
+        model="claude-3-5-sonnet-v2@20241022",
+        location="us-east5",
+        temperature=0.8,
+        project=os.environ["PROJECT_ID"],
+    )
+
+    res = llm.invoke(
+        [
+            {
+                "role": "user",
+                "content": [
+                    "Parse this pdf.",
+                    {"type": "image_url", "image_url": {"url": gcs_uri}},
+                ],
+            }
+        ]
+    )
+    assert len(res.content) > 100
+
+
+def test_pdf_byts():
+    gcs_uri = "gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf"
+    llm = ChatAnthropicVertex(
+        model="claude-3-5-sonnet-v2@20241022",
+        location="us-east5",
+        temperature=0.8,
+        project=os.environ["PROJECT_ID"],
+    )
+    image = load_image_from_gcs(gcs_uri, "kuligin-sandbox1")
+    image_data = image_bytes_to_b64_string(image.data, "ascii", "pdf")
+
+    res = llm.invoke(
+        [
+            {
+                "role": "user",
+                "content": [
+                    "Parse this pdf.",
+                    {"type": "image_url", "image_url": {"url": image_data}},
+                ],
+            }
+        ]
+    )
+    assert len(res.content) > 100
+
+
+def test_https_image():
+    uri = "https://picsum.photos/seed/picsum/200/300.jpg"
+
+    llm = ChatAnthropicVertex(
+        model="claude-3-5-sonnet-v2@20241022",
+        location="us-east5",
+        temperature=0.8,
+        project=os.environ["PROJECT_ID"],
+    )
+
+    res = llm.invoke(
+        [
+            {
+                "role": "user",
+                "content": [
+                    "Parse this pdf.",
+                    {"type": "image_url", "image_url": {"url": uri}},
+                ],
+            }
+        ]
+    )
+    assert len(res.content) > 10

libs/vertexai/tests/unit_tests/test_anthropic_utils.py

@@ -173,40 +173,6 @@ def test_format_message_anthropic_with_chain_of_thoughts():
     ]
 
 
-def test_format_message_anthropic_with_image_content():
-    """Test formatting a system message with chain of thoughts."""
-    message = SystemMessage(
-        content=[
-            {
-                "type": "image_url",
-                "image_url": {"url": "data:image/png;base64,/9j/4AAQSk"},
-            },
-            {
-                "type": "image_url",
-                "image_url": {"url": "https://your-valid-image-url.png"},
-            },
-        ]
-    )
-    result = _format_message_anthropic(message, project="test-project")
-    assert result == [
-        {
-            "type": "image",
-            "source": {
-                "type": "base64",
-                "media_type": "image/png",
-                "data": "/9j/4AAQSk",
-            },
-        },
-        {
-            "type": "image",
-            "source": {
-                "type": "url",
-                "url": "https://your-valid-image-url.png",
-            },
-        },
-    ]
-
-
 def test_format_messages_anthropic_with_system_string():
     """Test formatting messages with system message as string."""
     messages = [