Add Anthropic task budget support

This commit is contained in:
Jayesh Patel 2026-04-17 20:45:46 -04:00
parent e3c2b7561a
commit 965bc4a5a7
15 changed files with 89 additions and 32 deletions

View file

@ -51,6 +51,7 @@ if TYPE_CHECKING:
else:
AsyncIOScheduler = Any
class PrometheusLogger(CustomLogger):
# Class variables or attributes
@ -991,9 +992,7 @@ class PrometheusLogger(CustomLogger):
amount: float = 1.0,
) -> None:
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name=metric_name
),
supported_enum_labels=self.get_labels_for_metric(metric_name=metric_name),
enum_values=enum_values,
label_context=label_context,
)
@ -1118,7 +1117,9 @@ class PrometheusLogger(CustomLogger):
user_api_key = hash_token(user_api_key)
label_context = PrometheusLabelFactoryContext(enum_values) #amortized per request.
label_context = PrometheusLabelFactoryContext(
enum_values
) # amortized per request.
# increment total LLM requests and spend metric
self._increment_top_level_request_and_spend_metrics(
@ -3490,7 +3491,9 @@ def _prometheus_labels_from_context(
}
if UserAPIKeyLabelNames.END_USER.value in filtered_labels:
filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ctx.get_resolved_end_user()
filtered_labels[UserAPIKeyLabelNames.END_USER.value] = (
ctx.get_resolved_end_user()
)
for sk, val in ctx._custom_by_sanitized_key.items():
if sk in supported_enum_labels:

View file

@ -51,8 +51,7 @@ class PrometheusLabelFactoryContext:
self.enum_values = enum_values
enum_dict = enum_values.model_dump()
self._sanitized_enum: Dict[str, Optional[str]] = {
k: _sanitize_prometheus_label_value(v)
for k, v in enum_dict.items()
k: _sanitize_prometheus_label_value(v) for k, v in enum_dict.items()
}
self._custom_by_sanitized_key: Dict[str, Optional[str]] = {}
if enum_values.custom_metadata_labels is not None:

View file

@ -74,6 +74,7 @@ from litellm.utils import (
has_tool_call_blocks,
last_assistant_with_tool_calls_has_no_thinking_blocks,
supports_reasoning,
supports_task_budget,
token_counter,
)
@ -190,14 +191,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
v in model_lower for v in ("opus-4-6", "opus_4_6", "opus-4.6", "opus_4.6")
)
@staticmethod
def _is_opus_4_7_model(model: str) -> bool:
"""Check if the model is specifically Claude Opus 4.7."""
model_lower = model.lower()
return any(
v in model_lower for v in ("opus-4-7", "opus_4_7", "opus-4.7", "opus_4.7")
)
@staticmethod
def _supports_effort_level(model: str, level: str) -> bool:
"""Check ``supports_{level}_reasoning_effort`` in the model map.
@ -234,7 +227,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
"cache_control",
]
if AnthropicConfig._is_claude_4_7_model(model):
if supports_task_budget(
model=model,
custom_llm_provider=self.custom_llm_provider,
):
params.append("output_config")
if (
@ -1381,6 +1377,14 @@ 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(
@ -1563,9 +1567,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
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):
if not supports_task_budget(
model=model,
custom_llm_provider=self.custom_llm_provider,
):
raise ValueError(
f"output_config.task_budget is only supported by Claude Opus 4.7. "
f"output_config.task_budget is not supported by this model. "
f"Got model: {model}"
)
if not isinstance(task_budget, dict):

View file

@ -372,6 +372,14 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
if optional_params.get("speed") == "fast":
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.FAST_MODE_2026_02_01.value)
# Check for task budgets (Claude Opus 4.7+)
output_config = optional_params.get("output_config")
if (
isinstance(output_config, dict)
and output_config.get("task_budget") is not None
):
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.TASK_BUDGETS_2026_03_13.value)
# Check for advisor tool
tools = optional_params.get("tools")
if tools:

View file

@ -294,9 +294,7 @@ class Authenticator:
access_token_url = os.getenv(
"GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL
)
client_id = os.getenv(
"GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID
)
client_id = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
for attempt in range(max_attempts):
try:

View file

@ -1144,6 +1144,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
@ -9162,6 +9163,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,
"provider_specific_entry": {
@ -9194,6 +9196,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,
"provider_specific_entry": {

View file

@ -79,7 +79,9 @@ class BasePassthroughUtils:
for header_name, header_value in request_headers.items():
if header_name.lower().startswith(PASS_THROUGH_HEADER_PREFIX):
# Strip the 'x-pass-' prefix and normalize to lowercase
actual_header_name = header_name[len(PASS_THROUGH_HEADER_PREFIX) :].lower()
actual_header_name = header_name[
len(PASS_THROUGH_HEADER_PREFIX) :
].lower()
if actual_header_name in _PASS_THROUGH_PROTECTED_HEADERS or any(
actual_header_name.startswith(p)
for p in _PASS_THROUGH_PROTECTED_HEADER_PREFIXES

View file

@ -784,7 +784,7 @@ class UserAPIKeyLabelValues:
org_id: Optional[str] = None
org_alias: Optional[str] = None
#Added for test compatibility.
# Added for test compatibility.
def __init__(self, **kwargs: Any) -> None:
"""
Match former Pydantic behavior: unknown keys are ignored; ``api_key_hash`` maps to

View file

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

View file

