mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Add Anthropic task budget support
This commit is contained in:
parent
e3c2b7561a
commit
965bc4a5a7
15 changed files with 89 additions and 32 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
),
|
||||
|
|
|
|||
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue