From bfb4f4b2c399df60a541b25597d447fec23fea23 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:22:10 -0700 Subject: [PATCH] fix(databricks): accept service_tier on requests and carry the served tier on responses (#46111) * fix(databricks): accept service_tier on requests and carry the served tier on responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(databricks): mark the in-place service_tier store as rebind-ok Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(databricks): wrap the service_tier store to the line limit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(databricks): keep the rebind-ok suppression on the service_tier line Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(databricks): only set service_tier on responses when the upstream returned one Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(databricks): name the B010 suppression on the service_tier setattr Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(databricks): shorten the service_tier suppression to the line limit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): read service_tier from the Responses request only Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/databricks/chat/transformation.py | 4 ++ .../transformation.py | 2 +- .../test_databricks_chat_transformation.py | 39 +++++++++++++++++++ .../test_litellm_completion_responses.py | 11 ++++++ 4 files changed, 55 insertions(+), 1 deletion(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 50d3b82a989..dceec658824 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -285,6 +285,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): "tool_choice", "reasoning_effort", "thinking", + "service_tier", ] @staticmethod @@ -697,6 +698,9 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): model_response.id = completion_response["id"] model_response.created = completion_response["created"] setattr(model_response, "usage", Usage(**completion_response["usage"])) + served_tier: Final = completion_response.get("service_tier") + if isinstance(served_tier, str) and served_tier: + setattr(model_response, "service_tier", served_tier) # noqa: B010 # no declared field model_response.choices = self._transform_dbrx_choices( choices=completion_response["choices"], diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 4f99c570551..c2cc876970c 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -500,7 +500,7 @@ class LiteLLMCompletionResponsesConfig: "max_tokens": responses_api_request.get("max_output_tokens"), "stream": stream, "metadata": kwargs.get("metadata"), - "service_tier": kwargs.get("service_tier"), + "service_tier": responses_api_request.get("service_tier"), "safety_identifier": responses_api_request.get("safety_identifier"), "web_search_options": web_search_options, "response_format": response_format, 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 c6e0ecd2b41..fe76a667f3e 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -2216,3 +2216,42 @@ def test_completion_with_prompt_caching_anthropic_model( assert response.usage.total_tokens == 1666 assert response.usage.cache_read_input_tokens == 0 assert response.usage.cache_creation_input_tokens == 1545 + + +def test_get_optional_params_keeps_service_tier_for_databricks() -> None: + optional_params: Final = litellm.utils.get_optional_params( + model="databricks-claude-opus-5", + custom_llm_provider="databricks", + service_tier="priority", + drop_params=True, + ) + + assert optional_params["service_tier"] == "priority" + + +@pytest.mark.parametrize("served_tier", ["priority", "default", None]) +def test_transform_response_carries_served_service_tier(served_tier: str | None) -> None: + raw_response: Final = MagicMock() + raw_response.json.return_value = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "dbrx", + **({"service_tier": served_tier} if served_tier is not None else {}), + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + + result: Final = DatabricksConfig().transform_response( + model="databricks/dbrx", + raw_response=raw_response, + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.to_dict().get("service_tier") == served_tier diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index c7f20a1d959..7006feb22b2 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -5180,3 +5180,14 @@ def test_convert_response_function_tool_call_keeps_the_index_it_is_given(index: ) assert tool_call["index"] == index + + +def test_transform_request_reads_service_tier_from_responses_api_request() -> None: + result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="databricks/dbrx", + input="hello", + responses_api_request={"service_tier": "priority"}, + custom_llm_provider="databricks", + ) + + assert result["service_tier"] == "priority"