From 29647eb343f0381a3c85969ea391a23326f9a47b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 06:50:36 -0700 Subject: [PATCH 1/3] fix(databricks): join plain text blocks for non-Claude json_object requests --- .../llms/databricks/chat/transformation.py | 78 ++++++++++++- .../providers/test_databricks_chat_wire.py | 70 +++++++++++- .../test_databricks_chat_transformation.py | 104 ++++++++++++++++-- 3 files changed, 237 insertions(+), 15 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index a98cfdfe52d..8d625f46578 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -75,15 +75,67 @@ def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool: ) +_CONTENT_BLOCKS: Final = TypeAdapter(tuple[object, ...]) + + +def _plain_text(block: object) -> str | None: + match block: + case {"type": "text", "text": str() as text} if block == {"type": "text", "text": text}: + return text + case _: + return None + + +def _plain_texts(message: Mapping[str, object]) -> tuple[str, ...]: + match content := message.get("content"): + case list(): + blocks: Final = _CONTENT_BLOCKS.validate_python(content) + texts: Final = tuple(text for block in blocks if (text := _plain_text(block)) is not None) + return texts if len(texts) == len(blocks) else () + case _: + return () + + def _collapse_single_text_block(message: Mapping[str, object]) -> Mapping[str, object]: - content: Final = message.get("content") - match content: - case [{"type": "text", "text": str() as text}] if content == [{"type": "text", "text": text}]: + match _plain_texts(message): + case (text,): return {**message, "content": text} case _: return message +def _join_plain_text_blocks(message: AllMessageValues) -> AllMessageValues: + match _plain_texts(message): + case (): + return message + case texts: + return cast(AllMessageValues, {**message, "content": "".join(texts)}) # cast-ok: spread keeps the shape + + +def _json_object_ready_messages( + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited transform_request contract + response_format: object, +) -> list[AllMessageValues]: # mutable-ok: inherited transform_request contract + match response_format: + case {"type": "json_object"} if "claude" not in model: + return [_join_plain_text_blocks(message) for message in messages] + case _: + return messages + + +def _is_empty_text_block(block: object) -> bool: + match block: + case {"type": "text", "text": str() as text}: + return not text.strip() + case {"type": "text", "text": object() as text}: + return not text + case {"type": "text"}: + return True + case _: + return False + + def _sanitize_empty_content(message_dict: dict[str, object]) -> None: """ Remove or filter content so empty text blocks are not sent. @@ -102,9 +154,7 @@ def _sanitize_empty_content(message_dict: dict[str, object]) -> None: message_dict.pop("content") return filtered: Final = [ - block - for block in content - if not (isinstance(block, dict) and block.get("type") == "text" and not (block.get("text") or "").strip()) + block for block in _CONTENT_BLOCKS.validate_python(content) if not _is_empty_text_block(block) ] if not filtered: message_dict.pop("content") @@ -452,6 +502,22 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): return False + def transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + headers: dict[str, object], # mutable-ok: inherited contract + ) -> dict[str, object]: # mutable-ok: inherited contract + return super().transform_request( + model=model, + messages=_json_object_ready_messages(model, messages, optional_params.get("response_format")), + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + @overload def _transform_messages( self, messages: list[AllMessageValues], model: str, is_async: Literal[True] diff --git a/tests/integration/providers/test_databricks_chat_wire.py b/tests/integration/providers/test_databricks_chat_wire.py index 3d7f273279f..e2189d0bafe 100644 --- a/tests/integration/providers/test_databricks_chat_wire.py +++ b/tests/integration/providers/test_databricks_chat_wire.py @@ -313,9 +313,75 @@ def _shape_reply(identity: str, stream: bool) -> Reply: "/chat/completions", id="A5-responses-list-input", ), + pytest.param( + "json-object-two-blocks", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A6-chat-json-object-two-blocks-joined", + ), + pytest.param( + "json-object-two-blocks-stream", + "/v1/chat/completions", + "databricks-qwen35-122b-a10b", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + "stream": True, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A7-chat-stream-json-object-two-blocks-joined", + ), + pytest.param( + "json-object-system-blocks", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "system", "content": "Bravo rule."}, + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "system", "content": "Alpha rule.\n\nBravo rule."}, {"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A8-chat-json-object-system-blocks-keep-break", + ), + pytest.param( + "json-object-responses-two-parts", + "/v1/responses", + "databricks-qwen35-122b-a10b", + { + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": "Reply"}, {"type": "input_text", "text": " in JSON"}], + } + ], + "text": {"format": {"type": "json_object"}}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A9-responses-json-object-two-parts-joined", + ), ], ) -def test_databricks_non_claude_single_text_block_reaches_upstream_as_string( +def test_databricks_non_claude_plain_text_content_reaches_upstream_as_string( gateway: Gateway, case: str, endpoint: str, @@ -325,7 +391,7 @@ def test_databricks_non_claude_single_text_block_reaches_upstream_as_string( responses_mode: bool, upstream_path: str, ) -> None: - identity: Final = f"databricks-single-text-{case}-{uuid.uuid4().hex}" + identity: Final = f"databricks-plain-text-{case}-{uuid.uuid4().hex}" streaming: Final = request_body.get("stream") is True def respond(request: Request) -> Reply: diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index d25787d7028..24d9210ec76 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -483,13 +483,6 @@ def test_transform_request_preserves_unity_model_service_name(): [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], id="cache-control", ), - pytest.param( - "databricks-meta-llama-3-3-70b-instruct", - [{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}], - {"response_format": {"type": "json_object"}}, - [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}], - id="multiple-text-blocks", - ), pytest.param( "databricks-meta-llama-3-3-70b-instruct", [ @@ -547,6 +540,103 @@ def test_transform_request_collapses_single_plain_text_block_for_non_claude_mode assert outgoing_messages == [{**messages[0], "content": expected_content}] +_TWO_TEXT_BLOCKS: Final = [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}] +_JSON_OBJECT: Final = {"response_format": {"type": "json_object"}} +_IMAGE_BLOCK: Final = {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}} +_CACHED_BLOCK: Final = {"type": "text", "text": " in JSON", "cache_control": {"type": "ephemeral"}} + + +def _outgoing_messages( + model: str, messages: list[AllMessageValues], optional_params: dict[str, object] +) -> list[dict[str, object]]: + request: Final = TypeAdapter(dict[str, object]).validate_python( + DatabricksConfig().transform_request( # pyright: ignore[reportUnknownMemberType] # inherited hook is untyped + model=model, + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + ) + return TypeAdapter(list[dict[str, object]]).validate_python(request["messages"]) + + +@pytest.mark.parametrize( + "model,messages,optional_params,expected_content", + [ + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + "Reply in JSON", + id="json-object-two-blocks-joined", + ), + pytest.param( + "databricks-qwen35-122b-a10b", + [{"role": "user", "content": [{"type": "text", "text": "A"}, {"type": "text", "text": "B json"}] * 2}], + _JSON_OBJECT, + "AB jsonAB json", + id="json-object-four-blocks-joined-as-sent", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + {}, + _TWO_TEXT_BLOCKS, + id="no-response-format-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + {"response_format": {"type": "json_schema", "json_schema": {"name": "r", "schema": {"type": "object"}}}}, + _TWO_TEXT_BLOCKS, + id="json-schema-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply"}, _CACHED_BLOCK], + id="json-object-cache-control-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}, _IMAGE_BLOCK]}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply in JSON"}, _IMAGE_BLOCK], + id="json-object-text-and-image-stays-list", + ), + pytest.param( + "databricks-claude-sonnet-5", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + _TWO_TEXT_BLOCKS, + id="json-object-claude-stays-list", + ), + ], +) +def test_transform_request_joins_plain_text_blocks_only_for_non_claude_json_object( + model: str, + messages: list[AllMessageValues], + optional_params: dict[str, object], + expected_content: object, +) -> None: + assert _outgoing_messages(model, messages, optional_params) == [{"role": "user", "content": expected_content}] + + +def test_transform_request_json_object_join_keeps_the_break_between_system_messages() -> None: + messages: Final[list[AllMessageValues]] = [ + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "system", "content": "Bravo rule."}, + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + ] + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", messages, _JSON_OBJECT) == [ + {"role": "system", "content": "Alpha rule.\n\nBravo rule."}, + {"role": "user", "content": "Reply in JSON"}, + ] + + 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 From f3d459e844a5aa6855c72d4fe8a27c03d2083ac6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:37:05 -0700 Subject: [PATCH 2/3] fix(databricks): join only the last message for json_object and drop cache_control under drop_params Databricks checks for "json" only in the last message, so earlier turns now go out as sent and only the last message, plus any system messages merged into it, is joined into a string. Under drop_params a text block's cache_control is dropped so that message can be joined too, and without it the request still goes out as a list. --- .../llms/databricks/chat/transformation.py | 60 +++++++++++++-- .../providers/test_databricks_chat_wire.py | 53 +++++++++++-- .../test_databricks_chat_transformation.py | 76 +++++++++++++++++-- 3 files changed, 171 insertions(+), 18 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index e2a10b18c92..71d4132f10c 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -4,12 +4,14 @@ Translates from OpenAI's `/v1/chat/completions` to Databricks' `/chat/completion import os from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence +from itertools import takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload import httpx from pydantic import BaseModel, ConfigDict, TypeAdapter +import litellm from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( handle_invalid_parallel_tool_calls, @@ -109,22 +111,63 @@ def _collapse_single_text_block(message: Mapping[str, object]) -> Mapping[str, o return message -def _join_plain_text_blocks(message: AllMessageValues) -> AllMessageValues: - match _plain_texts(message): +def _without_text_cache_control(block: object) -> object: + match block: + case {"type": "text", "text": str() as text, "cache_control": object() as cache_control} if block == { + "type": "text", + "text": text, + "cache_control": cache_control, + }: + return {"type": "text", "text": text} + case _: + return block + + +def _without_cache_control(message: AllMessageValues) -> AllMessageValues: + uncached: Final = {key: value for key, value in message.items() if key != "cache_control"} + match content := uncached.get("content"): + case list(): + blocks: Final = [_without_text_cache_control(block) for block in _CONTENT_BLOCKS.validate_python(content)] + return cast(AllMessageValues, {**uncached, "content": blocks}) # cast-ok: spread keeps the shape + case _: + return cast(AllMessageValues, uncached) # cast-ok: only cache_control was left out + + +def _has_string_content(message: Mapping[str, object]) -> bool: + return isinstance(message.get("content"), str) + + +def _is_system(message: Mapping[str, object]) -> bool: + return message.get("role") == "system" + + +def _json_mode_message(message: AllMessageValues, drop_params: bool) -> AllMessageValues: + candidate: Final = _without_cache_control(message) if drop_params else message + match _plain_texts(candidate): case (): - return message + return candidate if _has_string_content(candidate) else message case texts: - return cast(AllMessageValues, {**message, "content": "".join(texts)}) # cast-ok: spread keeps the shape + return cast(AllMessageValues, {**candidate, "content": "".join(texts)}) # cast-ok: spread keeps the shape + + +def _json_checked_start(messages: Sequence[AllMessageValues]) -> int: + trailing_system: Final = sum(1 for _ in takewhile(_is_system, reversed(messages))) + return max(len(messages) - max(trailing_system, 1), 0) def _json_object_ready_messages( model: str, messages: list[AllMessageValues], # mutable-ok: inherited transform_request contract response_format: object, + drop_params: bool, ) -> list[AllMessageValues]: # mutable-ok: inherited transform_request contract match response_format: case {"type": "json_object"} if "claude" not in model: - return [_join_plain_text_blocks(message) for message in messages] + checked_start: Final = _json_checked_start(messages) + return [ + *messages[:checked_start], + *(_json_mode_message(message, drop_params) for message in messages[checked_start:]), + ] case _: return messages @@ -586,7 +629,12 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): ) -> dict[str, object]: # mutable-ok: inherited contract return super().transform_request( model=model, - messages=_json_object_ready_messages(model, messages, optional_params.get("response_format")), + messages=_json_object_ready_messages( + model, + messages, + optional_params.get("response_format"), + drop_params=litellm.drop_params is True or litellm_params.get("drop_params") is True, + ), optional_params=optional_params, litellm_params=litellm_params, headers=headers, diff --git a/tests/integration/providers/test_databricks_chat_wire.py b/tests/integration/providers/test_databricks_chat_wire.py index d280e5a6dce..1ba41d067c6 100644 --- a/tests/integration/providers/test_databricks_chat_wire.py +++ b/tests/integration/providers/test_databricks_chat_wire.py @@ -385,21 +385,64 @@ def _shape_reply(identity: str, stream: bool) -> Reply: id="A7-chat-stream-json-object-two-blocks-joined", ), pytest.param( - "json-object-system-blocks", + "json-object-history", "/v1/chat/completions", "databricks-meta-llama-3-3-70b-instruct", { "messages": [ - {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, - {"role": "system", "content": "Bravo rule."}, + {"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]}, + {"role": "assistant", "content": "Got it."}, {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, ], "response_format": {"type": "json_object"}, }, - [{"role": "system", "content": "Alpha rule.\n\nBravo rule."}, {"role": "user", "content": "Reply in JSON"}], + [ + {"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]}, + {"role": "assistant", "content": "Got it."}, + {"role": "user", "content": "Reply in JSON"}, + ], False, "/chat/completions", - id="A8-chat-json-object-system-blocks-keep-break", + id="A8-chat-json-object-earlier-turns-kept", + ), + pytest.param( + "json-object-trailing-system", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": "Say hi"}, + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "system", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Say hi"}, {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}], + False, + "/chat/completions", + id="A8b-chat-json-object-trailing-system-keeps-break", + ), + pytest.param( + "json-object-drop-params-cache-control", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Reply"}, + {"type": "text", "text": " in JSON", "cache_control": {"type": "ephemeral"}}, + ], + } + ], + "response_format": {"type": "json_object"}, + "drop_params": True, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A8c-chat-json-object-drop-params-drops-cache-control", ), pytest.param( "json-object-responses-two-parts", diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index 1d1f902017a..17e9bc17ab7 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -603,14 +603,14 @@ _CACHED_BLOCK: Final = {"type": "text", "text": " in JSON", "cache_control": {"t def _outgoing_messages( - model: str, messages: list[AllMessageValues], optional_params: dict[str, object] + model: str, messages: list[AllMessageValues], optional_params: dict[str, object], drop_params: bool = False ) -> list[dict[str, object]]: request: Final = TypeAdapter(dict[str, object]).validate_python( DatabricksConfig().transform_request( # pyright: ignore[reportUnknownMemberType] # inherited hook is untyped model=model, messages=messages, optional_params=optional_params, - litellm_params={}, + litellm_params={"drop_params": True} if drop_params else {}, headers={}, ) ) @@ -655,6 +655,13 @@ def _outgoing_messages( [{"type": "text", "text": "Reply"}, _CACHED_BLOCK], id="json-object-cache-control-stays-list", ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], + id="json-object-message-cache-control-stays-block", + ), pytest.param( "databricks-meta-llama-3-3-70b-instruct", [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}, _IMAGE_BLOCK]}], @@ -680,19 +687,74 @@ def test_transform_request_joins_plain_text_blocks_only_for_non_claude_json_obje assert _outgoing_messages(model, messages, optional_params) == [{"role": "user", "content": expected_content}] -def test_transform_request_json_object_join_keeps_the_break_between_system_messages() -> None: +def test_transform_request_json_object_keeps_earlier_turns_as_sent() -> None: + earlier_turn: Final[AllMessageValues] = { + "role": "user", + "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}], + } + reply: Final[AllMessageValues] = {"role": "assistant", "content": "Got it."} + + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [earlier_turn, reply, {"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + ) == [earlier_turn, reply, {"role": "user", "content": "Reply in JSON"}] + + +def test_transform_request_json_object_joins_a_trailing_system_run_and_keeps_its_break() -> None: messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": "Say hi"}, {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, - {"role": "system", "content": "Bravo rule."}, - {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + {"role": "system", "content": _TWO_TEXT_BLOCKS}, ] assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", messages, _JSON_OBJECT) == [ - {"role": "system", "content": "Alpha rule.\n\nBravo rule."}, - {"role": "user", "content": "Reply in JSON"}, + {"role": "user", "content": "Say hi"}, + {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}, ] +_CACHED_TEXT_AND_IMAGE: Final = [{**_CACHED_BLOCK, "text": "Reply in JSON"}, _IMAGE_BLOCK] + + +@pytest.mark.parametrize( + "message,expected", + [ + pytest.param( + {"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}, + {"role": "user", "content": "Reply in JSON"}, + id="block-cache-control-dropped-and-joined", + ), + pytest.param( + {"role": "user", "content": "Reply in JSON", "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "Reply in JSON"}, + id="message-cache-control-dropped", + ), + pytest.param( + {"role": "user", "content": _CACHED_TEXT_AND_IMAGE}, + {"role": "user", "content": _CACHED_TEXT_AND_IMAGE}, + id="image-keeps-cache-control", + ), + ], +) +def test_transform_request_json_object_drops_last_message_cache_control_under_drop_params( + message: AllMessageValues, expected: dict[str, object] +) -> None: + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", [message], _JSON_OBJECT, drop_params=True) == [ + expected + ] + + +def test_transform_request_json_object_reads_the_global_drop_params(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}], + _JSON_OBJECT, + ) == [{"role": "user", "content": "Reply in JSON"}] + + 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 From 1ecfd61dfa5375e9075090d7c2b86ffb14952762 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 23:04:07 -0700 Subject: [PATCH 3/3] fix(databricks): skip dropped assistant turns and blank parts when joining for json_object The json_object join picked the last message before the request transform drops bare assistant turns, so a trailing thinking-only or reasoning-only turn (a Responses input ending in a reasoning item included) left the real last user message as a list and Databricks still returned its 400. The pick now skips the turns that transform drops. The join also ran when only one part had text, which sent whitespace-only parts the existing empty block filter used to drop on requests that already worked. It now joins only two or more non-blank parts, so a request with one non-blank part goes out exactly as before. Without drop_params, a message with a message-level cache_control now goes out as sent instead of being rewritten into one cached block. --- .../llms/databricks/chat/transformation.py | 26 +++++-- .../providers/test_databricks_chat_wire.py | 39 ++++++++++ .../test_databricks_chat_transformation.py | 74 +++++++++++++++++++ 3 files changed, 132 insertions(+), 7 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 326c3f61e0a..bb5834f7dc7 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -144,16 +144,28 @@ def _is_system(message: Mapping[str, object]) -> bool: def _json_mode_message(message: AllMessageValues, drop_params: bool) -> AllMessageValues: candidate: Final = _without_cache_control(message) if drop_params else message - match _plain_texts(candidate): - case (): - return candidate if _has_string_content(candidate) else message - case texts: - return cast(AllMessageValues, {**candidate, "content": "".join(texts)}) # cast-ok: spread keeps the shape + texts: Final = _plain_texts(candidate) + if "cache_control" in candidate or not (texts or _has_string_content(candidate)): + return message + if sum(1 for text in texts if text.strip()) < 2: + return candidate + return cast(AllMessageValues, {**candidate, "content": "".join(texts)}) # cast-ok: spread keeps the shape + + +def _is_dropped_on_send(message: Mapping[str, object]) -> bool: + sanitized: Final = {**message} + _sanitize_empty_content(sanitized) + return _is_bare_assistant_message(sanitized) + + +def _is_system_or_dropped(message: Mapping[str, object]) -> bool: + return _is_system(message) or _is_dropped_on_send(message) def _json_checked_start(messages: Sequence[AllMessageValues]) -> int: - trailing_system: Final = sum(1 for _ in takewhile(_is_system, reversed(messages))) - return max(len(messages) - max(trailing_system, 1), 0) + trailing: Final = tuple(takewhile(_is_system_or_dropped, reversed(messages))) + last_sent_offset: Final = 0 if any(_is_system(message) for message in trailing) else 1 + return max(len(messages) - len(trailing) - last_sent_offset, 0) def _json_object_ready_messages( diff --git a/tests/integration/providers/test_databricks_chat_wire.py b/tests/integration/providers/test_databricks_chat_wire.py index 1ba41d067c6..79572a6184b 100644 --- a/tests/integration/providers/test_databricks_chat_wire.py +++ b/tests/integration/providers/test_databricks_chat_wire.py @@ -444,6 +444,26 @@ def _shape_reply(identity: str, stream: bool) -> Reply: "/chat/completions", id="A8c-chat-json-object-drop-params-drops-cache-control", ), + pytest.param( + "json-object-trailing-thinking-turn", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + { + "role": "assistant", + "content": "", + "thinking_blocks": [{"type": "thinking", "thinking": "Plan.", "signature": "sig"}], + }, + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A8d-chat-json-object-trailing-thinking-turn-dropped", + ), pytest.param( "json-object-responses-two-parts", "/v1/responses", @@ -462,6 +482,25 @@ def _shape_reply(identity: str, stream: bool) -> Reply: "/chat/completions", id="A9-responses-json-object-two-parts-joined", ), + pytest.param( + "json-object-responses-trailing-reasoning", + "/v1/responses", + "databricks-qwen35-122b-a10b", + { + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": "Reply"}, {"type": "input_text", "text": " in JSON"}], + }, + {"type": "reasoning", "id": "rs_1", "summary": [{"type": "summary_text", "text": "Plan."}]}, + ], + "text": {"format": {"type": "json_object"}}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A9b-responses-json-object-trailing-reasoning-dropped", + ), ], ) def test_databricks_non_claude_plain_text_content_reaches_upstream_as_string( diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index 8c01aa93a1d..9b9a6dae502 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -745,6 +745,80 @@ def test_transform_request_json_object_drops_last_message_cache_control_under_dr ] +_THINKING_BLOCK: Final = {"type": "thinking", "thinking": "Plan the reply.", "signature": "sig"} + + +@pytest.mark.parametrize( + "dropped_turn", + [ + pytest.param({"role": "assistant", "content": "", "thinking_blocks": [_THINKING_BLOCK]}, id="thinking-only"), + pytest.param({"role": "assistant", "content": None, "reasoning_content": "Plan."}, id="reasoning-only"), + pytest.param({"role": "assistant", "content": [{"type": "text", "text": " "}]}, id="blank-text-part"), + ], +) +def test_transform_request_json_object_joins_the_message_before_a_dropped_assistant_turn( + dropped_turn: AllMessageValues, +) -> None: + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}, dropped_turn], + _JSON_OBJECT, + ) == [{"role": "user", "content": "Reply in JSON"}] + + +def test_transform_request_json_object_joins_a_system_run_split_by_a_dropped_assistant_turn() -> None: + earlier_turn: Final[AllMessageValues] = { + "role": "user", + "content": [{"type": "text", "text": "Say"}, {"type": "text", "text": " hi"}], + } + messages: Final[list[AllMessageValues]] = [ + earlier_turn, + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "assistant", "content": "", "thinking_blocks": [_THINKING_BLOCK]}, + {"role": "system", "content": _TWO_TEXT_BLOCKS}, + ] + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", messages, _JSON_OBJECT) == [ + earlier_turn, + {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}, + ] + + +@pytest.mark.parametrize( + "content,expected_content", + [ + pytest.param( + [{"type": "text", "text": "Reply in JSON"}, {"type": "text", "text": "\n\n"}], + "Reply in JSON", + id="one-non-blank-part-goes-out-as-before", + ), + pytest.param( + [{"type": "text", "text": "Reply"}, {"type": "text", "text": " "}, {"type": "text", "text": "in JSON"}], + "Reply in JSON", + id="blank-separator-between-parts-kept", + ), + ], +) +def test_transform_request_json_object_joins_only_two_or_more_non_blank_parts( + content: list[object], expected_content: str +) -> None: + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", [{"role": "user", "content": content}], _JSON_OBJECT + ) == [{"role": "user", "content": expected_content}] + + +def test_transform_request_json_object_keeps_message_cache_control_parts_as_sent_without_drop_params() -> None: + message: Final[AllMessageValues] = { + "role": "user", + "content": _TWO_TEXT_BLOCKS, + "cache_control": {"type": "ephemeral"}, + } + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", [message], {}) == _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", [message], _JSON_OBJECT + ) + + def test_transform_request_json_object_reads_the_global_drop_params(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "drop_params", True)