From 22a14789d0acd5f3a203971a4d723fc93a23a501 Mon Sep 17 00:00:00 2001 From: jacoblee93 Date: Tue, 21 Jul 2026 13:56:02 -0700 Subject: [PATCH] Fix returned Anthropic reasoning blocks --- .../src/databricks_langchain/chat_models.py | 46 ++++++++++---- .../tests/unit_tests/test_chat_models.py | 62 ++++++++++++++----- 2 files changed, 83 insertions(+), 25 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/chat_models.py b/integrations/langchain/src/databricks_langchain/chat_models.py index 0bfa87f2..643a5f4a 100644 --- a/integrations/langchain/src/databricks_langchain/chat_models.py +++ b/integrations/langchain/src/databricks_langchain/chat_models.py @@ -902,6 +902,7 @@ def _stream( 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 @@ def _stream( } 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 @@ def _stream( 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 @@ async def _astream( 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 @@ async def _astream( } 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 @@ async def _astream( 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 @@ def _convert_dict_to_message( ) -> 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 @@ def _convert_dict_to_message_chunk( _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": diff --git a/integrations/langchain/tests/unit_tests/test_chat_models.py b/integrations/langchain/tests/unit_tests/test_chat_models.py index 5e85ff7e..fe44d7ad 100644 --- a/integrations/langchain/tests/unit_tests/test_chat_models.py +++ b/integrations/langchain/tests/unit_tests/test_chat_models.py @@ -1440,21 +1440,55 @@ def test_prepare_inputs_with_extra_params(): 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 ###