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] 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