fix: Fix returned Anthropic reasoning blocks

integrations/langchain/src/databricks_langchain/chat_models.py

@@ -547,7 +547,7 @@
                 content_blocks.append(item_dict)
 
                 try:
                     args = json.loads(item.arguments, strict=False)  # ty:ignore[unresolved-attribute, invalid-argument-type]: astral-sh/ty#1479 should fix this
                     error = None
                 except json.JSONDecodeError as e:
                     error = str(e)
@@ -575,8 +575,8 @@
                 content_blocks.append(
                     {
                         "role": "tool",
                         "content": item.output,  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
                         "tool_call_id": item.call_id,  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
                     }
                 )
             elif item.type in (
@@ -902,6 +902,7 @@
                 usage_chunk_emitted = True
         else:
             first_chunk_role = None
+            text_block_index = None
             stream: Stream[ChatCompletionChunk] = self.client.chat.completions.create(**data)
             for chunk in stream:
                 # Handle ChatAgent chunks that don't have choices but have delta
@@ -913,8 +914,12 @@
                     }
                     if hasattr(chunk, "custom_outputs"):
                         chunk_delta_dict["custom_outputs"] = chunk.custom_outputs
+                    if isinstance(content := chunk_delta_dict.get("content"), list):
+                        text_block_index = len(content)
                     chunk_message = _convert_dict_to_message_chunk(
-                        chunk_delta_dict, first_chunk_role
+                        chunk_delta_dict,
+                        first_chunk_role,
+                        text_block_index=text_block_index,
                     )
                     generation_chunk = ChatGenerationChunk(message=chunk_message)
                     if run_manager:
@@ -936,8 +941,13 @@
                         final_usage = usage  # store for usage chunk at end
                     # Use model_dump instead of manual dict reconstruction
                     chunk_delta_dict = chunk_delta.model_dump(exclude_unset=True)
+                    if isinstance(content := chunk_delta_dict.get("content"), list):
+                        text_block_index = len(content)
                     chunk_message = _convert_dict_to_message_chunk(
-                        chunk_delta_dict, first_chunk_role, usage=usage
+                        chunk_delta_dict,
+                        first_chunk_role,
+                        usage=usage,
+                        text_block_index=text_block_index,
                     )
                     generation_info = {}
                     if choice.finish_reason:
@@ -1010,6 +1020,7 @@
                 usage_chunk_emitted = True
         else:
             first_chunk_role = None
+            text_block_index = None
             stream = cast(
                 AsyncStream[ChatCompletionChunk],
                 await self.async_client.chat.completions.create(**data),
@@ -1024,8 +1035,12 @@
                     }
                     if hasattr(chunk, "custom_outputs"):
                         chunk_delta_dict["custom_outputs"] = chunk.custom_outputs
+                    if isinstance(content := chunk_delta_dict.get("content"), list):
+                        text_block_index = len(content)
                     chunk_message = _convert_dict_to_message_chunk(
-                        chunk_delta_dict, first_chunk_role
+                        chunk_delta_dict,
+                        first_chunk_role,
+                        text_block_index=text_block_index,
                     )
                     generation_chunk = ChatGenerationChunk(message=chunk_message)
                     if run_manager:
@@ -1047,8 +1062,13 @@
                         final_usage = usage  # store for usage chunk at end
                     # Use model_dump instead of manual dict reconstruction
                     chunk_delta_dict = chunk_delta.model_dump(exclude_unset=True)
+                    if isinstance(content := chunk_delta_dict.get("content"), list):
+                        text_block_index = len(content)
                     chunk_message = _convert_dict_to_message_chunk(
-                        chunk_delta_dict, first_chunk_role, usage=usage
+                        chunk_delta_dict,
+                        first_chunk_role,
+                        usage=usage,
+                        text_block_index=text_block_index,
                     )
                     generation_info = {}
                     if choice.finish_reason:
@@ -1608,9 +1628,9 @@
 ) -> HumanMessage | SystemMessage | ToolMessage | AIMessage | ChatMessage:
     role = _dict["role"]
     content = _dict.get("content") or ""
-    if not isinstance(content, str):
-        # for non-string content, serialize it into a string to maintain compatibility with downstream consumers
-        # for example, output parsers expect a string
+    if not isinstance(content, (str, list)):
+        # Preserve structured content blocks while retaining a safe fallback for
+        # unexpected provider content types.
         content = json.dumps(content)
 
     lc_message = None
@@ -1656,13 +1676,17 @@
     _dict: Mapping[str, Any],
     default_role: str | None,
     usage: CompletionUsage | dict[str, Any] | None = None,
+    text_block_index: int | None = None,
 ) -> BaseMessageChunk:
     role = _dict.get("role", default_role)
-    content = _dict.get("content") or ""
-    if not isinstance(content, str):
-        # for non-string content, serialize it into a string to maintain compatibility with downstream consumers
-        # for example, output parsers expect a string
+    content: Any = _dict.get("content") or ""
+    if text_block_index is not None and isinstance(content, str) and content:
+        content = [{"type": "text", "text": content, "index": text_block_index}]
+    elif not isinstance(content, (str, list)):
+        # Preserve structured content blocks while retaining a safe fallback for
+        # unexpected provider content types.
         content = json.dumps(content)
