vertexai: Fixed content coercion while streaming Claude thinking response (#834)

Author: mitraan-deshaw

libs/vertexai/langchain_google_vertexai/_anthropic_utils.py

@@ -425,3 +425,20 @@ def _tools_in_params(params: dict) -> bool:
     return "tools" in params or (
         "extra_body" in params and params["extra_body"].get("tools")
     )
+
+
+def _thinking_in_params(params: dict) -> bool:
+    return params.get("thinking", {}).get("type") == "enabled"
+
+
+def _documents_in_params(params: dict) -> bool:
+    for message in params.get("messages", []):
+        if isinstance(message.get("content"), list):
+            for block in message["content"]:
+                if (
+                    isinstance(block, dict)
+                    and block.get("type") == "document"
+                    and block.get("citations", {}).get("enabled")
+                ):
+                    return True
+    return False

libs/vertexai/langchain_google_vertexai/model_garden.py

@@ -55,8 +55,10 @@
     _extract_tool_calls,
 )
 from langchain_google_vertexai._anthropic_utils import (
+    _documents_in_params,
     _format_messages_anthropic,
     _make_message_chunk_from_anthropic_event,
+    _thinking_in_params,
     _tools_in_params,
     convert_to_anthropic_tool,
 )
@@ -378,7 +380,11 @@ def _stream_with_retry(**params: Any) -> Any:
             return self.client.messages.create(**params, stream=True)
 
         stream = _stream_with_retry(**params)
-        coerce_content_to_string = not _tools_in_params(params)
+        coerce_content_to_string = (
+            not _tools_in_params(params)
+            and not _documents_in_params(params)
+            and not _thinking_in_params(params)
+        )
         for event in stream:
             msg = _make_message_chunk_from_anthropic_event(
                 event,
@@ -414,7 +420,11 @@ async def _astream_with_retry(**params: Any) -> Any:
             return await self.async_client.messages.create(**params, stream=True)
 
         stream = await _astream_with_retry(**params)
-        coerce_content_to_string = not _tools_in_params(params)
+        coerce_content_to_string = (
+            not _tools_in_params(params)
+            and not _documents_in_params(params)
+            and not _thinking_in_params(params)
+        )
         async for event in stream:
             msg = _make_message_chunk_from_anthropic_event(
                 event,

libs/vertexai/tests/integration_tests/test_model_garden.py

@@ -142,6 +142,30 @@ def test_anthropic_stream() -> None:
         assert isinstance(chunk, AIMessageChunk)
 
 
+@pytest.mark.extended
+def test_anthropic_thinking_stream() -> None:
+    project = os.environ["PROJECT_ID"]
+    location = "us-east5"
+    model = ChatAnthropicVertex(
+        project=project,
+        location=location,
+        model_kwargs={
+            "thinking": {
+                "type": "enabled",
+                "budget_tokens": 1024,  # budget tokens >= 1024
+            },
+        },
+        max_tokens=2048,  # max_tokens must be greater than budget_tokens
+    )
+    question = (
+        "Hello, could you recommend a good movie for me to watch this evening, please?"
+    )
+    message = HumanMessage(content=question)
+    sync_response = model.stream([message], model="claude-3-7-sonnet@20250219")
+    for chunk in sync_response:
+        assert isinstance(chunk, AIMessageChunk)
+
+
 @pytest.mark.extended
 async def test_anthropic_async() -> None:
     project = os.environ["PROJECT_ID"]

libs/vertexai/tests/unit_tests/test_anthropic_utils.py

@@ -16,9 +16,11 @@
 from langchain_core.messages.tool import tool_call as create_tool_call
 
 from langchain_google_vertexai._anthropic_utils import (
+    _documents_in_params,
     _format_message_anthropic,
     _format_messages_anthropic,
     _make_message_chunk_from_anthropic_event,
+    _thinking_in_params,
 )
 
 
@@ -838,3 +840,54 @@ def test_make_thinking_message_chunk_from_anthropic_event() -> None:
     )
     assert isinstance(thinking_chunk, AIMessageChunk)
     assert isinstance(signature_chunk, AIMessageChunk)
+
+
+def test_thinking_in_params_true() -> None:
+    """Test _thinking_in_params when thinking.type is 'enabled'."""
+    params = {"thinking": {"type": "enabled", "budget_tokens": 1024}}
+
+    assert _thinking_in_params(params)
+
+
+def test_thinking_in_params_false_different_type() -> None:
+    """Test _thinking_in_params when thinking.type is 'disabled'."""
+    params = {"thinking": {"type": "disabled", "budget_tokens": 1024}}
+
+    assert not _thinking_in_params(params)
+
+
+def test_documents_in_params_true() -> None:
+    """Test _documents_in_params when document with citations is enabled."""
+    params = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [{"type": "document", "citations": {"enabled": True}}],
+            }
+        ]
+    }
+
+    assert _documents_in_params(params)
+
+
+def test_documents_in_params_false_citations_disabled() -> None:
+    """Test _documents_in_params when citations are not enabled."""
+    params = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [{"type": "document", "citations": {"enabled": False}}],
+            }
+        ]
+    }
+
+    assert not _documents_in_params(params)
+
+
+def test_documents_in_params_false_no_document() -> None:
+    """Test _documents_in_params when there are no documents."""
+    params = {
+        "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}]
+    }
+
+    assert not _documents_in_params(params)