@ -137,6 +137,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
supports_parallel_function_calling: Optional[bool]
supports_web_search: Optional[bool]
supports_reasoning: Optional[bool]
supports_task_budget: Optional[bool]
supports_url_context: Optional[bool]
supports_none_reasoning_effort: Optional[bool]
supports_xhigh_reasoning_effort: Optional[bool]

View file

@ -2738,6 +2738,15 @@ def supports_reasoning(model: str, custom_llm_provider: Optional[str] = None) ->
)
def supports_task_budget(model: str, custom_llm_provider: Optional[str] = None) -> bool:
"""
Check if the given model supports Anthropic task budgets and return a boolean value.
"""
return _supports_factory(
model=model, custom_llm_provider=custom_llm_provider, key="supports_task_budget"
)
def supports_native_structured_output(
model: str, custom_llm_provider: Optional[str] = None
) -> bool:
@ -5898,6 +5907,7 @@ def _get_model_info_helper( # noqa: PLR0915
supports_web_search=_model_info.get("supports_web_search", None),
supports_url_context=_model_info.get("supports_url_context", None),
supports_reasoning=_model_info.get("supports_reasoning", None),
supports_task_budget=_model_info.get("supports_task_budget", None),
supports_none_reasoning_effort=_model_info.get(
"supports_none_reasoning_effort", None
),

View file

@ -1144,6 +1144,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
@ -9162,6 +9163,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,
"provider_specific_entry": {
@ -9194,6 +9196,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,
"provider_specific_entry": {

View file

@ -8,6 +8,8 @@ sys.path.insert(
) # Adds the parent directory to the system path
from unittest.mock import MagicMock, patch
import litellm
from litellm import utils as litellm_utils
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
AnthropicMessagesConfig,
@ -16,6 +18,17 @@ from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.utils import ServerToolUse
def _enable_task_budget_for_opus_47(monkeypatch):
model_info = dict(litellm.model_cost.get("claude-opus-4-7-20260416", {}))
model_info["supports_task_budget"] = True
monkeypatch.setitem(
litellm.model_cost,
"claude-opus-4-7-20260416",
model_info,
)
litellm_utils._invalidate_model_cost_lowercase_map()
def test_response_format_transformation_unit_test():
config = AnthropicConfig()
@ -1559,8 +1572,9 @@ def test_effort_output_config_preservation():
assert result["output_config"]["effort"] == "medium"
def test_task_budget_output_config_preservation():
def test_task_budget_output_config_preservation(monkeypatch):
"""Test that output_config with task_budget is preserved for Claude Opus 4.7."""
_enable_task_budget_for_opus_47(monkeypatch)
config = AnthropicConfig()
messages = [{"role": "user", "content": "Run this agentic task"}]
@ -1607,8 +1621,9 @@ def test_task_budget_beta_header_injection():
assert ANTHROPIC_TASK_BUDGETS_BETA_HEADER in headers["anthropic-beta"]
def test_output_config_supported_for_claude_opus_47():
def test_output_config_supported_for_claude_opus_47(monkeypatch):
"""Test that output_config is accepted as a direct param for Claude Opus 4.7."""
_enable_task_budget_for_opus_47(monkeypatch)
config = AnthropicConfig()
supported_params = config.get_supported_openai_params(
@ -1665,7 +1680,7 @@ def test_task_budget_rejected_for_non_opus_47():
with pytest.raises(
ValueError,
match="output_config.task_budget is only supported by Claude Opus 4.7",
match="output_config.task_budget is not supported by this model",
):
config.transform_request(
model="claude-opus-4-6-20260205",
@ -3670,6 +3685,17 @@ def test_messages_path_advisor_beta_header_preserved_when_user_sends_it():
assert "advisor-tool-2026-03-01" in result.get("anthropic-beta", "")
def test_messages_path_task_budget_beta_header_injected():
"""task-budgets beta header is auto-injected in /messages path."""
config = AnthropicMessagesConfig()
headers: dict = {}
optional_params = {
"output_config": {"task_budget": {"type": "tokens", "total": 64000}}
}
result = config._update_headers_with_anthropic_beta(headers, optional_params)
assert "task-budgets-2026-03-13" in result.get("anthropic-beta", "")
def test_strip_advisor_blocks_when_no_advisor_tool():
"""
Auto-strip removes server_tool_use(advisor) + advisor_tool_result blocks when

View file

@ -34,13 +34,8 @@ def test_govcloud_cross_region_inference_prefix():
base_model = bedrock_model_info.get_base_model(
model="bedrock/us-gov.anthropic.claude-haiku-4-5-20251001-v1:0"
)
<<<<<<< worktree-rustling-wishing-kite
assert base_model == "anthropic.claude-3-5-sonnet-20240620-v1:0"
=======
assert base_model == "anthropic.claude-haiku-4-5-20251001-v1:0"
>>>>>>> main
# Test us-gov prefix is stripped correctly for different Claude versions
base_model = bedrock_model_info.get_base_model(
model="bedrock/us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0"

View file

@ -768,6 +768,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"supports_multimodal": {"type": "boolean"},
"uses_embed_content": {"type": "boolean"},
"supports_reasoning": {"type": "boolean"},
"supports_task_budget": {"type": "boolean"},
"supports_minimal_reasoning_effort": {"type": "boolean"},
"supports_none_reasoning_effort": {"type": "boolean"},
"supports_xhigh_reasoning_effort": {"type": "boolean"},