diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 60081c38073..bb5834f7dc7 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.anthropic_beta_gated_fields import log_beta_gated_strips, strip_beta_gated_thinking from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -81,15 +83,120 @@ 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 _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 + 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: 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( + 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: + 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 + + +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. @@ -108,9 +215,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") @@ -535,6 +640,27 @@ 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"), + 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, + ) + @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 b93cebdbc4c..79572a6184b 100644 --- a/tests/integration/providers/test_databricks_chat_wire.py +++ b/tests/integration/providers/test_databricks_chat_wire.py @@ -353,9 +353,157 @@ 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-history", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"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": "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-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-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", + "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", + ), + 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_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, @@ -365,7 +513,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 19234767b19..9b9a6dae502 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -539,13 +539,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", [ @@ -603,6 +596,239 @@ 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], 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={"drop_params": True} if drop_params else {}, + 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": "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]}], + _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_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": _TWO_TEXT_BLOCKS}, + ] + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", messages, _JSON_OBJECT) == [ + {"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 + ] + + +_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) + + 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