diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 6f0a18b0088..3866240030b 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -21,7 +21,7 @@ from litellm.llms.bedrock.common_utils import ( BedrockError, get_bedrock_base_model, ) -from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler, BedrockCountTokensRequestError from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -54,6 +54,7 @@ class CountTokensFailure: status_code: int message: str tokenizer_type: str + rejected_before_provider_call: bool = False CountTokensOutcome = CountedTokens | CountTokensFailure @@ -74,6 +75,7 @@ def _claude_count_translator(litellm_params: Mapping[str, object]) -> AnthropicC def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: return ( isinstance(outcome, CountTokensFailure) + and not outcome.rejected_before_provider_call and outcome.status_code in _RUNTIME_FALLBACK_STATUSES and is_mantle_claude_model(resolved_model) ) @@ -162,6 +164,13 @@ class BedrockTokenCounter(BaseTokenCounter): client=self._client, ) reply: Final = _CountTokensReply.model_validate(result) + except BedrockCountTokensRequestError as e: + return CountTokensFailure( + status_code=e.status_code, + message=e.message, + tokenizer_type=tokenizer_type, + rejected_before_provider_call=True, + ) except BedrockError as e: verbose_logger.debug( "%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message @@ -268,6 +277,7 @@ class BedrockTokenCounter(BaseTokenCounter): error=True, error_message=outcome.message, status_code=outcome.status_code, + rejected_before_provider_call=outcome.rejected_before_provider_call, ) case _: assert_never(outcome) diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 84067923788..eb29459f8b1 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -20,6 +20,11 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_h _DECODED_JSON: Final = TypeAdapter(object) +class BedrockCountTokensRequestError(BedrockError): + def __init__(self, problem: str) -> None: + super().__init__(status_code=400, message=problem) + + class BedrockCountTokensHandler(BedrockCountTokensConfig): """ Simplified handler for AWS Bedrock CountTokens API requests. @@ -45,10 +50,10 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): Returns: Dictionary containing token count response """ + problem: Final = self.count_tokens_request_problem(request_data) + if problem is not None: + raise BedrockCountTokensRequestError(problem) try: - # Validate the request - self.validate_count_tokens_request(request_data) - verbose_logger.debug("Processing CountTokens request for resolved model: %s", resolved_model) # Get AWS region using existing LiteLLM function diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 5e5c5516b94..0b589eda896 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -12,6 +12,8 @@ from litellm.llms.anthropic.common_utils import context_management_betas, tools_ from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url +from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensRequestError +from litellm.llms.bedrock.count_tokens.transformation import count_tokens_messages_problem from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens" @@ -58,6 +60,9 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM): ) -> dict[str, JsonValue]: try: request: Final = _COUNT_REQUEST.validate_python(request_data) + problem: Final = count_tokens_messages_problem(request["messages"]) + if problem is not None: + raise BedrockCountTokensRequestError(problem) aws_region_name: Final = await run_aws_signing( self.get_aws_region_name, optional_params=litellm_params, model=resolved_model, model_id=None ) diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index cd59e714ba0..aaeb57303f9 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -6,7 +6,7 @@ to AWS Bedrock's CountTokens API format and vice versa. """ import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Final, Literal from pydantic import JsonValue @@ -14,7 +14,11 @@ from pydantic import JsonValue from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import tools_without_custom_field from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM -from litellm.llms.bedrock.common_utils import get_bedrock_base_model +from litellm.llms.bedrock.common_utils import ( + get_bedrock_base_model, + get_bedrock_cross_region_inference_regions, + strip_bedrock_routing_prefix, +) # Placeholder satisfying the Anthropic InvokeModel schema's required # max_tokens field; CountTokens only counts input, so it has no effect @@ -46,6 +50,23 @@ def _to_converse_message(message: JsonValue) -> dict[str, JsonValue]: } +def _converse_message_problem(index: int, message: JsonValue) -> str | None: + if not isinstance(message, dict): + return f"Message {index} must be a dictionary" + if "role" not in message: + return f"Message {index} must have a 'role' field" + if "content" not in message and "tool_calls" in message: + return f"Message {index} has tool_calls and no content, which token counting cannot convert for this model" + if "content" not in message: + return f"Message {index} must have a 'content' field" + return None + + +def count_tokens_messages_problem(messages: Sequence[JsonValue]) -> str | None: + problems: Final = (_converse_message_problem(index, message) for index, message in enumerate(messages)) + return next((problem for problem in problems if problem is not None), None) + + def _sanitized_bedrock_tool_name(raw_name: JsonValue) -> str: name: Final = re.sub(r"[^a-zA-Z0-9_]", "_", raw_name if isinstance(raw_name, str) else "") prefixed: Final = name if not name or name[0].isalpha() else f"t_{name}" @@ -217,6 +238,9 @@ class BedrockCountTokensConfig(BaseAWSLLM): @staticmethod def count_tokens_model_id(model: str) -> str: + model_id: Final = strip_bedrock_routing_prefix(model) + if model_id in get_bedrock_cross_region_inference_regions(): + return model_id return get_bedrock_base_model(model).removeprefix("bedrock/") def get_bedrock_count_tokens_endpoint( @@ -265,42 +289,12 @@ class BedrockCountTokensConfig(BaseAWSLLM): return {"input_tokens": input_tokens} - def validate_count_tokens_request(self, request_data: Mapping[str, JsonValue]) -> None: - """ - Validate the incoming count tokens request. - Supports both Converse and InvokeModel input formats. - - Args: - request_data: The request payload - - Raises: - ValueError: If the request is invalid - """ + def count_tokens_request_problem(self, request_data: Mapping[str, JsonValue]) -> str | None: if not request_data.get("model"): - raise ValueError("model parameter is required") - - input_type: Final = self._detect_input_type(request_data) - - if input_type == "converse": - # Validate Converse format (messages-based) - messages: Final = request_data.get("messages", []) - if not messages: - raise ValueError("messages parameter is required for Converse input") - - if not isinstance(messages, list): - raise ValueError("messages must be a list") - - for i, message in enumerate(messages): - if not isinstance(message, dict): - raise ValueError(f"Message {i} must be a dictionary") - - if "role" not in message: - raise ValueError(f"Message {i} must have a 'role' field") - - if "content" not in message: - raise ValueError(f"Message {i} must have a 'content' field") - else: - # For InvokeModel format, we need at least some content to count tokens - # The content structure varies by model, so we do minimal validation - if len(request_data) <= 1: # Only has 'model' field - raise ValueError("Request must contain content to count tokens") + return "model parameter is required" + if self._detect_input_type(request_data) == "invokeModel": + return "Request must contain content to count tokens" if len(request_data) <= 1 else None + messages: Final = _json_list(request_data.get("messages")) + if not messages: + return "messages parameter is required for Converse input" + return count_tokens_messages_problem(messages) diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py index 591205b4602..19c9a042afb 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -6,6 +6,7 @@ import httpx import pytest from botocore.credentials import RefreshableCredentials +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -85,6 +86,42 @@ async def test_handle_count_tokens_request_looks_up_its_region_off_the_event_loo assert result == {"input_tokens": 7} +_COUNTED_MODEL: Final = "us.anthropic.claude-opus-4-8" + + +@pytest.mark.parametrize( + ("request_data", "problem"), + [ + ({"model": _COUNTED_MODEL}, "Request must contain content to count tokens"), + ({"model": _COUNTED_MODEL, "messages": []}, "messages parameter is required for Converse input"), + ({"model": _COUNTED_MODEL, "messages": [{"role": "user"}]}, "Message 0 must have a 'content' field"), + ({"model": _COUNTED_MODEL, "messages": [{"content": "hi"}]}, "Message 0 must have a 'role' field"), + ({"model": _COUNTED_MODEL, "messages": ["hi"]}, "Message 0 must be a dictionary"), + ( + {"model": _COUNTED_MODEL, "messages": [{"role": "user", "content": "hi"}, {"role": "assistant"}]}, + "Message 1 must have a 'content' field", + ), + ({"messages": [{"role": "user", "content": "hi"}]}, "model parameter is required"), + ], +) +@pytest.mark.asyncio +async def test_handle_count_tokens_request_answers_400_for_a_malformed_body_without_calling_bedrock( + request_data: dict[str, object], problem: str +) -> None: + client: Final = AsyncMock(spec=AsyncHTTPHandler) + + with pytest.raises(BedrockError) as exc_info: + await BedrockCountTokensHandler().handle_count_tokens_request( + request_data=request_data, + litellm_params={"aws_region_name": "us-east-1"}, + resolved_model=_COUNTED_MODEL, + client=client, + ) + + assert (exc_info.value.status_code, exc_info.value.message) == (400, problem) + client.post.assert_not_awaited() + + class TestBedrockCountTokensEndpoint: def _make_handler(self) -> BedrockCountTokensConfig: return BedrockCountTokensConfig() diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index d5b57e42146..018b0fde8a0 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -311,3 +311,37 @@ def test_count_tokens_endpoint_encodes_model_id(monkeypatch): endpoint == "https://bedrock-runtime.us-east-1.amazonaws.com/model/..%2F..%2Fmodel%2Fother%3Fx%3D1%23frag/count-tokens" ) + + +@pytest.mark.parametrize( + "request_data", + [ + {"model": "us.anthropic.claude-opus-4-8", "messages": [{"role": "user", "content": "hi"}]}, + {"model": "us.anthropic.claude-opus-4-8", "messages": "hi"}, + { + "model": "us.anthropic.claude-opus-4-8", + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}, {"content": "hi"}], + }, + ], +) +def test_count_tokens_request_problem_leaves_bodies_bedrock_can_judge_to_bedrock(request_data): + assert BedrockCountTokensConfig().count_tokens_request_problem(request_data) is None + + +def test_count_tokens_request_problem_names_tool_calls_instead_of_asking_for_content(): + problem = BedrockCountTokensConfig().count_tokens_request_problem( + { + "model": "us.amazon.nova-pro-v1:0", + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [{"id": "1", "type": "function", "function": {"name": "t", "arguments": "{}"}}], + }, + ], + } + ) + + assert problem is not None + assert "tool_calls" in problem + assert "must have a 'content' field" not in problem diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index 30063936b2b..de0c14558b3 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -74,6 +74,59 @@ async def test_claude_model_bedrock_runtime_cannot_count_is_counted_on_mantle() assert mantle_body["system"] == "You are a terse assistant." +@pytest.mark.asyncio +async def test_claude_model_message_without_content_is_rejected_without_calling_runtime_or_mantle() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=[{"role": "user"}], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + ) + + assert result is not None + assert (result.error, result.status_code, result.error_message, result.rejected_before_provider_call) == ( + True, + 400, + "Message 0 must have a 'content' field", + True, + ) + assert bedrock.posted_hosts() == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("message", "problem"), + [ + ({"role": "user"}, "Message 0 must have a 'content' field"), + ({"content": "hi"}, "Message 0 must have a 'role' field"), + ], +) +async def test_mantle_counted_model_rejects_a_malformed_message_without_calling_mantle( + message: dict[str, str], problem: str +) -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="bedrock_mantle/anthropic.claude-opus-5-5", + messages=[message], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-5-5", + ) + + assert result is not None + assert (result.error, result.status_code, result.error_message, result.rejected_before_provider_call) == ( + True, + 400, + problem, + True, + ) + assert bedrock.posted_hosts() == () + + @pytest.mark.asyncio async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None: bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 31999ee0379..0799ec15225 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -1663,6 +1663,21 @@ class TestGetBedrockBaseModel: assert get_bedrock_base_model(input_model) == expected +def test_region_only_deployment_takes_no_traffic_from_a_valid_model_group() -> None: + valid_model: Final = "bedrock/us.anthropic.claude-opus-4-6-v1" + router: Final = litellm.Router( + model_list=[ + {"model_name": "group", "litellm_params": {"model": "bedrock/us", "aws_region_name": "us-east-1"}}, + {"model_name": "group", "litellm_params": {"model": valid_model, "aws_region_name": "us-east-1"}}, + ], + ignore_invalid_deployments=True, + ) + + assert [deployment["litellm_params"]["model"] for deployment in router.get_model_list(model_name="group")] == [ + valid_model + ] + + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBedrockModelInfoWrappers: """Tests that BedrockModelInfo methods correctly wrap standalone functions.""" diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index 80e7e793f6e..117d148c287 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -13,6 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx from fastapi import Request, Response from fastapi.responses import StreamingResponse from fastapi.testclient import TestClient @@ -1222,6 +1223,49 @@ def test_count_tokens_rejected_before_counting_logs_the_400_it_answers( ) +def test_count_tokens_on_bedrock_answers_400_for_a_message_without_content_without_calling_bedrock( + monkeypatch: pytest.MonkeyPatch, httpx_transport: None, respx_mock: respx.MockRouter +) -> None: + from fastapi import FastAPI + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.anthropic_endpoints.endpoints import router + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-bedrock-claude", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-zeta-9", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + }, + "model_info": {"id": "count-bedrock-claude-deployment", "token_count_framing": 3}, + } + ] + ), + ) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() + + response: Final = TestClient(app).post( + "/v1/messages/count_tokens", json={"model": "count-bedrock-claude", "messages": [{"role": "user"}]} + ) + + assert response.status_code == 400, response.text + assert response.json()["error"]["type"] == "invalid_request_error" + assert response.json()["error"]["message"] == "Message 0 must have a 'content' field" + assert len(respx_mock.calls) == 0 + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 96d0a2ea8e0..70f5e855e16 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -7659,8 +7659,16 @@ _BEDROCK_COUNT_ALIAS: Final = "count-log-bedrock" _BEDROCK_COUNT_MODEL: Final = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +_BEDROCK_COUNT_BODY: Final = MappingProxyType( + {"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user", "content": "hi"}]} +) + + async def _bedrock_native_count( - monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, aws_reply: httpx.Response + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + aws_reply: httpx.Response, + request_body: Mapping[str, object] = _BEDROCK_COUNT_BODY, ) -> dict[str, object]: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( handle_bedrock_count_tokens, @@ -7700,7 +7708,7 @@ async def _bedrock_native_count( request=request, fastapi_response=Response(), user_api_key_dict=MagicMock(), - request_body={"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user", "content": "hi"}]}, + request_body=dict(request_body), ) @@ -7754,6 +7762,81 @@ async def test_bedrock_native_count_tokens_logs_the_provider_status_when_the_cou ) +@pytest.mark.parametrize( + ("request_body", "problem"), + [ + ({"model": _BEDROCK_COUNT_ALIAS}, "Request must contain content to count tokens"), + ({"model": _BEDROCK_COUNT_ALIAS, "messages": []}, "messages parameter is required for Converse input"), + ({"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user"}]}, "Message 0 must have a 'content' field"), + ({"model": _BEDROCK_COUNT_ALIAS, "messages": [{"content": "hi"}]}, "Message 0 must have a 'role' field"), + ], +) +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_answers_400_for_a_malformed_body_without_calling_bedrock( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + respx_mock: respx.MockRouter, + request_body: Mapping[str, object], + problem: str, +) -> None: + with pytest.raises(HTTPException) as exc_info: + await _bedrock_native_count( + monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}), request_body + ) + + assert (exc_info.value.status_code, exc_info.value.detail) == (400, {"error": problem}) + assert len(respx_mock.calls) == 0 + + +@pytest.mark.parametrize( + ("model", "model_id"), + [ + ("us", "us"), + ("global", "global"), + ("apac", "apac"), + ("bedrock/us", "us"), + ("converse/us", "us"), + ("invoke/us", "us"), + ("bedrock/converse/global", "global"), + ], +) +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_for_a_bare_region_prefix_answers_bedrocks_own_400( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + respx_mock: respx.MockRouter, + model: str, + model_id: str, +) -> None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIATESTACCESSKEY") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + invalid_model: Final = {"message": "The provided model identifier is invalid."} + respx_mock.post(host="bedrock-runtime.us-east-1.amazonaws.com", path=f"/model/{model_id}/count-tokens").mock( + return_value=httpx.Response(400, json=invalid_model) + ) + request: Final = Request({"type": "http", "method": "POST", "path": f"/bedrock/model/{model}/count-tokens", "headers": []}) + + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint=f"model/{model}/count-tokens", + request=request, + fastapi_response=Response(), + user_api_key_dict=MagicMock(), + request_body={"model": model, "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert exc_info.value.status_code == 400 + assert json.loads(exc_info.value.detail["error"]) == invalid_model + assert len(respx_mock.calls) == 1 + + @pytest.mark.asyncio async def test_bedrock_native_count_tokens_without_a_model_logs_the_400_it_answers( token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], @@ -7820,9 +7903,7 @@ async def test_bedrock_native_count_tokens_with_a_region_only_model_answers_the_ request_body={"model": "us"}, ) - assert exc_info.value.detail == { - "error": "CountTokens processing error: Request must contain content to count tokens" - } + assert exc_info.value.detail == {"error": "Request must contain content to count tokens"} assert token_count_log_lines() == ( { "litellm_call_id": "bedrock-native-region-only", diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 2f85d1c7516..694b257b1f9 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -1080,33 +1080,32 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): response=mock_response, ) - with patch.object(handler, "validate_count_tokens_request"): - with patch.object(handler, "_get_aws_region_name", return_value="us-west-2"): - with patch.object(handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}): - with patch.object( - handler, - "get_bedrock_count_tokens_endpoint", - return_value="https://example.com", - ): - with patch.object(handler, "_sign_request", return_value=({}, "{}")): - with patch("litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client") as mock_client: - mock_async_client = AsyncMock() - mock_async_client.post = AsyncMock(side_effect=http_error) - mock_client.return_value = mock_async_client + with patch.object(handler, "_get_aws_region_name", return_value="us-west-2"): + with patch.object(handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}): + with patch.object( + handler, + "get_bedrock_count_tokens_endpoint", + return_value="https://example.com", + ): + with patch.object(handler, "_sign_request", return_value=({}, "{}")): + with patch("litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client") as mock_client: + mock_async_client = AsyncMock() + mock_async_client.post = AsyncMock(side_effect=http_error) + mock_client.return_value = mock_async_client - with pytest.raises(BedrockError) as exc_info: - await handler.handle_count_tokens_request( - request_data={ - "model": "test", - "messages": [{"role": "user", "content": "hello"}], - }, - litellm_params={}, - resolved_model="anthropic.claude-3-sonnet", - ) + with pytest.raises(BedrockError) as exc_info: + await handler.handle_count_tokens_request( + request_data={ + "model": "test", + "messages": [{"role": "user", "content": "hello"}], + }, + litellm_params={}, + resolved_model="anthropic.claude-3-sonnet", + ) - assert exc_info.value.status_code == 403 - # Message should be the raw response text - assert exc_info.value.message == "Forbidden - Invalid credentials" + assert exc_info.value.status_code == 403 + # Message should be the raw response text + assert exc_info.value.message == "Forbidden - Invalid credentials" @pytest.mark.asyncio