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