From 1aee1a21edec8606f6128ae97afe4f3c71bfe01f Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Wed, 14 Oct 2026 05:22:02 -0400 Subject: [PATCH] fix(databricks): send a single plain text block as a string for non-Claude models (#45686) * fix(databricks): send a single plain text block as a string for non-Claude models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(databricks): simplify single-text-block helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(databricks): clarify json_object validation wording Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(databricks): integration cells for single text block content shape Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(databricks): rename list-content cells to stays_list Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(responses): expect plain string content for bridged Databricks messages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(databricks): cover the customer's Unity Catalog model from #45617 and #45618 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(databricks): drop casts from _collapse_single_text_block Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(databricks): assert the whole outgoing messages list in the collapse unit test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(databricks): cover system joins, fallback, SDK clients, and outage burst for one-block content --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../llms/databricks/chat/transformation.py | 12 +- .../providers/test_databricks_chat_wire.py | 577 ++++++++++++++++++ .../test_databricks_chat_transformation.py | 116 ++++ .../test_responses_api_bridge_flag.py | 7 +- 4 files changed, 706 insertions(+), 6 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 4b6f178fa1a..a98cfdfe52d 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -75,6 +75,15 @@ def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool: ) +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}]: + return {**message, "content": text} + case _: + return message + + def _sanitize_empty_content(message_dict: dict[str, object]) -> None: """ Remove or filter content so empty text blocks are not sent. @@ -463,6 +472,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): Databricks does not support: - 'name' in user message. - litellm's internal `thinking_blocks` / `reasoning_content` on assistant messages. + - non-Claude json_object "json" checks read only string content, so a single text block is sent as a string. """ new_messages = [] for idx, message in enumerate(messages): @@ -478,7 +488,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): _sanitize_empty_content(cast(dict[str, object], _message)) if _is_bare_assistant_message(_message): continue - new_messages.append(_message) + new_messages.append(_collapse_single_text_block(_message) if "claude" not in model else _message) if "claude" not in model: new_messages = _split_parallel_tool_calls( diff --git a/tests/integration/providers/test_databricks_chat_wire.py b/tests/integration/providers/test_databricks_chat_wire.py index f6192d32c31..3d7f273279f 100644 --- a/tests/integration/providers/test_databricks_chat_wire.py +++ b/tests/integration/providers/test_databricks_chat_wire.py @@ -3,11 +3,15 @@ import json import uuid from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor +from pathlib import Path from typing import Final, Literal, TypeAlias +import anthropic +import openai import pytest from integration._support.client import Gateway, eventually from integration._support.database import read_rows +from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter @@ -201,6 +205,226 @@ def test_databricks_stream_final_usage_chunk_reaches_client_and_spend_log(gatewa assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"], rows[0]["total_tokens"]) == (12011, 8, 12019) +def _shape_reply(identity: str, stream: bool) -> Reply: + if stream: + chunks: Final = ( + _frame( + identity, + [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": None}], + ), + _frame(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}]), + b"data: [DONE]\n\n", + ) + return Reply(content_type="text/event-stream", chunks=chunks) + + body: Final = { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": _BACKEND, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + return Reply(body=json.dumps(body).encode()) + + +@pytest.mark.parametrize( + "case,endpoint,provider_model,request_body,expected_messages,responses_mode,upstream_path", + [ + pytest.param( + "chat_nonstream", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A1-chat-nonstream", + ), + pytest.param( + "chat_stream", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}], + "response_format": {"type": "json_object"}, + "stream": True, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A2-chat-stream", + ), + pytest.param( + "responses-mode-bridge", + "/v1/chat/completions", + "databricks-meta-llama-3-1-8b-instruct", + { + "messages": [{"role": "user", "content": "Reply in JSON"}], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Reply in JSON"}], + True, + "/chat/completions", + id="A3-responses-mode", + ), + pytest.param( + "responses-mode-bridge-unity-catalog", + "/v1/chat/completions", + "system.ai.deepseek-v4-1-flash", + { + "messages": [{"role": "user", "content": "Reply in JSON with key a. hi"}], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Reply in JSON with key a. hi"}], + True, + "/ai-gateway/mlflow/v1/chat/completions", + id="A3b-responses-mode-unity-catalog-model", + ), + pytest.param( + "messages-endpoint", + "/v1/messages", + "databricks-meta-llama-3-3-70b-instruct", + { + "max_tokens": 32, + "messages": [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}], + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A4-messages", + ), + pytest.param( + "responses-endpoint", + "/v1/responses", + "databricks-meta-llama-3-3-70b-instruct", + {"input": [{"role": "user", "content": [{"type": "input_text", "text": "Reply in JSON"}]}]}, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A5-responses-list-input", + ), + ], +) +def test_databricks_non_claude_single_text_block_reaches_upstream_as_string( + gateway: Gateway, + case: str, + endpoint: str, + provider_model: str, + request_body: dict[str, JsonValue], + expected_messages: list[JsonValue], + responses_mode: bool, + upstream_path: str, +) -> None: + identity: Final = f"databricks-single-text-{case}-{uuid.uuid4().hex}" + streaming: Final = request_body.get("stream") is True + + def respond(request: Request) -> Reply: + assert request.method == "POST" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["messages"] == expected_messages, body["messages"] + return _shape_reply(identity, streaming) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"databricks/{provider_model}", + api_base=wire.url, + api_key=_API_KEY, + model_info={"mode": "responses"} if responses_mode else None, + ) + headers: Final = {"anthropic-version": "2023-06-01"} if endpoint == "/v1/messages" else {} + body: Final = {"model": model, **request_body} + if streaming: + with gateway.client.stream( + "POST", endpoint, json=body, headers={"Authorization": f"Bearer {gateway.key}", **headers} + ) as response: + assert response.status_code == 200, response.read() + lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: ")) + assert lines[-1] == "data: [DONE]", lines + assert any('"content":"ok"' in line for line in lines), lines + else: + response: Final = gateway.request("POST", endpoint, body, headers=headers) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["id"] + assert [(request.method, request.target) for request in wire.drain()] == [("POST", upstream_path)] + + +@pytest.mark.parametrize( + "case,provider_model,content,expected_content", + [ + pytest.param( + "cache-control", + "databricks-meta-llama-3-3-70b-instruct", + [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], + [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], + id="B1-cache-control-stays-list", + ), + pytest.param( + "two-text-blocks", + "databricks-meta-llama-3-3-70b-instruct", + [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}], + [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}], + id="B2-two-text-blocks-stays-list", + ), + pytest.param( + "text-and-image", + "databricks-meta-llama-3-3-70b-instruct", + [ + {"type": "text", "text": "Reply in JSON"}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}}, + ], + [ + {"type": "text", "text": "Reply in JSON"}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}}, + ], + id="B3-text-and-image-stays-list", + ), + pytest.param( + "claude", + "databricks-claude-opus-5-5", + [{"type": "text", "text": "Reply in JSON"}], + [{"type": "text", "text": "Reply in JSON"}], + id="B4-claude-stays-list", + ), + ], +) +def test_databricks_list_content_stays_list( + gateway: Gateway, + case: str, + provider_model: str, + content: list[JsonValue], + expected_content: list[JsonValue], +) -> None: + identity: Final = f"databricks-list-content-{case}-{uuid.uuid4().hex}" + expected_messages: Final = [{"role": "user", "content": expected_content}] + + def respond(request: Request) -> Reply: + assert request.method == "POST" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["messages"] == expected_messages, body["messages"] + return _shape_reply(identity, False) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{provider_model}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": content}]}, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["id"] + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + + def test_databricks_json_schema_refs_are_preserved_in_chat_completions(gateway: Gateway) -> None: def respond(request: Request) -> Reply: assert request.method == "POST" @@ -868,3 +1092,356 @@ def test_databricks_json_schema_burst_through_an_upstream_outage_keeps_every_ref ) assert tuple(_outcome(*result) for result in results) == expected assert tuple(_outcome(*result) for result in recovered) == ((200, _EXPECTED_JSON),) * len(recovery) + + +_LLAMA_BACKEND: Final = "databricks-meta-llama-3-3-70b-instruct" +_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} + + +def _one_block(text: str) -> list[JsonValue]: + return [{"type": "text", "text": text}] + + +def _plain_reply(_request: Request) -> Reply: + return _shape_reply(f"databricks-plain-{uuid.uuid4().hex}", False) + + +def _sent_messages(request: Request) -> JsonValue: + return _JSON_OBJECT.validate_json(request.body)["messages"] + + +def _tool_call(identity: str, city: str) -> dict[str, JsonValue]: + return { + "id": identity, + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps({"city": city})}, + } + + +@pytest.mark.parametrize( + "messages,expected_messages", + [ + pytest.param( + [ + {"role": "system", "content": _one_block("Alpha rule.")}, + {"role": "system", "content": _one_block("Bravo rule.")}, + {"role": "user", "content": _one_block("Reply in JSON")}, + ], + [ + {"role": "system", "content": "Alpha rule.\n\nBravo rule."}, + {"role": "user", "content": "Reply in JSON"}, + ], + id="C1-two-one-block-system-messages", + ), + pytest.param( + [ + {"role": "system", "content": _one_block("Alpha rule.")}, + {"role": "developer", "content": _one_block("Bravo rule.")}, + {"role": "user", "content": "Reply in JSON"}, + ], + [ + {"role": "system", "content": "Alpha rule.\n\nBravo rule."}, + {"role": "user", "content": "Reply in JSON"}, + ], + id="C2-system-then-developer-one-block", + ), + pytest.param( + [ + {"role": "system", "content": _one_block("Alpha rule.")}, + {"role": "system", "content": "Bravo rule."}, + {"role": "user", "content": "Reply in JSON"}, + ], + [ + {"role": "system", "content": "Alpha rule.\n\nBravo rule."}, + {"role": "user", "content": "Reply in JSON"}, + ], + id="C3-one-block-then-string-system", + ), + pytest.param( + [ + {"role": "system", "content": _one_block("Alpha rule.")}, + { + "role": "system", + "content": [{"type": "text", "text": "Bravo", "cache_control": {"type": "ephemeral"}}], + }, + {"role": "user", "content": "Reply in JSON"}, + ], + [ + { + "role": "system", + "content": [ + {"type": "text", "text": "Alpha rule."}, + {"type": "text", "text": "Bravo", "cache_control": {"type": "ephemeral"}}, + ], + }, + {"role": "user", "content": "Reply in JSON"}, + ], + id="C4-one-block-then-cache-marked-system-stays-parts", + ), + pytest.param( + [ + {"role": "user", "content": _one_block("Weather in Paris and Rome?")}, + { + "role": "assistant", + "content": _one_block("Checking both."), + "tool_calls": [_tool_call("call-paris", "Paris"), _tool_call("call-rome", "Rome")], + }, + {"role": "tool", "tool_call_id": "call-paris", "content": _one_block("18C")}, + {"role": "tool", "tool_call_id": "call-rome", "content": _one_block("22C")}, + ], + [ + {"role": "user", "content": "Weather in Paris and Rome?"}, + {"role": "assistant", "content": "Checking both.", "tool_calls": [_tool_call("call-paris", "Paris")]}, + {"role": "tool", "tool_call_id": "call-paris", "content": "18C"}, + {"role": "assistant", "tool_calls": [_tool_call("call-rome", "Rome")]}, + {"role": "tool", "tool_call_id": "call-rome", "content": "22C"}, + ], + id="C5-parallel-tool-results-split-and-collapsed", + ), + pytest.param( + [ + {"role": "user", "content": [{"type": "text", "text": "Reply in JSON", "cache_control": None}]}, + {"role": "assistant", "content": [{"type": "text", "text": "ok", "annotations": []}]}, + {"role": "user", "content": _one_block(" ")}, + {"role": "user", "content": _one_block("x" * 5120)}, + ], + [ + {"role": "user", "content": [{"type": "text", "text": "Reply in JSON", "cache_control": None}]}, + {"role": "assistant", "content": [{"type": "text", "text": "ok", "annotations": []}]}, + {"role": "user"}, + {"role": "user", "content": "x" * 5120}, + ], + id="F1-extra-keys-stay-list-blank-dropped-5kb-collapsed", + ), + ], +) +def test_databricks_non_claude_message_shapes_reach_upstream( + gateway: Gateway, messages: list[JsonValue], expected_messages: list[JsonValue] +) -> None: + with wire_server(_plain_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_LLAMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": messages, "cache": _NO_CACHE}, + ) + received: Final = wire.drain() + assert response.status_code == 200, response.text + assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")] + assert _sent_messages(received[0]) == expected_messages + + +def test_databricks_responses_instructions_and_one_block_developer_reach_upstream_joined(gateway: Gateway) -> None: + with wire_server(_plain_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_LLAMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/responses", + { + "model": model, + "instructions": "Alpha rule.", + "input": [ + {"role": "developer", "content": [{"type": "input_text", "text": "Bravo rule."}]}, + {"role": "user", "content": [{"type": "input_text", "text": "Reply in JSON"}]}, + ], + }, + ) + received: Final = wire.drain() + assert response.status_code == 200, response.text + assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")] + assert _sent_messages(received[0]) == [ + {"role": "system", "content": "Alpha rule.\n\nBravo rule."}, + {"role": "user", "content": "Reply in JSON"}, + ] + + +def test_databricks_fallback_sends_the_collapsed_string_to_both_deployments(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + fallback_backend: Final = "databricks-meta-llama-3-1-8b-instruct" + + def respond(request: Request) -> Reply: + if _JSON_OBJECT.validate_json(request.body)["model"] == _LLAMA_BACKEND: + return Reply(status=500, body=b'{"error_code": "INTERNAL_ERROR", "message": "primary deployment is down"}') + return _shape_reply(f"databricks-fallback-{marker}", False) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + primary: Final = scenario.model(model=f"databricks/{_LLAMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + fallback: Final = scenario.model(model=f"databricks/{fallback_backend}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": primary, + "messages": [{"role": "user", "content": _one_block(f"Reply in JSON {marker}")}], + "response_format": {"type": "json_object"}, + "fallbacks": [fallback], + "num_retries": 0, + "cache": _NO_CACHE, + }, + ) + received: Final = wire.drain() + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["id"] == f"databricks-fallback-{marker}" + assert [_JSON_OBJECT.validate_json(request.body)["model"] for request in received] == [ + _LLAMA_BACKEND, + fallback_backend, + ] + assert [_sent_messages(request) for request in received] == [ + [{"role": "user", "content": f"Reply in JSON {marker}"}], + ] * 2 + + +def _proxy_root(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def test_databricks_openai_sdk_sync_one_block_reaches_upstream_as_string_and_logs_spend(gateway: Gateway) -> None: + identity: Final = f"databricks-openai-sdk-{uuid.uuid4().hex}" + with ( + wire_server(lambda _request: _shape_reply(identity, False)) as wire, + gateway.scenario() as scenario, + openai.OpenAI(base_url=f"{_proxy_root(gateway)}/v1", api_key=gateway.key, max_retries=0) as client, + ): + model: Final = scenario.model(model=f"databricks/{_LLAMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + completion: Final = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}], + response_format={"type": "json_object"}, + extra_body={"cache": _NO_CACHE}, + ) + received: Final = wire.drain() + assert (completion.id, completion.choices[0].message.content) == (identity, "ok") + assert [_sent_messages(request) for request in received] == [[{"role": "user", "content": "Reply in JSON"}]] + rows: Final = eventually( + lambda: read_rows('SELECT status, call_type FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), + lambda values: len(values) == 1, + seconds=70, + ) + assert (rows[0]["status"], rows[0]["call_type"]) == ("success", "acompletion") + + +async def test_databricks_openai_sdk_async_stream_one_block_reaches_upstream_as_string(gateway: Gateway) -> None: + identity: Final = f"databricks-openai-async-{uuid.uuid4().hex}" + with wire_server(lambda _request: _shape_reply(identity, True)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_LLAMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + async with openai.AsyncOpenAI( + base_url=f"{_proxy_root(gateway)}/v1", api_key=gateway.key, max_retries=0 + ) as client: + stream: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}], + response_format={"type": "json_object"}, + stream=True, + extra_body={"cache": _NO_CACHE}, + ) + chunks: Final = [chunk async for chunk in stream] + received: Final = wire.drain() + assert {chunk.id for chunk in chunks} == {identity} + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "ok" + assert [_sent_messages(request) for request in received] == [[{"role": "user", "content": "Reply in JSON"}]] + + +async def test_databricks_anthropic_sdk_async_stream_one_block_system_and_user_reach_upstream_as_strings( + gateway: Gateway, +) -> None: + identity: Final = f"databricks-anthropic-async-{uuid.uuid4().hex}" + with wire_server(lambda _request: _shape_reply(identity, True)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_LLAMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + async with anthropic.AsyncAnthropic( + base_url=_proxy_root(gateway), api_key=gateway.key, max_retries=0 + ) as client: + async with client.messages.stream( + model=model, + max_tokens=32, + system=[{"type": "text", "text": "Reply in JSON."}], + messages=[{"role": "user", "content": [{"type": "text", "text": "What is 2+2?"}]}], + ) as stream: + final: Final = await stream.get_final_message() + received: Final = wire.drain() + assert [block.text for block in final.content if block.type == "text"] == ["ok"] + assert [_sent_messages(request) for request in received] == [ + [{"role": "system", "content": "Reply in JSON."}, {"role": "user", "content": "What is 2+2?"}] + ] + + +def test_databricks_one_block_burst_through_an_upstream_outage_reaches_upstream_as_strings(gateway: Gateway) -> None: + plan: Final = tuple(_burst_request(index) for index in range(24)) + + def respond(request: Request) -> Reply: + if _OUTAGE.encode() in request.body: + error: Final = {"error_code": "TEMPORARILY_UNAVAILABLE", "message": _OUTAGE_MESSAGE} + return Reply(status=503, body=json.dumps(error).encode()) + body: Final = _JSON_OBJECT.validate_json(request.body) + return _upstream_reply(_JSON_SCHEMA_BACKEND, body.get("stream") is True) + + def one_block_body(endpoint: _Endpoint, model: str, stream: bool, prompt: str) -> dict[str, JsonValue]: + body: Final = _client_body(endpoint, model, stream, prompt) + match endpoint: + case "responses": + return {**body, "input": [{"role": "user", "content": [{"type": "input_text", "text": prompt}]}]} + case "chat" | "messages": + return {**body, "messages": [{"role": "user", "content": _one_block(prompt)}]} + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + prompts: Final = tuple( + f"{_JSON_SCHEMA_PROMPT} {index} {_OUTAGE if down else ''}".strip() + for index, (_e, _s, down) in enumerate(plan) + ) + bodies: Final = tuple( + one_block_body(endpoint, model, stream, prompt) for (endpoint, stream, _down), prompt in zip(plan, prompts) + ) + with ThreadPoolExecutor(max_workers=12) as pool: + results: Final = tuple(pool.map(_call, itertools.repeat(gateway), (e for e, _s, _d in plan), bodies)) + received: Final = wire.drain() + assert sorted(json.dumps(_sent_messages(request)) for request in received) == sorted( + json.dumps([{"role": "user", "content": prompt}]) for prompt in prompts + ) + expected: Final = tuple( + (503, _OUTAGE_MESSAGE) if down else (200, _EXPECTED_JSON) for _endpoint, _stream, down in plan + ) + assert tuple(_outcome(*result) for result in results) == expected + + +def test_databricks_transform_request_shows_the_collapsed_string(gateway: Gateway, tmp_path: Path) -> None: + environment: Final = { + "DATABRICKS_API_BASE": "https://databricks.invalid/serving-endpoints", + "DATABRICKS_API_KEY": _API_KEY, + } + with owned_proxy_process(gateway, tmp_path, environment) as owned: + response: Final = owned.gateway.request( + "POST", + "/utils/transform_request", + { + "call_type": "completion", + "request_body": { + "model": f"databricks/{_LLAMA_BACKEND}", + "messages": [{"role": "user", "content": _one_block("Reply in JSON")}], + }, + }, + ) + assert response.status_code == 200, response.text + body: Final = _JSON_OBJECT.validate_json(response.content) + assert body["error"] is None, body + assert body["raw_request_api_base"] == "https://databricks.invalid/serving-endpoints/chat/completions" + sent: Final = body["raw_request_body"] + assert isinstance(sent, dict), body + assert sent["messages"] == [{"role": "user", "content": "Reply in JSON"}] + + +def test_databricks_non_string_text_block_is_rejected_before_reaching_upstream(gateway: Gateway) -> None: + with wire_server(_plain_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"databricks/{_LLAMA_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": [{"type": "text", "text": 42}]}], + "cache": _NO_CACHE, + }, + ) + received: Final = wire.drain() + assert response.status_code == 400, response.text + assert "'int' object has no attribute 'strip'" in response.text + assert received == () 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 fe76a667f3e..d25787d7028 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest import respx +from pydantic import TypeAdapter from fastapi.testclient import TestClient from pydantic import BaseModel, ConfigDict @@ -23,6 +24,7 @@ from litellm.llms.databricks.chat.transformation import ( DatabricksConfig, _sanitize_empty_content, ) +from litellm.types.llms.openai import AllMessageValues import asyncio from unittest.mock import Mock from litellm._version import version @@ -432,6 +434,119 @@ def test_transform_request_preserves_unity_model_service_name(): assert result["model"] == "system.ai.kimi-k3" +@pytest.mark.parametrize( + "model,messages,optional_params,expected_content", + [ + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}], + {"response_format": {"type": "json_object"}}, + "Reply in JSON", + id="user", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "system", "content": [{"type": "text", "text": "Reply in JSON"}]}], + {"response_format": {"type": "json_object"}}, + "Reply in JSON", + id="system", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "assistant", "content": [{"type": "text", "text": "Reply in JSON"}]}], + {"response_format": {"type": "json_object"}}, + "Reply in JSON", + id="assistant", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [ + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "text", "text": "Reply in JSON"}], + } + ], + {"response_format": {"type": "json_object"}}, + "Reply in JSON", + id="tool", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [ + { + "role": "user", + "content": [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], + } + ], + {"response_format": {"type": "json_object"}}, + [{"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", + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Reply in JSON"}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}}, + ], + } + ], + {"response_format": {"type": "json_object"}}, + [ + {"type": "text", "text": "Reply in JSON"}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}}, + ], + id="text-and-image", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": "Reply in JSON"}], + {"response_format": {"type": "json_object"}}, + "Reply in JSON", + id="plain-string", + ), + pytest.param( + "databricks-claude-sonnet-5", + [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}], + {"response_format": {"type": "json_object"}}, + [{"type": "text", "text": "Reply in JSON"}], + id="claude", + ), + ], +) +def test_transform_request_collapses_single_plain_text_block_for_non_claude_models( + model: str, + messages: list[AllMessageValues], + optional_params: dict[str, object], + expected_content: object, +) -> None: + config: Final[DatabricksConfig] = DatabricksConfig() + result: Final[dict[str, object]] = TypeAdapter(dict[str, object]).validate_python( + config.transform_request( # pyright: ignore[reportUnknownMemberType] # inherited hook is untyped + model=model, + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + ) + outgoing_messages: Final[list[dict[str, object]]] = TypeAdapter(list[dict[str, object]]).validate_python( + result["messages"] + ) + + assert outgoing_messages == [{**messages[0], "content": expected_content}] + + 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 @@ -2117,6 +2232,7 @@ def mock_chat_streaming_response_chunks() -> List[str]: "databricks-meta-llama-3-3-70b-instruct", "databricks-qwen35-122b-a10b", "databricks-gpt-oss-120b", + "system.ai.deepseek-v4-1-flash", ), ) def test_databricks_non_claude_json_schema_preserves_refs(model: str) -> None: diff --git a/tests/unit/responses/test_responses_api_bridge_flag.py b/tests/unit/responses/test_responses_api_bridge_flag.py index fb1361c1f49..699144a31bb 100644 --- a/tests/unit/responses/test_responses_api_bridge_flag.py +++ b/tests/unit/responses/test_responses_api_bridge_flag.py @@ -356,11 +356,8 @@ class TestUseResponsesApiBridgeFlag: assert upstream.call_count == 1 request_body: Final = json.loads(upstream.calls[0].request.read()) assert request_body["messages"] == [ - { - "role": "system", - "content": [{"type": "text", "text": "You are terse."}, {"type": "text", "text": "Skills: none."}], - }, - {"role": "user", "content": [{"type": "text", "text": "Hello"}]}, + {"role": "system", "content": "You are terse.\n\nSkills: none."}, + {"role": "user", "content": "Hello"}, ] assert "client_metadata" not in request_body assert request_body["chat_template_kwargs"] == {"thinking": True}