Add GCS image support in ChatAnthropicVertex model (#808)

Author: SauravP97

libs/vertexai/langchain_google_vertexai/_anthropic_utils.py

@@ -31,6 +31,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._utils import load_image_from_gcs
+
 if TYPE_CHECKING:
     from anthropic.types import (
         RawMessageStreamEvent,  # type: ignore[unused-ignore, import-not-found]
@@ -44,7 +47,7 @@
 }
 
 
-def _format_image(image_url: str) -> Dict:
+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>.+)$"
     match = re.match(regex, image_url)
@@ -60,6 +63,16 @@ def _format_image(image_url: str) -> Dict:
             "type": "url",
             "url": image_url,
         }
+    elif image_url.startswith("gs://"):
+        # Gets image and encodes to base64.
+        image = load_image_from_gcs(image_url, project)
+        return {
+            "type": "base64",
+            "media_type": image._mime_type(),
+            "data": image_bytes_to_b64_string(
+                image.data(), "ascii", image._mime_type().split("/")[-1]
+            ),
+        }
     else:
         raise ValueError(
             "Anthropic only supports base64-encoded images and urls currently."
@@ -83,7 +96,9 @@ def _format_text_content(text: str) -> Dict[str, Union[str, Dict[str, Any]]]:
     return content
 
 
-def _format_message_anthropic(message: Union[HumanMessage, AIMessage, SystemMessage]):
+def _format_message_anthropic(
+    message: Union[HumanMessage, AIMessage, SystemMessage], project: Optional[str]
+):
     """Format a message for Anthropic API.
 
     Args:
@@ -152,7 +167,7 @@ def _format_message_anthropic(message: Union[HumanMessage, AIMessage, SystemMess
 
                 if block["type"] == "image_url":
                     # convert format
-                    source = _format_image(block["image_url"]["url"])
+                    source = _format_image(block["image_url"]["url"], project)
                     content.append({"type": "image", "source": source})
                     continue
 
@@ -183,6 +198,7 @@ def _format_message_anthropic(message: Union[HumanMessage, AIMessage, SystemMess
 
 def _format_messages_anthropic(
     messages: List[BaseMessage],
+    project: Optional[str],
 ) -> Tuple[Optional[Dict[str, Any]], List[Dict]]:
     """Formats messages for anthropic."""
     system_messages: Optional[Dict[str, Any]] = None
@@ -193,12 +209,12 @@ def _format_messages_anthropic(
         if message.type == "system":
             if i != 0:
                 raise ValueError("System message must be at beginning of message list.")
-            fm = _format_message_anthropic(message)
+            fm = _format_message_anthropic(message, project)
             if fm:
                 system_messages = fm
             continue
 
-        fm = _format_message_anthropic(message)
+        fm = _format_message_anthropic(message, project)
         if not fm:
             continue
         formatted_messages.append(fm)

libs/vertexai/langchain_google_vertexai/model_garden.py

@@ -248,7 +248,9 @@ def _format_params(
         stop: Optional[List[str]] = None,
         **kwargs: Any,
     ) -> Dict[str, Any]:
-        system_message, formatted_messages = _format_messages_anthropic(messages)
+        system_message, formatted_messages = _format_messages_anthropic(
+            messages, self.project
+        )
         params = self._default_params
         params.update(kwargs)
         if kwargs.get("model_name"):

libs/vertexai/tests/unit_tests/test_anthropic_utils.py

@@ -27,7 +27,7 @@ def test_format_message_anthropic_with_cache_control_in_kwargs():
     message = HumanMessage(
         content="Hello", additional_kwargs={"cache_control": {"type": "semantic"}}
     )
-    result = _format_message_anthropic(message)
+    result = _format_message_anthropic(message, project="test-project")
     assert result == {
         "role": "user",
         "content": [
@@ -43,7 +43,7 @@ def test_format_message_anthropic_with_cache_control_in_block():
             {"type": "text", "text": "Hello", "cache_control": {"type": "semantic"}}
         ]
     )
-    result = _format_message_anthropic(message)
+    result = _format_message_anthropic(message, project="test-project")
     assert result == {
         "role": "user",
         "content": [
@@ -61,7 +61,7 @@ def test_format_message_anthropic_with_mixed_blocks():
             "Plain text",
         ]
     )
-    result = _format_message_anthropic(message)
+    result = _format_message_anthropic(message, project="test-project")
     assert result == {
         "role": "user",
         "content": [
@@ -81,7 +81,9 @@ def test_format_messages_anthropic_with_system_cache_control():
         ),
         HumanMessage(content="Hello"),
     ]
-    system_messages, formatted_messages = _format_messages_anthropic(messages)
+    system_messages, formatted_messages = _format_messages_anthropic(
+        messages, project="test-project"
+    )
 
     assert system_messages == [
         {
@@ -102,7 +104,7 @@ def test_format_message_anthropic_system():
         content="System message",
         additional_kwargs={"cache_control": {"type": "ephemeral"}},
     )
-    result = _format_message_anthropic(message)
+    result = _format_message_anthropic(message, project="test-project")
     assert result == [
         {
             "type": "text",
@@ -124,7 +126,7 @@ def test_format_message_anthropic_system_list():
             {"type": "text", "text": "System rule 2"},
         ]
     )
-    result = _format_message_anthropic(message)
+    result = _format_message_anthropic(message, project="test-project")
     assert result == [
         {
             "type": "text",
@@ -156,7 +158,7 @@ def test_format_message_anthropic_with_chain_of_thoughts():
             },
         ]
     )
-    result = _format_message_anthropic(message)
+    result = _format_message_anthropic(message, project="test-project")
     assert result == [
         {
             "type": "text",
@@ -185,7 +187,7 @@ def test_format_message_anthropic_with_image_content():
             },
         ]
     )
-    result = _format_message_anthropic(message)
+    result = _format_message_anthropic(message, project="test-project")
     assert result == [
         {
             "type": "image",
@@ -211,7 +213,9 @@ def test_format_messages_anthropic_with_system_string():
         SystemMessage(content="System message"),
         HumanMessage(content="Hello"),
     ]
-    system_messages, formatted_messages = _format_messages_anthropic(messages)
+    system_messages, formatted_messages = _format_messages_anthropic(
+        messages, project="test-project"
+    )
 
     assert system_messages == [{"type": "text", "text": "System message"}]
 
@@ -235,7 +239,9 @@ def test_format_messages_anthropic_with_system_list():
         ),
         HumanMessage(content="Hello"),
     ]
-    system_messages, formatted_messages = _format_messages_anthropic(messages)
+    system_messages, formatted_messages = _format_messages_anthropic(
+        messages, project="test-project"
+    )
 
     assert system_messages == [
         {
@@ -266,7 +272,9 @@ def test_format_messages_anthropic_with_system_mixed_list():
         ),
         HumanMessage(content="Hello"),
     ]
-    system_messages, formatted_messages = _format_messages_anthropic(messages)
+    system_messages, formatted_messages = _format_messages_anthropic(
+        messages, project="test-project"
+    )
 
     assert system_messages == [
         {"type": "text", "text": "Plain system rule"},
@@ -308,7 +316,9 @@ def test_format_messages_anthropic_with_mixed_messages():
             additional_kwargs={"cache_control": {"type": "semantic"}},
         ),
     ]
-    system_messages, formatted_messages = _format_messages_anthropic(messages)
+    system_messages, formatted_messages = _format_messages_anthropic(
+        messages, project="test-project"
+    )
 
     assert system_messages == [
         {
@@ -806,7 +816,9 @@ def test_format_messages_anthropic(
     source_history, expected_sm, expected_history
 ) -> None:
     """Test the original format_messages_anthropic functionality."""
-    sm, result_history = _format_messages_anthropic(source_history)
+    sm, result_history = _format_messages_anthropic(
+        source_history, project="test-project"
+    )
 
     for result, expected in zip(result_history, expected_history):
         assert result == expected