From bddf3ea1b2a817f24f2e89228743dec8f7d8f56f Mon Sep 17 00:00:00 2001 From: Jayesh Patel Date: Fri, 17 Apr 2026 21:31:39 -0400 Subject: [PATCH] Address task budget PR feedback --- litellm/llms/anthropic/chat/transformation.py | 16 ++++++------ ...odel_prices_and_context_window_backup.json | 4 +++ model_prices_and_context_window.json | 4 +++ .../test_anthropic_chat_transformation.py | 25 +++++++++++++++++++ 4 files changed, 41 insertions(+), 8 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index f702fb74637..d35f0a7c787 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1377,14 +1377,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): self._ensure_beta_header( headers, ANTHROPIC_BETA_HEADER_VALUES.FAST_MODE_2026_02_01.value ) - output_config = optional_params.get("output_config") - if ( - isinstance(output_config, dict) - and output_config.get("task_budget") is not None - ): - self._ensure_beta_header( - headers, ANTHROPIC_BETA_HEADER_VALUES.TASK_BUDGETS_2026_03_13.value - ) for tool in _tools: if tool.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE: self._ensure_beta_header( @@ -1533,6 +1525,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): self._apply_output_config( data=data, model=model, optional_params=optional_params ) + output_config = optional_params.get("output_config") + if ( + isinstance(output_config, dict) + and output_config.get("task_budget") is not None + ): + self._ensure_beta_header( + headers, ANTHROPIC_BETA_HEADER_VALUES.TASK_BUDGETS_2026_03_13.value + ) return data diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a49e711f70e..ab11b4bca21 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1173,6 +1173,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_task_budget": true, "supports_xhigh_reasoning_effort": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true @@ -1201,6 +1202,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_task_budget": true, "supports_xhigh_reasoning_effort": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true @@ -1229,6 +1231,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_task_budget": true, "supports_xhigh_reasoning_effort": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true @@ -1257,6 +1260,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_task_budget": true, "supports_xhigh_reasoning_effort": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a49e711f70e..ab11b4bca21 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1173,6 +1173,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_task_budget": true, "supports_xhigh_reasoning_effort": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true @@ -1201,6 +1202,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_task_budget": true, "supports_xhigh_reasoning_effort": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true @@ -1229,6 +1231,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_task_budget": true, "supports_xhigh_reasoning_effort": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true @@ -1257,6 +1260,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_task_budget": true, "supports_xhigh_reasoning_effort": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 1edc588feb3..eb3bbe4fa4a 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1696,6 +1696,31 @@ def test_task_budget_rejected_for_non_opus_47(): ) +def test_task_budget_rejection_does_not_mutate_headers(): + """Test that invalid task_budget requests do not leak beta headers.""" + config = AnthropicConfig() + messages = [{"role": "user", "content": "Test"}] + headers = {"anthropic-beta": "existing-beta"} + + with pytest.raises( + ValueError, + match="output_config.task_budget is not supported by this model", + ): + config.transform_request( + model="claude-opus-4-6-20260205", + messages=messages, + optional_params={ + "output_config": { + "task_budget": {"type": "tokens", "total": 64000}, + } + }, + litellm_params={}, + headers=headers, + ) + + assert headers == {"anthropic-beta": "existing-beta"} + + def test_effort_beta_header_injection(): """Test that effort beta header is automatically added when output_config is detected.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo