diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index d6007a2d56e..1e4836654a1 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1981,9 +1981,28 @@ def get_model_from_request( if vertex_match: model = vertex_match.group(1) + if route.lower().startswith("/bedrock"): + bedrock_model: Final = _model_from_bedrock_route(route) + return model if bedrock_model is None else bedrock_model + return model +def _model_from_bedrock_route(route: str) -> str | None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _extract_model_from_bedrock_endpoint, + is_bedrock_count_tokens_endpoint, + ) + + bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE) + if is_bedrock_count_tokens_endpoint(bedrock_endpoint): + return None + try: + return _extract_model_from_bedrock_endpoint(bedrock_endpoint) + except ValueError: + return None + + def abbreviate_api_key(api_key: str) -> str: if len(api_key) < MINIMUM_CUSTOM_KEY_LENGTH: return "sk-..." diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 32da2658b99..b95547e2b54 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -712,6 +712,10 @@ BEDROCK_ENDPOINT_ACTIONS: Final = { BEDROCK_STREAMING_ACTIONS: Final = {"invoke-with-response-stream", "converse-stream"} +def is_bedrock_count_tokens_endpoint(endpoint: str) -> bool: + return "count_tokens" in endpoint or "count-tokens" in endpoint + + def _extract_model_from_bedrock_endpoint(endpoint: str) -> str: """ Extract model name from Bedrock endpoint path. @@ -986,8 +990,7 @@ async def bedrock_llm_proxy_route( request_body: Final = await _read_request_body(request=request) - # Special handling for count_tokens endpoints - if "count_tokens" in endpoint or "count-tokens" in endpoint: + if is_bedrock_count_tokens_endpoint(endpoint): return await handle_bedrock_count_tokens( endpoint=endpoint, request=request, diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index ed2bc87597c..152f3befa15 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1389,12 +1389,15 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int: def _count_input_tokens(request_body: dict, model: str) -> int | None: try: if "messages" in request_body: - return litellm.token_counter( - model=model, - messages=request_body.get("messages") or [], - tools=request_body.get("tools"), - tool_choice=request_body.get("tool_choice"), - ) + try: + return litellm.token_counter( + model=model, + messages=request_body.get("messages") or (), + tools=request_body.get("tools"), + tool_choice=request_body.get("tool_choice"), + ) + except ValueError: + return _count_text_tokens(model=model, text=request_body.get("messages")) if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: @@ -1442,11 +1445,7 @@ def _estimate_output_tokens( if _is_input_only_route(route=route): return 0 - requested: int | None = None - for key in ("max_completion_tokens", "max_tokens", "max_output_tokens"): - requested = _to_int(request_body.get(key)) - if requested is not None: - break + requested: Final = _requested_output_tokens(request_body) # Clamp at min(requested-or-default, model_max-or-default). Two purposes: # (1) Without an explicit cap we still need a finite reservation so the @@ -1457,9 +1456,19 @@ def _estimate_output_tokens( # at the cap — the model can only physically emit max_output_tokens # anyway, so reserving more is both wasteful and a DoS surface. model_ceiling: Final = _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK - if requested is None: - requested = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK - return min(requested, model_ceiling) + return min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK if requested is None else requested, model_ceiling) + + +_OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens") + + +def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: + inference_config: Final = request_body.get("inferenceConfig") + candidates: Final = ( + *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), + inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None, + ) + return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) def _count_text_tokens(model: str, text: object) -> int: diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index f513f397b64..a996de4d40c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -428,6 +428,106 @@ def test_get_model_from_request_openai_deployment_route_still_works(): ) +def test_get_model_from_request_bedrock_converse_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/converse", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_invoke_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_v2_converse_stream_passthrough(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/v2/model/us.anthropic.claude-sonnet-4-6/converse-stream", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_model_id_with_slashes(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/model/aws/anthropic/model-name/invoke", + ) + == "aws/anthropic/model-name" + ) + + +def test_get_model_from_request_bedrock_unparseable_endpoint_returns_none(): + assert ( + get_model_from_request( + request_data={}, + route="/bedrock/agents/some-agent-route", + ) + is None + ) + + +def test_get_model_from_request_bedrock_url_model_overrides_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/converse", + ) + == "us.anthropic.claude-opus-4-6-v1" + ) + + +def test_get_model_from_request_bedrock_invoke_url_model_overrides_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/invoke", + ) + == "us.anthropic.claude-opus-4-6-v1" + ) + + +def test_get_model_from_request_bedrock_count_tokens_uses_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/v1/messages/count_tokens", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_uppercase_count_tokens_segment_is_not_count_tokens(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-haiku-4-5-20251001-v1:0"}, + route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke/COUNT_TOKENS", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/agents/some-agent-route", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + def test_get_model_from_request_includes_file_endpoint_header_model(): assert ( get_model_from_request( diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index f65f68812a2..c7de8c943ad 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -5,7 +5,7 @@ import pytest from litellm.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.spend_tracking.budget_reservation import reserve_budget_for_request +from litellm.proxy.spend_tracking.budget_reservation import estimate_request_max_cost, reserve_budget_for_request from litellm.proxy.utils import ProxyLogging TOKEN_COUNTING_ROUTES: Final = ( @@ -46,3 +46,32 @@ async def test_non_exempt_llm_route_still_reserves_budget(): assert reservation is not None assert reservation["reserved_cost"] > 0 + + +BEDROCK_SONNET: Final = "us.anthropic.claude-sonnet-4-6" +CONVERSE_BODY: Final = { + "messages": [{"role": "user", "content": [{"text": "Reply with one word: pong"}]}], + "inferenceConfig": {"maxTokens": 5}, +} +INVOKE_BODY: Final = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 5, + "messages": [{"role": "user", "content": "Reply with one word: pong"}], +} + + +def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window(): + converse_cost: Final = estimate_request_max_cost( + request_body=CONVERSE_BODY, + route=f"/bedrock/model/{BEDROCK_SONNET}/converse", + llm_router=None, + input_token_counts={}, + ) + invoke_cost: Final = estimate_request_max_cost( + request_body=INVOKE_BODY, + route=f"/bedrock/model/{BEDROCK_SONNET}/invoke", + llm_router=None, + input_token_counts={}, + ) + assert converse_cost is not None and invoke_cost is not None + assert invoke_cost < converse_cost < 2 * invoke_cost