From 6c232f3b85b8a0e7cb79ce116d6f6c014809632e Mon Sep 17 00:00:00 2001 From: pokepoke81 <4258646+pokepoke81@users.noreply.github.com> Date: Fri, 14 Aug 2026 11:31:54 -0400 Subject: [PATCH] fix(databricks): keep a single reasoning block and build the message immutably A message arriving with a reasoning block already in its content and signed thinking_blocks alongside it gained a second reasoning block, which Databricks rejects. Keep the existing block and drop the redundant keys. Also builds the transformed message in one shot rather than copying and mutating, per the repo's immutability convention for new code. --- .../llms/databricks/chat/transformation.py | 60 ++++++++++--------- .../test_databricks_chat_transformation.py | 20 +++++++ 2 files changed, 53 insertions(+), 27 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 1355c0a7bd6..ee4993bcbf3 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -481,38 +481,44 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): Databricks rejects both message-level keys outright, and constrains the replacement to exactly one reasoning block holding exactly one summary entry, whose signature is - required. Blocks without a signature cannot be replayed, so the keys are dropped rather - than sent in a form the API refuses. reasoning_content is a plain-text mirror of the - same thinking, so it carries nothing the signed summary does not. + required. Blocks without a signature cannot be replayed, and a message already carrying + a reasoning block must not gain a second one, so in both cases the keys are dropped + rather than sent in a form the API refuses. reasoning_content is a plain-text mirror of + the same thinking, so it carries nothing the signed summary does not. """ - transformed_message: Final = cast(dict[str, Any], message.copy()) - thinking_blocks: Final = transformed_message.pop("thinking_blocks", None) - transformed_message.pop("reasoning_content", None) - signed_block: Final = next( - (block for block in thinking_blocks or [] if isinstance(block, dict) and block.get("signature")), - None, + stripped: Final = cast( + dict[str, Any], + {key: value for key, value in message.items() if key not in ("thinking_blocks", "reasoning_content")}, ) - if signed_block is None: - return cast(AllMessageValues, transformed_message) - - content: Final = transformed_message.get("content") + content: Final = stripped.get("content") existing_blocks: Final = ( [{"type": "text", "text": content}] if isinstance(content, str) and content else list(content or []) ) - transformed_message["content"] = [ - { - "type": "reasoning", - "summary": [ - { - "type": "summary_text", - "text": signed_block.get("thinking") or "", - "signature": signed_block["signature"], - } - ], - }, - *existing_blocks, - ] - return cast(AllMessageValues, transformed_message) + if any(isinstance(block, dict) and block.get("type") == "reasoning" for block in existing_blocks): + return cast(AllMessageValues, stripped) + + signed_block: Final = next( + ( + block + for block in message.get("thinking_blocks") or [] + if isinstance(block, dict) and block.get("signature") + ), + None, + ) + if signed_block is None: + return cast(AllMessageValues, stripped) + + reasoning_block: Final = { + "type": "reasoning", + "summary": [ + { + "type": "summary_text", + "text": signed_block.get("thinking") or "", + "signature": signed_block["signature"], + } + ], + } + return cast(AllMessageValues, {**stripped, "content": [reasoning_block, *existing_blocks]}) @staticmethod def extract_content_str( diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index e4beb3bbfc4..6800c61d4e1 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -606,3 +606,23 @@ def test_transform_messages_strips_both_reasoning_fields(): assert all("thinking_blocks" not in m and "reasoning_content" not in m for m in result) assert result[1]["content"][0]["summary"][0]["signature"] == "sig-abc" + + +def test_existing_reasoning_block_is_not_duplicated(): + """Databricks rejects a message holding two reasoning blocks, so a message that already + carries one keeps it and the redundant thinking_blocks are dropped instead.""" + already_converted = { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "earlier", "signature": "sig-earlier"}], + } + + result = DatabricksConfig()._move_reasoning_into_content_block( + { + "role": "assistant", + "content": [already_converted, {"type": "text", "text": "391"}], + "thinking_blocks": [{"type": "thinking", "thinking": "later", "signature": "sig-later"}], + } + ) + + assert "thinking_blocks" not in result + assert [b for b in result["content"] if b.get("type") == "reasoning"] == [already_converted]