Merge pull request #39409 from BerriAI/litellm_databricks_strip_thinking_blocks

fix(databricks): strip thinking_blocks and reasoning_content from outbound messages

litellm/litellm_core_utils/prompt_templates/common_utils.py

@@ -1554,6 +1554,22 @@ def with_prompt_cache_breakpoint(target: _MarkedT, marker: object) -> _MarkedT:
     return cast(_MarkedT, marked)  # cast-ok: same block shape as the input plus the marker key
 
 
+LITELLM_INTERNAL_MESSAGE_FIELDS: Final = frozenset({"thinking_blocks", "reasoning_content", "provider_specific_fields"})
+
+
+def strip_litellm_internal_message_fields(message: AllMessageValues) -> AllMessageValues:
+    """Drop the fields litellm attaches to assistant messages (e.g. when translating Anthropic thinking
+    blocks) that OpenAI-compatible endpoints with strict schemas reject as extra inputs."""
+    if LITELLM_INTERNAL_MESSAGE_FIELDS.isdisjoint(message):
+        return message
+    return cast(  # cast-ok: same TypedDict minus internal keys
+        AllMessageValues,
+        {  # mutable-ok: provider transforms mutate message dicts in place downstream
+            key: value for key, value in message.items() if key not in LITELLM_INTERNAL_MESSAGE_FIELDS
+        },
+    )
+
+
 def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Any:
     """
     Filters a value from a dictionary

litellm/llms/databricks/chat/transformation.py

@@ -3,7 +3,7 @@ Translates from OpenAI's `/v1/chat/completions` to Databricks' `/chat/completion
 """
 
 import os
-from collections.abc import AsyncIterator, Coroutine, Iterator
+from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
 from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
 
 import httpx
@@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
     _should_convert_tool_call_to_json_mode,
 )
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    strip_litellm_internal_message_fields,
     strip_name_from_message,
 )
 from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
@@ -55,6 +56,14 @@ from ...openai_like.chat.transformation import OpenAILikeChatConfig
 from ..common_utils import DatabricksBase, DatabricksException
 
 
+def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool:
+    """Databricks rejects assistant messages with neither content nor tool calls, e.g. a replayed
+    thinking-only turn once its `thinking_blocks` are stripped."""
+    return message_dict.get("role") == "assistant" and not any(
+        message_dict.get(key) for key in ("content", "tool_calls", "function_call")
+    )
+
+
 def _sanitize_empty_content(message_dict: dict[str, Any]) -> None:
     """
     Remove or filter content so empty text blocks are not sent.
@@ -423,6 +432,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
         """
         Databricks does not support:
         - 'name' in user message.
+        - litellm's internal `thinking_blocks` / `reasoning_content` on assistant messages.
         """
         new_messages = []
         for idx, message in enumerate(messages):
@@ -431,10 +441,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
             else:
                 _message = message
             _message = strip_name_from_message(_message, allowed_name_roles=["user"])
+            _message = strip_litellm_internal_message_fields(_message)
             # 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)
             _sanitize_empty_content(cast(dict[str, Any], _message))
+            if _is_bare_assistant_message(_message):
+                continue
             new_messages.append(_message)
 
         if "claude" not in model:

tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py

@@ -255,6 +255,82 @@ def test_transform_messages_sanitizes_empty_content():
     assert result[1]["content"] == "Hi"
 
 
+def test_transform_request_strips_thinking_blocks_and_reasoning_content():
+    """Regression for LIT-6762: replaying an assistant turn that litellm decorated with
+    `thinking_blocks` / `reasoning_content` made Databricks 400 with
+    'messages.N.thinking_blocks: Extra inputs are not permitted'."""
+    config = DatabricksConfig()
+    messages = [
+        {"role": "user", "content": "hi"},
+        {
+            "role": "assistant",
+            "content": "Hello! How can I help?",
+            "thinking_blocks": [
+                {"type": "thinking", "thinking": "greet briefly", "signature": "sig_abc", "cache_control": {}}
+            ],
+            "reasoning_content": "greet briefly",
+            "provider_specific_fields": {"foo": "bar"},
+        },
+        {"role": "user", "content": "thanks"},
+    ]
+
+    result = config.transform_request(
+        model="databricks-claude-opus-5",
+        messages=messages,
+        optional_params={},
+        litellm_params={},
+        headers={},
+    )["messages"]
+
+    assert result[1] == {"role": "assistant", "content": "Hello! How can I help?"}
+    assert not any(
+        key in message
+        for message in result
+        for key in ("thinking_blocks", "reasoning_content", "provider_specific_fields")
+    )
+    assert "thinking_blocks" in messages[1]
+
+
+def test_transform_request_drops_thinking_only_assistant_turn_but_keeps_tool_call_turn():
+    """A replayed thinking-only assistant turn has nothing left once `thinking_blocks` are stripped, so it must be
+    dropped instead of being sent as a bare {"role": "assistant"}. A thinking + tool_use turn keeps its tool_calls."""
+    config = DatabricksConfig()
+    tool_call = {"id": "call_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}
+    messages = [
+        {"role": "user", "content": "hi"},
+        {
+            "role": "assistant",
+            "content": None,
+            "thinking_blocks": [{"type": "thinking", "thinking": "hmm", "signature": "sig_1"}],
+            "reasoning_content": "hmm",
+        },
+        {"role": "user", "content": "again"},
+        {
+            "role": "assistant",
+            "content": None,
+            "thinking_blocks": [{"type": "thinking", "thinking": "call f", "signature": "sig_2"}],
+            "reasoning_content": "call f",
+            "tool_calls": [tool_call],
+        },
+        {"role": "tool", "tool_call_id": "call_1", "content": "ok"},
+    ]
+
+    result = config.transform_request(
+        model="databricks-claude-opus-5",
+        messages=messages,
+        optional_params={},
+        litellm_params={},
+        headers={},
+    )["messages"]
+
+    assert result == [
+        {"role": "user", "content": "hi"},
+        {"role": "user", "content": "again"},
+        {"role": "assistant", "tool_calls": [tool_call]},
+        {"role": "tool", "tool_call_id": "call_1", "content": "ok"},
+    ]
+
+
 def _parallel_tool_calls():
     return [
         {