mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(databricks): join only the last message for json_object and drop cache_control under drop_params
Databricks checks for "json" only in the last message, so earlier turns now go out as sent and only the last message, plus any system messages merged into it, is joined into a string. Under drop_params a text block's cache_control is dropped so that message can be joined too, and without it the request still goes out as a list.
This commit is contained in:
mateo-berri
parent
f3d4419379
commit
f3d459e844
3 changed files with 171 additions and 18 deletions
|
|
@ -4,12 +4,14 @@ Translates from OpenAI's `/v1/chat/completions` to Databricks' `/chat/completion
|
|||
|
||||
import os
|
||||
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
|
||||
from itertools import takewhile
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
|
||||
handle_invalid_parallel_tool_calls,
|
||||
|
|
@ -109,22 +111,63 @@ def _collapse_single_text_block(message: Mapping[str, object]) -> Mapping[str, o
|
|||
return message
|
||||
|
||||
|
||||
def _join_plain_text_blocks(message: AllMessageValues) -> AllMessageValues:
|
||||
match _plain_texts(message):
|
||||
def _without_text_cache_control(block: object) -> object:
|
||||
match block:
|
||||
case {"type": "text", "text": str() as text, "cache_control": object() as cache_control} if block == {
|
||||
"type": "text",
|
||||
"text": text,
|
||||
"cache_control": cache_control,
|
||||
}:
|
||||
return {"type": "text", "text": text}
|
||||
case _:
|
||||
return block
|
||||
|
||||
|
||||
def _without_cache_control(message: AllMessageValues) -> AllMessageValues:
|
||||
uncached: Final = {key: value for key, value in message.items() if key != "cache_control"}
|
||||
match content := uncached.get("content"):
|
||||
case list():
|
||||
blocks: Final = [_without_text_cache_control(block) for block in _CONTENT_BLOCKS.validate_python(content)]
|
||||
return cast(AllMessageValues, {**uncached, "content": blocks}) # cast-ok: spread keeps the shape
|
||||
case _:
|
||||
return cast(AllMessageValues, uncached) # cast-ok: only cache_control was left out
|
||||
|
||||
|
||||
def _has_string_content(message: Mapping[str, object]) -> bool:
|
||||
return isinstance(message.get("content"), str)
|
||||
|
||||
|
||||
def _is_system(message: Mapping[str, object]) -> bool:
|
||||
return message.get("role") == "system"
|
||||
|
||||
|
||||
def _json_mode_message(message: AllMessageValues, drop_params: bool) -> AllMessageValues:
|
||||
candidate: Final = _without_cache_control(message) if drop_params else message
|
||||
match _plain_texts(candidate):
|
||||
case ():
|
||||
return message
|
||||
return candidate if _has_string_content(candidate) else message
|
||||
case texts:
|
||||
return cast(AllMessageValues, {**message, "content": "".join(texts)}) # cast-ok: spread keeps the shape
|
||||
return cast(AllMessageValues, {**candidate, "content": "".join(texts)}) # cast-ok: spread keeps the shape
|
||||
|
||||
|
||||
def _json_checked_start(messages: Sequence[AllMessageValues]) -> int:
|
||||
trailing_system: Final = sum(1 for _ in takewhile(_is_system, reversed(messages)))
|
||||
return max(len(messages) - max(trailing_system, 1), 0)
|
||||
|
||||
|
||||
def _json_object_ready_messages(
|
||||
model: str,
|
||||
messages: list[AllMessageValues], # mutable-ok: inherited transform_request contract
|
||||
response_format: object,
|
||||
drop_params: bool,
|
||||
) -> list[AllMessageValues]: # mutable-ok: inherited transform_request contract
|
||||
match response_format:
|
||||
case {"type": "json_object"} if "claude" not in model:
|
||||
return [_join_plain_text_blocks(message) for message in messages]
|
||||
checked_start: Final = _json_checked_start(messages)
|
||||
return [
|
||||
*messages[:checked_start],
|
||||
*(_json_mode_message(message, drop_params) for message in messages[checked_start:]),
|
||||
]
|
||||
case _:
|
||||
return messages
|
||||
|
||||
|
|
@ -586,7 +629,12 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
) -> dict[str, object]: # mutable-ok: inherited contract
|
||||
return super().transform_request(
|
||||
model=model,
|
||||
messages=_json_object_ready_messages(model, messages, optional_params.get("response_format")),
|
||||
messages=_json_object_ready_messages(
|
||||
model,
|
||||
messages,
|
||||
optional_params.get("response_format"),
|
||||
drop_params=litellm.drop_params is True or litellm_params.get("drop_params") is True,
|
||||
),
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
|
|
|
|||
|
|
@ -385,21 +385,64 @@ def _shape_reply(identity: str, stream: bool) -> Reply:
|
|||
id="A7-chat-stream-json-object-two-blocks-joined",
|
||||
),
|
||||
pytest.param(
|
||||
"json-object-system-blocks",
|
||||
"json-object-history",
|
||||
"/v1/chat/completions",
|
||||
"databricks-meta-llama-3-3-70b-instruct",
|
||||
{
|
||||
"messages": [
|
||||
{"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]},
|
||||
{"role": "system", "content": "Bravo rule."},
|
||||
{"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]},
|
||||
{"role": "assistant", "content": "Got it."},
|
||||
{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]},
|
||||
],
|
||||
"response_format": {"type": "json_object"},
|
||||
},
|
||||
[{"role": "system", "content": "Alpha rule.\n\nBravo rule."}, {"role": "user", "content": "Reply in JSON"}],
|
||||
[
|
||||
{"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]},
|
||||
{"role": "assistant", "content": "Got it."},
|
||||
{"role": "user", "content": "Reply in JSON"},
|
||||
],
|
||||
False,
|
||||
"/chat/completions",
|
||||
id="A8-chat-json-object-system-blocks-keep-break",
|
||||
id="A8-chat-json-object-earlier-turns-kept",
|
||||
),
|
||||
pytest.param(
|
||||
"json-object-trailing-system",
|
||||
"/v1/chat/completions",
|
||||
"databricks-meta-llama-3-3-70b-instruct",
|
||||
{
|
||||
"messages": [
|
||||
{"role": "user", "content": "Say hi"},
|
||||
{"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]},
|
||||
{"role": "system", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]},
|
||||
],
|
||||
"response_format": {"type": "json_object"},
|
||||
},
|
||||
[{"role": "user", "content": "Say hi"}, {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}],
|
||||
False,
|
||||
"/chat/completions",
|
||||
id="A8b-chat-json-object-trailing-system-keeps-break",
|
||||
),
|
||||
pytest.param(
|
||||
"json-object-drop-params-cache-control",
|
||||
"/v1/chat/completions",
|
||||
"databricks-meta-llama-3-3-70b-instruct",
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Reply"},
|
||||
{"type": "text", "text": " in JSON", "cache_control": {"type": "ephemeral"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
"response_format": {"type": "json_object"},
|
||||
"drop_params": True,
|
||||
},
|
||||
[{"role": "user", "content": "Reply in JSON"}],
|
||||
False,
|
||||
"/chat/completions",
|
||||
id="A8c-chat-json-object-drop-params-drops-cache-control",
|
||||
),
|
||||
pytest.param(
|
||||
"json-object-responses-two-parts",
|
||||
|
|
|
|||
|
|
@ -603,14 +603,14 @@ _CACHED_BLOCK: Final = {"type": "text", "text": " in JSON", "cache_control": {"t
|
|||
|
||||
|
||||
def _outgoing_messages(
|
||||
model: str, messages: list[AllMessageValues], optional_params: dict[str, object]
|
||||
model: str, messages: list[AllMessageValues], optional_params: dict[str, object], drop_params: bool = False
|
||||
) -> list[dict[str, object]]:
|
||||
request: Final = TypeAdapter(dict[str, object]).validate_python(
|
||||
DatabricksConfig().transform_request( # pyright: ignore[reportUnknownMemberType] # inherited hook is untyped
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
litellm_params={"drop_params": True} if drop_params else {},
|
||||
headers={},
|
||||
)
|
||||
)
|
||||
|
|
@ -655,6 +655,13 @@ def _outgoing_messages(
|
|||
[{"type": "text", "text": "Reply"}, _CACHED_BLOCK],
|
||||
id="json-object-cache-control-stays-list",
|
||||
),
|
||||
pytest.param(
|
||||
"databricks-meta-llama-3-3-70b-instruct",
|
||||
[{"role": "user", "content": "Reply in JSON", "cache_control": {"type": "ephemeral"}}],
|
||||
_JSON_OBJECT,
|
||||
[{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}],
|
||||
id="json-object-message-cache-control-stays-block",
|
||||
),
|
||||
pytest.param(
|
||||
"databricks-meta-llama-3-3-70b-instruct",
|
||||
[{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}, _IMAGE_BLOCK]}],
|
||||
|
|
@ -680,19 +687,74 @@ def test_transform_request_joins_plain_text_blocks_only_for_non_claude_json_obje
|
|||
assert _outgoing_messages(model, messages, optional_params) == [{"role": "user", "content": expected_content}]
|
||||
|
||||
|
||||
def test_transform_request_json_object_join_keeps_the_break_between_system_messages() -> None:
|
||||
def test_transform_request_json_object_keeps_earlier_turns_as_sent() -> None:
|
||||
earlier_turn: Final[AllMessageValues] = {
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}],
|
||||
}
|
||||
reply: Final[AllMessageValues] = {"role": "assistant", "content": "Got it."}
|
||||
|
||||
assert _outgoing_messages(
|
||||
"databricks-meta-llama-3-3-70b-instruct",
|
||||
[earlier_turn, reply, {"role": "user", "content": _TWO_TEXT_BLOCKS}],
|
||||
_JSON_OBJECT,
|
||||
) == [earlier_turn, reply, {"role": "user", "content": "Reply in JSON"}]
|
||||
|
||||
|
||||
def test_transform_request_json_object_joins_a_trailing_system_run_and_keeps_its_break() -> None:
|
||||
messages: Final[list[AllMessageValues]] = [
|
||||
{"role": "user", "content": "Say hi"},
|
||||
{"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]},
|
||||
{"role": "system", "content": "Bravo rule."},
|
||||
{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]},
|
||||
{"role": "system", "content": _TWO_TEXT_BLOCKS},
|
||||
]
|
||||
|
||||
assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", messages, _JSON_OBJECT) == [
|
||||
{"role": "system", "content": "Alpha rule.\n\nBravo rule."},
|
||||
{"role": "user", "content": "Reply in JSON"},
|
||||
{"role": "user", "content": "Say hi"},
|
||||
{"role": "system", "content": "Alpha rule.\n\nReply in JSON"},
|
||||
]
|
||||
|
||||
|
||||
_CACHED_TEXT_AND_IMAGE: Final = [{**_CACHED_BLOCK, "text": "Reply in JSON"}, _IMAGE_BLOCK]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"message,expected",
|
||||
[
|
||||
pytest.param(
|
||||
{"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]},
|
||||
{"role": "user", "content": "Reply in JSON"},
|
||||
id="block-cache-control-dropped-and-joined",
|
||||
),
|
||||
pytest.param(
|
||||
{"role": "user", "content": "Reply in JSON", "cache_control": {"type": "ephemeral"}},
|
||||
{"role": "user", "content": "Reply in JSON"},
|
||||
id="message-cache-control-dropped",
|
||||
),
|
||||
pytest.param(
|
||||
{"role": "user", "content": _CACHED_TEXT_AND_IMAGE},
|
||||
{"role": "user", "content": _CACHED_TEXT_AND_IMAGE},
|
||||
id="image-keeps-cache-control",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_transform_request_json_object_drops_last_message_cache_control_under_drop_params(
|
||||
message: AllMessageValues, expected: dict[str, object]
|
||||
) -> None:
|
||||
assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", [message], _JSON_OBJECT, drop_params=True) == [
|
||||
expected
|
||||
]
|
||||
|
||||
|
||||
def test_transform_request_json_object_reads_the_global_drop_params(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", True)
|
||||
|
||||
assert _outgoing_messages(
|
||||
"databricks-meta-llama-3-3-70b-instruct",
|
||||
[{"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}],
|
||||
_JSON_OBJECT,
|
||||
) == [{"role": "user", "content": "Reply in JSON"}]
|
||||
|
||||
|
||||
def test_transform_request_strips_thinking_blocks_and_reasoning_content():
|
||||
"""Regression for LIT-6762: replaying an assistant turn that litellm decorated with
|
||||
`thinking_blocks` / `reasoning_content` made Databricks 400 with
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue