From e3c2b7561aca1093597cab922b0ccf4f3ca64af4 Mon Sep 17 00:00:00 2001 From: Jayesh Patel Date: Fri, 17 Apr 2026 20:34:55 -0400 Subject: [PATCH] Add Anthropic task budget support --- litellm/llms/anthropic/chat/transformation.py | 32 ++++- litellm/llms/anthropic/common_utils.py | 29 +++++ litellm/types/llms/anthropic.py | 13 +- litellm/utils.py | 8 ++ .../test_anthropic_chat_transformation.py | 122 ++++++++++++++++++ 5 files changed, 202 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index cd5bb731717..590d0d37d89 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -234,6 +234,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "cache_control", ] + if AnthropicConfig._is_claude_4_7_model(model): + params.append("output_config") + if ( "claude-3-7-sonnet" in model or AnthropicConfig._is_claude_4_6_model(model) @@ -1105,7 +1108,18 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "max": "max", } mapped_effort = effort_map.get(value, value) - optional_params["output_config"] = {"effort": mapped_effort} + output_config = optional_params.get("output_config") + if not isinstance(output_config, dict): + output_config = {} + output_config["effort"] = mapped_effort + optional_params["output_config"] = output_config + elif param == "output_config" and isinstance(value, dict): + output_config = optional_params.get("output_config") + if isinstance(output_config, dict): + output_config.update(value) + optional_params["output_config"] = output_config + else: + optional_params["output_config"] = value elif param == "web_search_options" and isinstance(value, dict): hosted_web_search_tool = self.map_web_search_tool( cast(OpenAIWebSearchOptions, value) @@ -1528,6 +1542,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if not output_config or not isinstance(output_config, dict): return effort = output_config.get("effort") + task_budget = output_config.get("task_budget") valid_efforts = ["high", "medium", "low", "xhigh", "max"] if effort and effort not in valid_efforts: raise ValueError( @@ -1547,6 +1562,21 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): raise ValueError( f"effort='xhigh' is not supported by this model. Got model: {model}" ) + if task_budget is not None: + if not self._is_opus_4_7_model(model): + raise ValueError( + f"output_config.task_budget is only supported by Claude Opus 4.7. " + f"Got model: {model}" + ) + if not isinstance(task_budget, dict): + raise ValueError("output_config.task_budget must be a dictionary") + if task_budget.get("type") != "tokens": + raise ValueError("output_config.task_budget.type must be 'tokens'") + 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" + ) data["output_config"] = output_config def _transform_response_for_json_mode( diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index b095d40156d..7b00e8bfa3e 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -310,6 +310,20 @@ class AnthropicModelInfo(BaseLLMModelInfo): return False + def is_task_budget_used(self, optional_params: Optional[dict]) -> bool: + """ + Check if output_config.task_budget is being used. + """ + if not optional_params: + return False + + output_config = optional_params.get("output_config") + if not output_config or not isinstance(output_config, dict): + return False + + task_budget = output_config.get("task_budget") + return bool(task_budget and isinstance(task_budget, dict)) + def is_code_execution_tool_used(self, tools: Optional[List]) -> bool: """ Check if code execution tool is being used. @@ -382,15 +396,20 @@ class AnthropicModelInfo(BaseLLMModelInfo): List of beta header strings """ from litellm.types.llms.anthropic import ANTHROPIC_EFFORT_BETA_HEADER + from litellm.types.llms.anthropic import ANTHROPIC_TASK_BUDGETS_BETA_HEADER betas = [] # Detect features effort_used = self.is_effort_used(optional_params, model) + task_budget_used = self.is_task_budget_used(optional_params) if effort_used: betas.append(ANTHROPIC_EFFORT_BETA_HEADER) # effort-2025-11-24 + if task_budget_used: + betas.append(ANTHROPIC_TASK_BUDGETS_BETA_HEADER) + if computer_tool_used: beta_header = self.get_computer_tool_beta_header(computer_tool_used) betas.append(beta_header) @@ -423,6 +442,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): programmatic_tool_calling_used: bool = False, input_examples_used: bool = False, effort_used: bool = False, + task_budget_used: bool = False, is_vertex_request: bool = False, user_anthropic_beta_headers: Optional[List[str]] = None, code_execution_tool_used: bool = False, @@ -454,6 +474,13 @@ class AnthropicModelInfo(BaseLLMModelInfo): betas.add(ANTHROPIC_EFFORT_BETA_HEADER) + if task_budget_used: + from litellm.types.llms.anthropic import ( + ANTHROPIC_TASK_BUDGETS_BETA_HEADER, + ) + + betas.add(ANTHROPIC_TASK_BUDGETS_BETA_HEADER) + # Code execution tool uses a separate beta header if code_execution_tool_used: betas.add("code-execution-2025-08-25") @@ -535,6 +562,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): ) input_examples_used = self.is_input_examples_used(tools=tools) effort_used = self.is_effort_used(optional_params=optional_params, model=model) + task_budget_used = self.is_task_budget_used(optional_params=optional_params) code_execution_tool_used = self.is_code_execution_tool_used(tools=tools) container_with_skills_used = self.is_container_with_skills_used( optional_params=optional_params @@ -557,6 +585,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): programmatic_tool_calling_used=programmatic_tool_calling_used, input_examples_used=input_examples_used, effort_used=effort_used, + task_budget_used=task_budget_used, code_execution_tool_used=code_execution_tool_used, container_with_skills_used=container_with_skills_used, ) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index e3f63d05742..4d99c419039 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -36,10 +36,18 @@ class AnthropicOutputSchema(TypedDict, total=False): schema: Required[dict] +class AnthropicTaskBudget(TypedDict, total=False): + """Soft token budget for an entire Claude agentic loop.""" + + type: Required[Literal["tokens"]] + total: Required[int] + + class AnthropicOutputConfig(TypedDict, total=False): """Configuration for controlling Claude's output behavior.""" - effort: Literal["high", "medium", "low"] + effort: Literal["high", "medium", "low", "xhigh", "max"] + task_budget: AnthropicTaskBudget class AnthropicMessagesTool(TypedDict, total=False): @@ -677,6 +685,9 @@ ANTHROPIC_TOOL_SEARCH_BETA_HEADER = "advanced-tool-use-2025-11-20" # Effort beta header constant ANTHROPIC_EFFORT_BETA_HEADER = "effort-2025-11-24" +# Task budget beta header constant +ANTHROPIC_TASK_BUDGETS_BETA_HEADER = "task-budgets-2026-03-13" + # OAuth constants ANTHROPIC_OAUTH_TOKEN_PREFIX = "sk-ant-oat" ANTHROPIC_OAUTH_BETA_HEADER = "oauth-2025-04-20" diff --git a/litellm/utils.py b/litellm/utils.py index 2125875ee1d..b7d6f5ca00f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4804,6 +4804,14 @@ def add_provider_specific_params_to_optional_params( k=k, additional_drop_params=additional_drop_params ): continue + if isinstance(optional_params.get(k), dict) and isinstance( + passed_params[k], dict + ): + optional_params[k] = { + **optional_params[k], + **passed_params[k], + } + continue optional_params[k] = passed_params[k] return optional_params 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 a7f5f92ab05..089eb4a4ac8 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 @@ -1559,6 +1559,128 @@ def test_effort_output_config_preservation(): assert result["output_config"]["effort"] == "medium" +def test_task_budget_output_config_preservation(): + """Test that output_config with task_budget is preserved for Claude Opus 4.7.""" + config = AnthropicConfig() + + messages = [{"role": "user", "content": "Run this agentic task"}] + optional_params = { + "output_config": { + "effort": "high", + "task_budget": {"type": "tokens", "total": 64000}, + } + } + + result = config.transform_request( + model="claude-opus-4-7-20260416", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result["output_config"]["effort"] == "high" + assert result["output_config"]["task_budget"] == { + "type": "tokens", + "total": 64000, + } + + +def test_task_budget_beta_header_injection(): + """Test that task budget beta header is added when task_budget is detected.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + from litellm.types.llms.anthropic import ANTHROPIC_TASK_BUDGETS_BETA_HEADER + + model_info = AnthropicModelInfo() + optional_params = { + "output_config": {"task_budget": {"type": "tokens", "total": 64000}} + } + + task_budget_used = model_info.is_task_budget_used(optional_params=optional_params) + assert task_budget_used is True + + headers = model_info.get_anthropic_headers( + api_key="test-key", task_budget_used=task_budget_used + ) + + assert "anthropic-beta" in headers + assert ANTHROPIC_TASK_BUDGETS_BETA_HEADER in headers["anthropic-beta"] + + +def test_output_config_supported_for_claude_opus_47(): + """Test that output_config is accepted as a direct param for Claude Opus 4.7.""" + config = AnthropicConfig() + + supported_params = config.get_supported_openai_params( + model="claude-opus-4-7-20260416" + ) + + assert "output_config" in supported_params + + +def test_output_config_maps_direct_param(): + """Test that output_config maps from non-default params.""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={ + "reasoning_effort": "high", + "output_config": {"task_budget": {"type": "tokens", "total": 64000}}, + }, + optional_params={}, + model="claude-opus-4-7-20260416", + drop_params=False, + ) + + assert result["output_config"]["effort"] == "high" + assert result["output_config"]["task_budget"] == { + "type": "tokens", + "total": 64000, + } + + +def test_output_config_maps_from_completion_optional_params(): + """Test that output_config is accepted through the public optional params path.""" + from litellm.utils import get_optional_params + + result = get_optional_params( + model="claude-opus-4-7-20260416", + custom_llm_provider="anthropic", + messages=[{"role": "user", "content": "Run this agentic task"}], + reasoning_effort="high", + output_config={"task_budget": {"type": "tokens", "total": 64000}}, + ) + + assert result["output_config"]["effort"] == "high" + assert result["output_config"]["task_budget"] == { + "type": "tokens", + "total": 64000, + } + + +def test_task_budget_rejected_for_non_opus_47(): + """Test that task_budget is rejected for models other than Claude Opus 4.7.""" + config = AnthropicConfig() + messages = [{"role": "user", "content": "Test"}] + + with pytest.raises( + ValueError, + match="output_config.task_budget is only supported by Claude Opus 4.7", + ): + config.transform_request( + model="claude-opus-4-6-20260205", + messages=messages, + optional_params={ + "output_config": { + "effort": "high", + "task_budget": {"type": "tokens", "total": 64000}, + } + }, + litellm_params={}, + headers={}, + ) + + 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