Merge 29647eb343 into 8d6b8dab3d

litellm/llms/databricks/chat/transformation.py

@@ -79,15 +79,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.
@@ -106,9 +158,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")
@@ -528,6 +578,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]

tests/integration/providers/test_databricks_chat_wire.py

@@ -353,9 +353,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,
@@ -365,7 +431,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:

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