From 97dd4fc63be2c21efe7ea9ee501cb9aa5d13bebd Mon Sep 17 00:00:00 2001 From: Kcstring Date: Sun, 19 Apr 2026 20:36:30 +0800 Subject: [PATCH] feat(anthropic): support task_budget in output_config with auto beta header injection Adds support for the Anthropic task-budgets-2026-03-13 beta feature. - Add TASK_BUDGETS_2026_03_13 to ANTHROPIC_BETA_HEADER_VALUES enum - Auto-inject 'task-budgets-2026-03-13' beta header when output_config.task_budget is set - Validate task_budget shape in _apply_output_config (type must be 'tokens', total must be positive int) - 7 new unit tests covering header injection, merging, and validation Closes #25971 --- litellm/llms/anthropic/chat/transformation.py | 25 +++++ litellm/types/llms/anthropic.py | 1 + .../test_anthropic_chat_transformation.py | 103 ++++++++++++++++++ 3 files changed, 129 insertions(+) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index d7ce6a5f8de..6eeb505cc51 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1334,6 +1334,11 @@ 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") or {} + if isinstance(_output_config, dict) and _output_config.get("task_budget"): + 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( @@ -1505,6 +1510,26 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) data["output_config"] = output_config + # Validate task_budget if present + task_budget = output_config.get("task_budget") + if task_budget is not None: + if not isinstance(task_budget, dict): + raise ValueError( + "output_config.task_budget must be a dict, " + f"e.g. {{'type': 'tokens', 'total': 64000}}. Got: {type(task_budget)}" + ) + if task_budget.get("type") != "tokens": + raise ValueError( + "output_config.task_budget.type must be 'tokens'. " + f"Got: {task_budget.get('type')!r}" + ) + total = task_budget.get("total") + if not isinstance(total, int) or total <= 0: + raise ValueError( + "output_config.task_budget.total must be a positive integer. " + f"Got: {total!r}" + ) + def _transform_response_for_json_mode( self, json_mode: Optional[bool], diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 6d428173862..61d881567ac 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -677,6 +677,7 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum): ADVANCED_TOOL_USE_2025_11_20 = "advanced-tool-use-2025-11-20" FAST_MODE_2026_02_01 = "fast-mode-2026-02-01" ADVISOR_TOOL_2026_03_01 = "advisor-tool-2026-03-01" + TASK_BUDGETS_2026_03_13 = "task-budgets-2026-03-13" # Tool search beta header constant (for Anthropic direct API and Microsoft Foundry) 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 6b1b9ced245..0d2d4e64777 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 @@ -3601,3 +3601,106 @@ def test_strip_advisor_blocks_no_op_when_no_advisor_blocks(): original_content = [dict(b) for b in messages[1]["content"]] result = strip_advisor_blocks_from_messages(messages) assert result[1]["content"] == original_content + + +# --------------------------------------------------------------------------- +# task_budget tests +# --------------------------------------------------------------------------- + + +def test_task_budget_beta_header_injected(): + """output_config with task_budget should add the task-budgets beta header.""" + config = AnthropicConfig() + headers = config.update_headers_with_optional_anthropic_beta( + headers={}, + optional_params={ + "output_config": { + "effort": "high", + "task_budget": {"type": "tokens", "total": 64000}, + } + }, + ) + assert "task-budgets-2026-03-13" in headers.get("anthropic-beta", ""), ( + f"task-budgets beta header missing from: {headers}" + ) + + +def test_task_budget_beta_header_not_injected_without_task_budget(): + """output_config without task_budget should NOT add the task-budgets beta header.""" + config = AnthropicConfig() + headers = config.update_headers_with_optional_anthropic_beta( + headers={}, + optional_params={"output_config": {"effort": "high"}}, + ) + assert "task-budgets-2026-03-13" not in headers.get("anthropic-beta", "") + + +def test_task_budget_beta_header_merges_with_existing(): + """task_budget beta header should merge with pre-existing anthropic-beta values.""" + config = AnthropicConfig() + headers = {"anthropic-beta": "context-1m-2025-08-07"} + result = config.update_headers_with_optional_anthropic_beta( + headers=headers, + optional_params={ + "output_config": {"task_budget": {"type": "tokens", "total": 32000}} + }, + ) + beta = result["anthropic-beta"] + assert "context-1m-2025-08-07" in beta + assert "task-budgets-2026-03-13" in beta + + +def test_apply_output_config_task_budget_valid(): + """Valid task_budget passes validation and is included in data.""" + config = AnthropicConfig() + data: dict = {} + config._apply_output_config( + data=data, + model="claude-opus-4-7-20251101", + optional_params={ + "output_config": { + "effort": "high", + "task_budget": {"type": "tokens", "total": 64000}, + } + }, + ) + assert data["output_config"]["task_budget"] == {"type": "tokens", "total": 64000} + + +def test_apply_output_config_task_budget_invalid_type(): + """task_budget with wrong type field raises ValueError.""" + config = AnthropicConfig() + with pytest.raises(ValueError, match="task_budget.type must be 'tokens'"): + config._apply_output_config( + data={}, + model="claude-opus-4-7-20251101", + optional_params={ + "output_config": { + "task_budget": {"type": "characters", "total": 64000} + } + }, + ) + + +def test_apply_output_config_task_budget_invalid_total(): + """task_budget with non-positive total raises ValueError.""" + config = AnthropicConfig() + with pytest.raises(ValueError, match="task_budget.total must be a positive integer"): + config._apply_output_config( + data={}, + model="claude-opus-4-7-20251101", + optional_params={ + "output_config": {"task_budget": {"type": "tokens", "total": -1}} + }, + ) + + +def test_apply_output_config_task_budget_not_a_dict(): + """task_budget that is not a dict raises ValueError.""" + config = AnthropicConfig() + with pytest.raises(ValueError, match="task_budget must be a dict"): + config._apply_output_config( + data={}, + model="claude-opus-4-7-20251101", + optional_params={"output_config": {"task_budget": 64000}}, + )