Merge pull request #40449 from BerriAI/litellm_databricks_reasoning_content

fix(databricks): keep top-level reasoning_content from OpenAI-compatible gateway models

litellm/llms/databricks/chat/transformation.py

@@ -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 (
+    _extract_reasoning_content,  # pyright: ignore[reportPrivateUsage]  # same import as the OpenAI transformation
     strip_litellm_internal_message_fields,
     strip_name_from_message,
 )
@@ -23,7 +24,9 @@ from litellm.types.llms.anthropic import AllAnthropicToolsValues
 from litellm.types.llms.databricks import (
     AllDatabricksContentValues,
     DatabricksChoice,
+    DatabricksDelta,
     DatabricksFunction,
+    DatabricksMessage,
     DatabricksResponse,
     DatabricksTool,
 )
@@ -534,6 +537,19 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
                             thinking_blocks.append(thinking_block)
         return reasoning_content, thinking_blocks
 
+    @staticmethod
+    def extract_top_level_reasoning_content(delta: DatabricksDelta) -> str | None:
+        return delta.get("reasoning_content")
+
+    @staticmethod
+    def resolve_reasoning_and_content(
+        message: DatabricksMessage, block_reasoning_content: str | None
+    ) -> tuple[str | None, str | None]:
+        content_str: Final = DatabricksConfig.extract_content_str(message["content"])
+        if block_reasoning_content is not None:
+            return block_reasoning_content, content_str
+        return _extract_reasoning_content({**message, "content": content_str})
+
     @staticmethod
     def extract_citations(
         content: AllDatabricksContentValues | None,
@@ -577,14 +593,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
                     finish_reason = "stop"
 
             if translated_message is None:
-                ## get the content str
-                content_str = DatabricksConfig.extract_content_str(choice["message"]["content"])
-
-                ## get the reasoning content
                 (
-                    reasoning_content,
+                    block_reasoning_content,
                     thinking_blocks,
                 ) = DatabricksConfig.extract_reasoning_content(choice["message"].get("content"))
+                reasoning_content, content_str = DatabricksConfig.resolve_reasoning_and_content(
+                    choice["message"], block_reasoning_content
+                )
 
                 citations = DatabricksConfig.extract_citations(choice["message"].get("content"))
 
@@ -738,12 +753,16 @@ class DatabricksChatResponseIterator(BaseModelResponseIterator):
 
                 # extract the reasoning content
                 (
-                    reasoning_content,
+                    block_reasoning_content,
                     thinking_blocks,
                 ) = DatabricksConfig.extract_reasoning_content(choice["delta"].get("content"))
 
                 choice["delta"]["content"] = content_str
-                choice["delta"]["reasoning_content"] = reasoning_content
+                choice["delta"]["reasoning_content"] = (
+                    block_reasoning_content
+                    if block_reasoning_content is not None
+                    else DatabricksConfig.extract_top_level_reasoning_content(choice["delta"])
+                )
                 choice["delta"]["thinking_blocks"] = thinking_blocks
                 translated_choices.append(choice)
             return ModelResponseStream(

litellm/types/llms/databricks.py

@@ -2,6 +2,7 @@ from typing import Any, Literal
 
 from pydantic import BaseModel
 from typing_extensions import (
+    ReadOnly,
     Required,
     TypedDict,
 )
@@ -57,6 +58,14 @@ class DatabricksMessage(TypedDict, total=False):
     role: Required[str]
     content: Required[AllDatabricksContentValues]
     tool_calls: list[DatabricksTool] | None
+    reasoning_content: ReadOnly[str | None]
+    reasoning: ReadOnly[str | None]
+
+
+class DatabricksDelta(TypedDict, total=False):
+    role: ReadOnly[str]
+    content: ReadOnly[AllDatabricksContentValues | None]
+    reasoning_content: ReadOnly[str | None]
 
 
 class DatabricksChoice(TypedDict, total=False):

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

@@ -590,3 +590,87 @@ 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"
+
+
+@pytest.mark.parametrize("reasoning_key", ["reasoning_content", "reasoning"])
+def test_transform_choices_surfaces_top_level_reasoning_content(reasoning_key: str) -> None:
+    config = DatabricksConfig()
+    databricks_choices = [
+        {
+            "message": {
+                "role": "assistant",
+                "content": "391",
+                reasoning_key: "We need answer just number. 17*23=391.",
+            },
+            "index": 0,
+            "finish_reason": "stop",
+        }
+    ]
+
+    choices = config._transform_dbrx_choices(choices=databricks_choices)
+
+    assert choices[0].message.content == "391"
+    assert choices[0].message.reasoning_content == "We need answer just number. 17*23=391."
+    assert getattr(choices[0].message, "thinking_blocks", None) is None
+
+
+def test_transform_choices_parses_think_tags_in_string_content():
+    config = DatabricksConfig()
+    databricks_choices = [
+        {
+            "message": {"role": "assistant", "content": "<think>17 times 23</think>391"},
+            "index": 0,
+            "finish_reason": "stop",
+        }
+    ]
+
+    choices = config._transform_dbrx_choices(choices=databricks_choices)
+
+    assert choices[0].message.content == "391"
+    assert choices[0].message.reasoning_content == "17 times 23"
+
+
+def test_transform_choices_prefers_reasoning_blocks_over_top_level_field():
+    config = DatabricksConfig()
+    databricks_choices = [
+        {
+            "message": {
+                "role": "assistant",
+                "content": [
+                    {"type": "reasoning", "summary": [{"type": "summary_text", "text": "from block"}]},
+                    {"type": "text", "text": "391"},
+                ],
+                "reasoning_content": "from field",
+            },
+            "index": 0,
+            "finish_reason": "stop",
+        }
+    ]
+
+    choices = config._transform_dbrx_choices(choices=databricks_choices)
+
+    assert choices[0].message.reasoning_content == "from block"
+    assert choices[0].message.content == "391"
+
+
+@pytest.mark.parametrize("reasoning_key", ["reasoning_content", "reasoning"])
+def test_chunk_parser_surfaces_top_level_reasoning_delta(reasoning_key: str) -> None:
+    iterator = DatabricksChatResponseIterator(None, sync_stream=True)
+    chunk = {
+        "id": "1",
+        "object": "chat.completion.chunk",
+        "created": 0,
+        "model": "lit-qa-deepseek-v4-flash",
+        "choices": [
+            {
+                "delta": {"role": "assistant", "content": None, reasoning_key: "We need answer"},
+                "index": 0,
+                "finish_reason": None,
+            }
+        ],
+    }
+
+    parsed = iterator.chunk_parser(chunk)
+
+    assert parsed.choices[0].delta.reasoning_content == "We need answer"
+    assert parsed.choices[0].delta.content is None