From b4d921529b77f23469c45d84cc7e0d9efdf61d45 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 29 Jun 2026 10:43:34 -0700 Subject: [PATCH 1/2] fix(databricks): split parallel tool calls so each tool message follows tool_calls Databricks OpenAI-compatible serving (e.g. GPT models) 400s with "messages with role 'tool' must be a response to a preceeding message with 'tool_calls'" when an assistant turn makes parallel tool calls. LiteLLM faithfully sends one assistant message holding all tool_calls followed by one 'tool' message per result, so every result after the first is preceded by another 'tool' message rather than the assistant tool_calls message, which Databricks rejects. Re-emit each result immediately after an assistant message that carries only its matching tool_call, turning assistant(tool_calls=[A, B]), tool(A), tool(B) into assistant(tool_calls=[A]), tool(A), assistant(tool_calls=[B]), tool(B). The rewrite is a no-op when the turn is already valid (single call), the group is incomplete, or ids don't line up, so no tool call is ever dropped. Scoped to non-Claude models, matching the existing OpenAI-shaped transformation path. --- .../llms/databricks/chat/transformation.py | 58 ++++++ .../test_databricks_chat_transformation.py | 168 +++++++++++++++++- 2 files changed, 220 insertions(+), 6 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index a2c1d41022be..dff86fd67c5c 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -40,10 +40,13 @@ ) from litellm.types.llms.openai import ( AllMessageValues, + ChatCompletionAssistantMessage, + ChatCompletionAssistantToolCall, ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock, ChatCompletionToolChoiceFunctionParam, ChatCompletionToolChoiceObjectParam, + ChatCompletionToolMessage, ChatCompletionToolParam, ) from litellm.types.utils import ( @@ -92,6 +95,58 @@ def _sanitize_empty_content(message_dict: dict[str, Any]) -> None: message_dict["content"] = filtered +def _split_parallel_tool_calls(messages: List[AllMessageValues]) -> List[AllMessageValues]: + """ + Databricks (OpenAI-compatible serving) rejects a ``tool`` message unless the + message immediately before it carries ``tool_calls``. A single assistant turn + with parallel tool calls is followed by one ``tool`` message per call, so every + result after the first is preceded by another ``tool`` message and 400s. Re-emit + each result right after an assistant message holding only its matching call: + ``assistant(tool_calls=[A, B]), tool(A), tool(B)`` becomes + ``assistant(tool_calls=[A]), tool(A), assistant(tool_calls=[B]), tool(B)``. + + Left untouched (no-op) when the turn is already valid or the history is + malformed, so no tool call is ever dropped. + """ + + def _expand( + assistant: ChatCompletionAssistantMessage, + calls_by_id: dict[Optional[str], ChatCompletionAssistantToolCall], + tool_messages: List[ChatCompletionToolMessage], + ) -> Iterator[AllMessageValues]: + for position, tool_message in enumerate(tool_messages): + matched_call = calls_by_id[tool_message["tool_call_id"]] + if position == 0: + yield cast(AllMessageValues, {**assistant, "tool_calls": [matched_call]}) + else: + yield ChatCompletionAssistantMessage(role="assistant", tool_calls=[matched_call]) + yield tool_message + + def _generate() -> Iterator[AllMessageValues]: + index = 0 + while index < len(messages): + message = messages[index] + tool_calls = message.get("tool_calls") if message["role"] == "assistant" else None + if not tool_calls or len(tool_calls) < 2: + yield message + index += 1 + continue + end = index + 1 + while end < len(messages) and messages[end]["role"] == "tool": + end += 1 + tool_messages = cast(List[ChatCompletionToolMessage], messages[index + 1 : end]) + calls_by_id = {call["id"]: call for call in tool_calls} + result_ids = {tool_message["tool_call_id"] for tool_message in tool_messages} + if len(tool_messages) == len(tool_calls) and set(calls_by_id) == result_ids: + yield from _expand(cast(ChatCompletionAssistantMessage, message), calls_by_id, tool_messages) + index = end + else: + yield message + index += 1 + + return list(_generate()) + + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -385,6 +440,9 @@ def _transform_messages( _sanitize_empty_content(cast(dict[str, Any], _message)) new_messages.append(_message) + if "claude" not in model: + new_messages = _split_parallel_tool_calls(cast(List[AllMessageValues], new_messages)) + if is_async: return super()._transform_messages(messages=new_messages, model=model, is_async=cast(Literal[True], True)) else: diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 79e354d86217..cfdb76a97f46 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -5,9 +5,7 @@ import pytest from fastapi.testclient import TestClient -sys.path.insert( - 0, os.path.abspath("../../../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path from unittest.mock import MagicMock, patch from litellm.llms.databricks.chat.transformation import ( @@ -255,8 +253,166 @@ def test_transform_messages_sanitizes_empty_content(): {"role": "user", "content": [{"type": "text", "text": ""}]}, {"role": "user", "content": "Hi"}, ] - result = config._transform_messages( - messages=messages, model="databricks-claude", is_async=False - ) + result = config._transform_messages(messages=messages, model="databricks-claude", is_async=False) assert "content" not in result[0] assert result[1]["content"] == "Hi" + + +def _parallel_tool_calls(): + return [ + { + "id": "call_A", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "SF"}'}, + }, + { + "id": "call_B", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "NYC"}'}, + }, + ] + + +def _assert_every_tool_message_follows_tool_calls(messages): + for index, message in enumerate(messages): + if message.get("role") == "tool": + previous = messages[index - 1] if index > 0 else {} + assert previous.get("role") == "assistant" and previous.get("tool_calls"), ( + f"tool message at index {index} is not preceded by an assistant message with tool_calls: {messages}" + ) + + +def _declared_tool_call_ids(messages): + return sorted( + call["id"] + for message in messages + if message.get("role") == "assistant" and message.get("tool_calls") + for call in message["tool_calls"] + ) + + +def test_transform_request_splits_parallel_tool_calls_for_gpt(): + """Regression for LIT-3984: Databricks 400s with 'messages with role tool must + be a response to a preceeding message with tool_calls' because parallel tool + calls send consecutive tool messages. Each result must be re-paired with an + assistant tool_calls message holding only its matching call.""" + config = DatabricksConfig() + messages = [ + {"role": "user", "content": "weather in SF and NYC?"}, + {"role": "assistant", "content": "checking", "tool_calls": _parallel_tool_calls()}, + {"role": "tool", "tool_call_id": "call_A", "content": "sunny"}, + {"role": "tool", "tool_call_id": "call_B", "content": "rainy"}, + ] + + result = config.transform_request( + model="gpt-5.4-mini", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + )["messages"] + + _assert_every_tool_message_follows_tool_calls(result) + assert _declared_tool_call_ids(result) == ["call_A", "call_B"] + assistant_tool_call_messages = [m for m in result if m.get("role") == "assistant" and m.get("tool_calls")] + assert all(len(m["tool_calls"]) == 1 for m in assistant_tool_call_messages), ( + "each split assistant message must declare exactly one tool call" + ) + tool_messages = [m for m in result if m.get("role") == "tool"] + assert [m["tool_call_id"] for m in tool_messages] == ["call_A", "call_B"] + for tool_message, assistant_message in zip(tool_messages, assistant_tool_call_messages): + assert assistant_message["tool_calls"][0]["id"] == tool_message["tool_call_id"] + + +def test_transform_request_pairs_out_of_order_parallel_results(): + config = DatabricksConfig() + messages = [ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "content": "checking", "tool_calls": _parallel_tool_calls()}, + {"role": "tool", "tool_call_id": "call_B", "content": "rainy"}, + {"role": "tool", "tool_call_id": "call_A", "content": "sunny"}, + ] + + result = config.transform_request( + model="gpt-5.4-mini", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + )["messages"] + + _assert_every_tool_message_follows_tool_calls(result) + for index, message in enumerate(result): + if message.get("role") == "tool": + assert result[index - 1]["tool_calls"][0]["id"] == message["tool_call_id"] + + +def test_transform_request_leaves_single_tool_call_untouched(): + config = DatabricksConfig() + messages = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_A", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_A", "content": "sunny"}, + ] + + result = config.transform_request( + model="gpt-5.4-mini", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + )["messages"] + + assert len(result) == 3 + _assert_every_tool_message_follows_tool_calls(result) + assert _declared_tool_call_ids(result) == ["call_A"] + + +def test_transform_request_does_not_drop_tool_calls_on_incomplete_results(): + config = DatabricksConfig() + messages = [ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "content": "checking", "tool_calls": _parallel_tool_calls()}, + {"role": "tool", "tool_call_id": "call_A", "content": "sunny"}, + {"role": "user", "content": "thanks"}, + ] + + result = config.transform_request( + model="gpt-5.4-mini", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + )["messages"] + + assert _declared_tool_call_ids(result) == ["call_A", "call_B"] + + +def test_transform_request_keeps_parallel_tool_calls_for_claude(): + config = DatabricksConfig() + messages = [ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "content": "checking", "tool_calls": _parallel_tool_calls()}, + {"role": "tool", "tool_call_id": "call_A", "content": "sunny"}, + {"role": "tool", "tool_call_id": "call_B", "content": "rainy"}, + ] + + result = config.transform_request( + model="databricks-claude-3-7-sonnet", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + )["messages"] + + assert len([m for m in result if m.get("role") == "assistant"]) == 1 From 93246fc642768011b12a1e9ec01a2823151b61cf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 29 Jun 2026 13:29:36 -0700 Subject: [PATCH 2/2] style(databricks): use builtin list generics in parallel tool-call split Switch the List[...] annotations introduced by _split_parallel_tool_calls to lowercase list[...] so the UP006 strict-rule budget stays within its ceiling. --- litellm/llms/databricks/chat/transformation.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index dff86fd67c5c..ba8c312ea51c 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -95,7 +95,7 @@ def _sanitize_empty_content(message_dict: dict[str, Any]) -> None: message_dict["content"] = filtered -def _split_parallel_tool_calls(messages: List[AllMessageValues]) -> List[AllMessageValues]: +def _split_parallel_tool_calls(messages: list[AllMessageValues]) -> list[AllMessageValues]: """ Databricks (OpenAI-compatible serving) rejects a ``tool`` message unless the message immediately before it carries ``tool_calls``. A single assistant turn @@ -112,7 +112,7 @@ def _split_parallel_tool_calls(messages: List[AllMessageValues]) -> List[AllMess def _expand( assistant: ChatCompletionAssistantMessage, calls_by_id: dict[Optional[str], ChatCompletionAssistantToolCall], - tool_messages: List[ChatCompletionToolMessage], + tool_messages: list[ChatCompletionToolMessage], ) -> Iterator[AllMessageValues]: for position, tool_message in enumerate(tool_messages): matched_call = calls_by_id[tool_message["tool_call_id"]] @@ -134,7 +134,7 @@ def _generate() -> Iterator[AllMessageValues]: end = index + 1 while end < len(messages) and messages[end]["role"] == "tool": end += 1 - tool_messages = cast(List[ChatCompletionToolMessage], messages[index + 1 : end]) + tool_messages = cast(list[ChatCompletionToolMessage], messages[index + 1 : end]) calls_by_id = {call["id"]: call for call in tool_calls} result_ids = {tool_message["tool_call_id"] for tool_message in tool_messages} if len(tool_messages) == len(tool_calls) and set(calls_by_id) == result_ids: @@ -441,7 +441,7 @@ def _transform_messages( new_messages.append(_message) if "claude" not in model: - new_messages = _split_parallel_tool_calls(cast(List[AllMessageValues], new_messages)) + new_messages = _split_parallel_tool_calls(cast(list[AllMessageValues], new_messages)) if is_async: return super()._transform_messages(messages=new_messages, model=model, is_async=cast(Literal[True], True))