diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 8a625569cfa..7e61b231764 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -428,6 +428,8 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): # Move message-level cache_control into a content block when content is a string. if "cache_control" in _message and isinstance(_message.get("content"), str): _message = self._move_cache_control_into_string_content_block(_message) + if "thinking_blocks" in _message or "reasoning_content" in _message: + _message = self._move_reasoning_into_content_block(_message) _sanitize_empty_content(cast(dict[str, Any], _message)) new_messages.append(_message) @@ -464,6 +466,48 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): ] return cast(AllMessageValues, transformed_message) + def _move_reasoning_into_content_block(self, message: AllMessageValues) -> AllMessageValues: + """ + Converts LiteLLM's message-level reasoning fields into the reasoning content block + Databricks accepts, so extended thinking survives a multi-turn round trip. + + Transforms: + {"role": "assistant", "content": "text", "thinking_blocks": [{"thinking": "t", "signature": "s"}]} + Into: + {"role": "assistant", "content": [ + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "t", "signature": "s"}]}, + {"type": "text", "text": "text"}, + ]} + + 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. It also rejects a reasoning block as the final block of an assistant message, + which happens when the turn was thinking plus a tool call: the call moves to + message-level tool_calls and leaves nothing behind it. So the block is only injected + when other content follows it, and a message already carrying a reasoning block never + gains a second one. Otherwise 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. + """ + dropped: Final = ("thinking_blocks", "reasoning_content") + stripped: Final = {k: v for k, v in message.items() if k not in dropped} # mutable-ok: outbound provider JSON + content: Final = stripped.get("content") + as_text: Final = ({"type": "text", "text": content},) # mutable-ok: outbound provider JSON + listed: Final = tuple(content) if isinstance(content, list) else () + existing: Final = as_text if isinstance(content, str) and content else listed + signed: Final = next( + (b for b in message.get("thinking_blocks") or () if isinstance(b, dict) and b.get("signature")), + None, + ) + holds_reasoning: Final = any(isinstance(b, dict) and b.get("type") == "reasoning" for b in existing) + replace: Final = signed is not None and not holds_reasoning and bool(existing) + text: Final = (signed.get("thinking") or "") if signed is not None else "" + signature: Final = (signed.get("signature") or "") if signed is not None else "" + entry: Final = {"type": "summary_text", "text": text, "signature": signature} # mutable-ok: provider JSON + block: Final = {"type": "reasoning", "summary": [entry]} # mutable-ok: outbound provider JSON + rebuilt: Final = {**stripped, "content": [block, *existing]} # mutable-ok: outbound provider JSON + return cast(AllMessageValues, rebuilt if replace else stripped) # cast-ok: same TypedDict, shape unchanged + @staticmethod def extract_content_str( content: AllDatabricksContentValues | None, 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 41fb2589655..a06f586de32 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 @@ -493,3 +493,167 @@ def test_chunk_parser_without_usage_still_parses_content(): assert result.id == "chatcmpl-test" assert result.model == "databricks-claude-sonnet-5" assert result.choices[0]["delta"]["content"] == "hi" + + +def _thinking_message(thinking_blocks, content="391"): + return {"role": "assistant", "content": content, "thinking_blocks": thinking_blocks} + + +def test_thinking_blocks_become_a_reasoning_content_block(): + """Databricks rejects the message-level thinking_blocks key outright, so extended + thinking never survived a second turn. It must become a reasoning content block + carrying exactly one signed summary entry.""" + result = DatabricksConfig()._move_reasoning_into_content_block( + _thinking_message([{"type": "thinking", "thinking": "let me work it out", "signature": "sig-abc"}]) + ) + + assert "thinking_blocks" not in result + assert result["content"] == [ + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "let me work it out", "signature": "sig-abc"}], + }, + {"type": "text", "text": "391"}, + ] + + +def test_thinking_blocks_conversion_preserves_existing_content_blocks(): + result = DatabricksConfig()._move_reasoning_into_content_block( + _thinking_message( + [{"type": "thinking", "thinking": "", "signature": "sig-abc"}], + content=[{"type": "text", "text": "391"}, {"type": "text", "text": "trailing"}], + ) + ) + + assert result["content"][0]["type"] == "reasoning" + assert result["content"][1:] == [{"type": "text", "text": "391"}, {"type": "text", "text": "trailing"}] + + +def test_thinking_blocks_collapse_to_a_single_summary_entry(): + """Databricks accepts exactly one reasoning block with one summary entry; sending two + of either is rejected, so extra blocks are dropped rather than sent in a refused shape.""" + result = DatabricksConfig()._move_reasoning_into_content_block( + _thinking_message( + [ + {"type": "thinking", "thinking": "first", "signature": "sig-1"}, + {"type": "thinking", "thinking": "second", "signature": "sig-2"}, + ] + ) + ) + + reasoning_blocks = [b for b in result["content"] if b.get("type") == "reasoning"] + assert len(reasoning_blocks) == 1 + assert reasoning_blocks[0]["summary"] == [{"type": "summary_text", "text": "first", "signature": "sig-1"}] + + +@pytest.mark.parametrize( + "thinking_blocks", + [ + None, + [], + [{"type": "redacted_thinking", "data": "opaque"}], + ], + ids=["none", "empty", "unsigned_only"], +) +def test_unreplayable_thinking_blocks_are_dropped(thinking_blocks): + """A signature is mandatory, so blocks that lack one cannot be replayed. The key is + still removed, since leaving it would fail the request outright.""" + result = DatabricksConfig()._move_reasoning_into_content_block(_thinking_message(thinking_blocks)) + + assert "thinking_blocks" not in result + assert result["content"] == "391" + + +def test_transform_messages_strips_thinking_blocks(): + messages = [ + {"role": "user", "content": "what is 17 * 23?"}, + _thinking_message([{"type": "thinking", "thinking": "t", "signature": "sig-abc"}]), + {"role": "user", "content": "now double it"}, + ] + + result = DatabricksConfig()._transform_messages(messages=messages, model="databricks-claude-opus-5") + + assert all("thinking_blocks" not in m for m in result) + assert result[1]["content"][0]["type"] == "reasoning" + + +def test_reasoning_content_is_dropped_even_when_empty(): + """reasoning_content is an empty string on Databricks responses, so a truthiness check + would leave it in place and the request would still be rejected.""" + result = DatabricksConfig()._move_reasoning_into_content_block( + {"role": "assistant", "content": "391", "reasoning_content": ""} + ) + + assert "reasoning_content" not in result + assert result["content"] == "391" + + +def test_transform_messages_strips_both_reasoning_fields(): + """A verbatim assistant turn from a Databricks thinking response carries both keys.""" + messages = [ + {"role": "user", "content": "what is 17 * 23?"}, + { + "role": "assistant", + "content": "391", + "reasoning_content": "", + "thinking_blocks": [{"type": "thinking", "thinking": "", "signature": "sig-abc"}], + }, + {"role": "user", "content": "now double it"}, + ] + + result = DatabricksConfig()._transform_messages(messages=messages, model="databricks-claude-opus-5") + + 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] + + +def test_reasoning_block_is_not_emitted_as_the_only_content_block(): + """Databricks rejects a reasoning block as the final block of an assistant message. That + is exactly a thinking-plus-tool-call turn, where the call moves to message-level + tool_calls and leaves nothing after the reasoning block, so the block is dropped.""" + result = DatabricksConfig()._move_reasoning_into_content_block( + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "bash", "arguments": "{}"}}], + "thinking_blocks": [{"type": "thinking", "thinking": "", "signature": "sig-abc"}], + } + ) + + assert "thinking_blocks" not in result + assert not any(isinstance(b, dict) and b.get("type") == "reasoning" for b in result.get("content") or []) + assert result["tool_calls"][0]["id"] == "call_1" + + +def test_reasoning_block_still_emitted_when_text_follows(): + """The normal case must keep working: text after the reasoning block is valid.""" + result = DatabricksConfig()._move_reasoning_into_content_block( + { + "role": "assistant", + "content": "391", + "thinking_blocks": [{"type": "thinking", "thinking": "t", "signature": "sig-abc"}], + } + ) + + assert result["content"][0]["type"] == "reasoning" + assert result["content"][1] == {"type": "text", "text": "391"}