diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 60081c38073..53674907341 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -4,6 +4,7 @@ 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 chain, takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload @@ -12,6 +13,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter 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.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( handle_invalid_parallel_tool_calls, should_convert_tool_call_to_json_mode, @@ -65,7 +67,7 @@ from ...anthropic.chat.transformation import ( tools_without_unsupported_tool_search, ) from ...openai_like.chat.transformation import OpenAILikeChatConfig -from ..common_utils import DatabricksBase, DatabricksException +from ..common_utils import DATABRICKS_HOSTED_MODEL_PREFIXES, DatabricksBase, DatabricksException _RESPONSE_FORMAT_ADAPTER: ( # mutable-ok: mirrors the dict return contract of get_json_schema_from_pydantic_object Final[TypeAdapter[dict[str, object] | None]] @@ -81,15 +83,139 @@ def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool: ) +_CONTENT_BLOCKS: Final = TypeAdapter(tuple[object, ...]) + + +def _content_blocks(message: object) -> tuple[object, ...]: + match message.model_dump(exclude_none=True) if isinstance(message, BaseModel) else message: + case {"content": list() as content}: + return _CONTENT_BLOCKS.validate_python(content) + case _: + return () + + +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, ...]: + blocks: Final = _content_blocks(message) + texts: Final = tuple(text for block in blocks if (text := _plain_text(block)) is not None) + return texts if len(texts) == len(blocks) else () + + 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 and model.startswith(DATABRICKS_HOSTED_MODEL_PREFIXES): + 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 _non_string_text_type(block: object) -> str | None: + match block: + case {"type": "text", "text": str()}: + return None + case {"type": "text", "text": object() as text} if text: + return type(text).__name__ + case _: + return None + + +def _first_non_string_text_type(messages: Sequence[AllMessageValues]) -> str | None: + blocks: Final = chain.from_iterable(map(_content_blocks, messages)) + return next((type_name for block in blocks if (type_name := _non_string_text_type(block))), None) + + def _sanitize_empty_content(message_dict: dict[str, object]) -> None: """ Remove or filter content so empty text blocks are not sent. @@ -108,9 +234,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 +659,35 @@ 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 + match _first_non_string_text_type(messages): + case str() as type_name: + raise DatabricksException( + status_code=400, + message=f"A text content block needs a string 'text', got {type_name}", + ) + case None: + pass + return super().transform_request( + model=model, + messages=_json_object_ready_messages( + model, + messages, + optional_params.get("response_format"), + drop_params=drops_unsupported_params(litellm_params.get("drop_params")), + ), + 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/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index b1562bef6b8..e8313d9a003 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -56,6 +56,7 @@ DATABRICKS_WORKSPACE_ROUTE_PREFIXES: Final = ( DATABRICKS_AI_GATEWAY_MESSAGES_PATH, ) DATABRICKS_NATIVE_MESSAGES_MODEL_PREFIXES: Final = ("databricks-claude-", "system.ai.claude-") +DATABRICKS_HOSTED_MODEL_PREFIXES: Final = ("databricks-", "system.ai.") DATABRICKS_ENDPOINT_PATHS: Final = MappingProxyType( { "chat_completions": "/chat/completions", diff --git a/tests/integration/providers/test_databricks_chat_wire.py b/tests/integration/providers/test_databricks_chat_wire.py index b93cebdbc4c..3d9d749eb04 100644 --- a/tests/integration/providers/test_databricks_chat_wire.py +++ b/tests/integration/providers/test_databricks_chat_wire.py @@ -353,9 +353,172 @@ 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-external-endpoint", + "/v1/chat/completions", + "openai-gpt-external", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}], + False, + "/chat/completions", + id="A8e-chat-json-object-external-endpoint-kept-as-sent", + ), + 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 +528,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: @@ -1453,7 +1616,7 @@ def test_databricks_non_string_text_block_is_rejected_before_reaching_upstream(g ) received: Final = wire.drain() assert response.status_code == 400, response.text - assert "'int' object has no attribute 'strip'" in response.text + assert "A text content block needs a string 'text', got int" 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 7a6f9c7dc18..33c1e7a2e22 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -1,14 +1,14 @@ import copy import json from collections.abc import AsyncIterator, Iterator -from typing import Final +from typing import Final, cast from unittest.mock import MagicMock, patch import httpx import pytest import respx from openai.lib.streaming.chat import ChatCompletionStreamState -from openai.types.chat import ChatCompletionChunk +from openai.types.chat import ChatCompletionChunk, ChatCompletionMessage from pydantic import TypeAdapter from fastapi.testclient import TestClient from pydantic import BaseModel, ConfigDict @@ -21,6 +21,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.databricks.common_utils import DatabricksException from litellm.llms.databricks.chat.transformation import ( DatabricksChatResponseIterator, DatabricksConfig, @@ -539,13 +540,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 +597,283 @@ 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", + ), + pytest.param( + "system.ai.deepseek-v4-flash-0731", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + "Reply in JSON", + id="json-object-ai-gateway-model-joined", + ), + pytest.param( + "openai-gpt-external", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + _TWO_TEXT_BLOCKS, + id="json-object-external-endpoint-stays-list", + ), + ], +) +def test_transform_request_joins_plain_text_blocks_only_for_hosted_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}] + + +@pytest.mark.parametrize("response_format", [None, {"type": "json_object"}], ids=["no-response-format", "json-object"]) +def test_transform_request_rejects_a_text_block_whose_text_is_a_truthy_non_string( + response_format: dict[str, str] | None, +) -> None: + with pytest.raises(DatabricksException) as error: + _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}, {"type": "text", "text": 7}]}], + {} if response_format is None else {"response_format": response_format}, + ) + + assert error.value.status_code == 400 + + +def test_transform_request_accepts_an_openai_message_object_in_history() -> None: + earlier_answer: Final = ChatCompletionMessage(role="assistant", content="Earlier answer") + + outgoing: Final = _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [ + {"role": "user", "content": "Earlier question"}, + cast(AllMessageValues, earlier_answer), # cast-ok: SDK callers pass response messages back as history + {"role": "user", "content": "Next question"}, + ], + {}, + ) + + assert [message.get("content") for message in outgoing] == ["Earlier question", "Earlier answer", "Next question"] + + +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