+    content = cast(str | list[str | dict[Any, Any]], content)
 
     lc_chunk = None
     if role == "user":
@@ -1733,15 +1757,15 @@
         item = chunk.item  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
         if item.type == "function_call_output":
             lc_chunk = ToolMessageChunk(
                 content=item.output,  # ty:ignore[unresolved-attribute, invalid-argument-type]: astral-sh/ty#1479 should fix this
                 tool_call_id=item.call_id,  # ty:ignore[unresolved-attribute, invalid-argument-type]: astral-sh/ty#1479 should fix this
             )
         elif item.type == "function_call":
             id = item.call_id  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
             tool_call_chunks.append(
                 tool_call_chunk(
                     name=item.name,  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
                     args=item.arguments,  # ty:ignore[unresolved-attribute, invalid-argument-type]: astral-sh/ty#1479 should fix this
                     id=item.call_id,  # ty:ignore[unresolved-attribute]: astral-sh/ty#1479 should fix this
                 )
             )

integrations/langchain/tests/unit_tests/test_chat_models.py

@@ -738,8 +738,8 @@
         id=ID,
         tool_calls=[
             {
                 "name": tool_calls[0]["function"]["name"],  # type: ignore[index]
                 "args": json.loads(tool_calls[0]["function"]["arguments"]),  # type: ignore[index]
                 "id": ID,
                 "type": "tool_call",
             }
@@ -1440,21 +1440,55 @@
     assert result["param2"] == "value2"
 
 
-def test_convert_dict_to_message_with_non_string_content():
-    """Test _convert_dict_to_message handles non-string content by JSON encoding it."""
-    # Test with list of dict content (matching gpt oss)
-    message_dict = {
-        "role": "assistant",
-        "content": [
-            {"type": "reasoning", "summary": [{"type": "summary_text", "text": "asdf"}]},
-            {"type": "text", "text": "asdf"},
-        ],
-    }
-    result = _convert_dict_to_message(message_dict, None)
-    expected = AIMessage(
-        content='[{"type": "reasoning", "summary": [{"type": "summary_text", "text": "asdf"}]}, {"type": "text", "text": "asdf"}]'
+def test_convert_dict_to_message_preserves_structured_content():
+    content = [
+        {
+            "type": "reasoning",
+            "summary": [{"type": "summary_text", "text": "", "signature": "abc"}],
+        },
+        {"type": "text", "text": "The answer"},
+    ]
+
+    result = _convert_dict_to_message({"role": "assistant", "content": content}, None)
+
+    assert result.content == content
+    assert result.text == "The answer"
+
+
+def test_convert_dict_to_message_chunk_preserves_structured_content():
+    reasoning = [
+        {
+            "type": "reasoning",
+            "summary": [{"type": "summary_text", "text": "", "signature": "abc"}],
+        }
+    ]
+
+    reasoning_chunk = _convert_dict_to_message_chunk(
+        {"role": "assistant", "content": reasoning}, None
     )
-    assert result == expected
+    first_text_chunk = _convert_dict_to_message_chunk(
+        {"content": "The answer"}, "assistant", text_block_index=len(reasoning)
+    )
+    second_text_chunk = _convert_dict_to_message_chunk(
+        {"content": " has spaces"}, "assistant", text_block_index=len(reasoning)
+    )
+    result = reasoning_chunk + first_text_chunk + second_text_chunk
+
+    assert result.content == [
+        *reasoning,
+        {"type": "text", "text": "The answer has spaces", "index": 1},
+    ]
+    assert result.text == "The answer has spaces"
+
+
+def test_convert_dict_to_message_serializes_unsupported_content():
+    content = {"unexpected": "value"}
+
+    message = _convert_dict_to_message({"role": "assistant", "content": content}, None)
+    chunk = _convert_dict_to_message_chunk({"role": "assistant", "content": content}, None)
+
+    assert message.content == json.dumps(content)
+    assert chunk.content == json.dumps(content)
 
 
 ### Test custom_inputs and custom_outputs functionality ###
@@ -1778,7 +1812,7 @@
         input_tokens=100,
         output_tokens=50,
         total_tokens=150,
         input_tokens_details=InputTokensDetails(cached_tokens=cached_tokens),
         output_tokens_details=OutputTokensDetails(reasoning_tokens=reasoning_tokens),
     )
 
@@ -1856,7 +1890,7 @@
         input_tokens=100,
         output_tokens=50,
         total_tokens=150,
         input_tokens_details=InputTokensDetails(cached_tokens=25),
         output_tokens_details=OutputTokensDetails(reasoning_tokens=10),
     )
     response = Mock()
@@ -1906,7 +1940,7 @@
             input_tokens=100,
             output_tokens=50,
             total_tokens=150,
             input_tokens_details=InputTokensDetails(cached_tokens=25),
             output_tokens_details=OutputTokensDetails(reasoning_tokens=10),
         )
     else:
@@ -2083,7 +2117,7 @@
                 input_tokens=100,
                 output_tokens=50,
                 total_tokens=150,
                 input_tokens_details=InputTokensDetails(cached_tokens=25),
                 output_tokens_details=OutputTokensDetails(reasoning_tokens=10),
             ),
         )