Address task budget PR feedback

This commit is contained in:
Jayesh Patel 2026-04-17 21:31:39 -04:00
parent 965bc4a5a7
commit bddf3ea1b2
4 changed files with 41 additions and 8 deletions

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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