mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
Merge a48db2e7d7 into 22a349ee70
This commit is contained in:
pokepoke81
GitHub
commit
5a4ba9dfa0
2 changed files with 208 additions and 0 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue