diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 97f28b9fbc5..068cf4f2e67 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -59,8 +59,7 @@ def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: class BedrockTokenCounter(BaseTokenCounter): - """Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens - on bedrock-mantle for Claude models bedrock-runtime answers 400 or 404 for""" + """Counts on the configured Bedrock endpoint, with Mantle retry when Runtime rejects a Claude model""" def __init__( self, @@ -127,17 +126,22 @@ class BedrockTokenCounter(BaseTokenCounter): **({"tools": tools} if tools else {}), **({"system": system} if system else {}), } - resolved_model: Final = get_bedrock_base_model(model_to_use) + use_mantle: Final = model_to_use.startswith("bedrock_mantle/") + resolved_model: Final = get_bedrock_base_model(model_to_use.removeprefix("bedrock_mantle/")) - runtime: Final = await self._count_with( - self._runtime_handler, RUNTIME_TOKENIZER_TYPE, request_data, litellm_params, resolved_model + primary: Final = await self._count_with( + self._mantle_handler if use_mantle else self._runtime_handler, + MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE, + request_data, + litellm_params, + resolved_model, ) outcome: Final = ( await self._count_with( self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model ) - if _runtime_rejected_claude_model(runtime, resolved_model) - else runtime + if not use_mantle and _runtime_rejected_claude_model(primary, resolved_model) + else primary ) match outcome: case CountedTokens(): diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index 6f805616820..22b35bd1001 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -9,7 +9,7 @@ import litellm from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy._types import ProxyException -from litellm.types.token_count import TokenCountSourceFields, TokenCountTokenizer +from litellm.types.token_count import TokenCountProvider, TokenCountSourceFields, TokenCountTokenizer from litellm.types.utils import TokenCountResponse CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404}) @@ -46,6 +46,7 @@ def _with_hint(message: str, hint: str) -> str: class TokenCountModel: model: str region: str | None + litellm_provider: TokenCountProvider | None = None @dataclass(frozen=True, slots=True) @@ -100,9 +101,27 @@ def _non_blank(value: object) -> str | None: return value if isinstance(value, str) and value.strip() else None +def _count_provider(value: object) -> TokenCountProvider | None: + match value: + case "bedrock": + return "bedrock" + case "bedrock_mantle": + return "bedrock_mantle" + case _: + return None + + def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: model: Final = _non_blank(value.get("model")) - return None if model is None else TokenCountModel(model=model, region=_non_blank(value.get("region"))) + return ( + None + if model is None + else TokenCountModel( + model=model, + region=_non_blank(value.get("region")), + litellm_provider=_count_provider(value.get("litellm_provider")), + ) + ) def _geography(region: str) -> str: @@ -112,7 +131,10 @@ def _geography(region: str) -> str: def _within_deployment_geography( count_model: TokenCountModel | None, deployment_region: Callable[[], str] ) -> TokenCountModel | None: - if count_model is None or count_model.region is None or count_model.model.startswith("global."): + if count_model is None or count_model.region is None: + return count_model + model: Final = count_model.model.removeprefix("bedrock/").removeprefix("bedrock_mantle/") + if model.startswith("global."): return count_model if _geography(deployment_region()) == _geography(count_model.region): return count_model diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cf291d65a5b..3fb7043aa2c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14534,9 +14534,14 @@ async def _provider_token_count( Returns None when the count should be estimated locally, or the provider's failure when the framed local estimate stands in for it and that failure is the answer if the estimate cannot read the request.""" count_model: Final = settings.count_model - counted: Final = await ( - count_with(model_to_use, None) if count_model is None else count_with(count_model.model, count_model.region) + counting_model: Final = ( + model_to_use + if count_model is None + else count_model.model + if count_model.litellm_provider is None + else f"{count_model.litellm_provider}/{count_model.model.removeprefix('bedrock/').removeprefix('bedrock_mantle/')}" ) + counted: Final = await count_with(counting_model, None if count_model is None else count_model.region) if counted is None: return None if isinstance(counted, TokenCountResponse): diff --git a/litellm/types/token_count.py b/litellm/types/token_count.py index 3b21401066c..90ccefea177 100644 --- a/litellm/types/token_count.py +++ b/litellm/types/token_count.py @@ -4,11 +4,13 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"] TokenCountTokenizer: TypeAlias = Literal["kimi_k3"] +TokenCountProvider: TypeAlias = Literal["bedrock", "bedrock_mantle"] class TokenCountModelInfo(TypedDict): model: ReadOnly[str] region: ReadOnly[NotRequired[str]] + litellm_provider: ReadOnly[NotRequired[TokenCountProvider]] class TokenCountSourceFields(TypedDict, total=False): diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 45724d0d520..7a8b2c0cc7d 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -1253,9 +1253,16 @@ "region": { "type": "string", "pattern": "\\S" + }, + "litellm_provider": { + "type": "string", + "enum": [ + "bedrock", + "bedrock_mantle" + ] } }, - "description": "Model that counts tokens for this one, and the AWS region its count API runs in." + "description": "Model that counts tokens for this one, its counting provider, and the AWS region its count API runs in." }, "token_count_tokenizer": { "type": "string", 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 8a21027872a..dd0d6429f8a 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 @@ -174,3 +174,50 @@ async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> N "system": "You are a terse assistant.", "tools": tools, } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "host", "tokenizer", "expected_tokens"), + ( + ("bedrock_mantle/anthropic.claude-opus-5-5", MANTLE_HOST, "bedrock_mantle_api", 271), + ("bedrock/anthropic.claude-3-haiku-20240307-v1:0", RUNTIME_HOST, "bedrock_api", 193), + ), +) +async def test_qualified_counting_model_uses_its_endpoint_without_a_failed_probe( + model: str, host: str, tokenizer: str, expected_tokens: int +) -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 193}), mantle=(200, {"input_tokens": 271})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use=model, + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="served-model", + ) + + assert result is not None + assert result.total_tokens == expected_tokens + assert result.tokenizer_type == tokenizer + assert bedrock.posted_hosts() == (host,) + if host == MANTLE_HOST: + assert _JSON_BODY.validate_json(bedrock.requests[0].content)["model"] == model.split("/", 1)[1] + + +@pytest.mark.asyncio +async def test_explicit_mantle_failure_does_not_probe_runtime() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 193}), mantle=(404, UNSUPPORTED)) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="bedrock_mantle/anthropic.claude-opus-5-5", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + ) + + assert result is not None + assert result.error is True + assert result.status_code == 404 + assert result.tokenizer_type == "bedrock_mantle_api" + assert bedrock.posted_hosts() == (MANTLE_HOST,) diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py index 8bd5ced1b58..d7314f1a6d3 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -4,6 +4,7 @@ from typing import Final import pytest +from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy.common_utils.token_count_settings import ( ProviderCountFailure, TokenCountModel, @@ -33,6 +34,13 @@ def _in_us_west_2() -> str: {"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}}, TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")), ), + ( + {"token_count_model": {"model": "custom-model", "litellm_provider": "bedrock_mantle"}}, + TokenCountSettings( + framing=97, + count_model=TokenCountModel(model="custom-model", region=None, litellm_provider="bedrock_mantle"), + ), + ), ( {"token_count_model": {"model": "custom-model"}}, TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), @@ -94,13 +102,15 @@ def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geogr assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region) -def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography() -> None: +@pytest.mark.parametrize("prefix", ("", "bedrock/", "bedrock_mantle/")) +def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography(prefix: str) -> None: + model: Final = f"{prefix}global.anthropic.claude-opus-5-5" settings: Final = resolve_token_count_settings( configured={}, - default={"token_count_model": {"model": "global.anthropic.claude-opus-5-5", "region": "us-east-1"}}, + default={"token_count_model": {"model": model, "region": "us-east-1"}}, deployment_region=lambda: "eu-west-1", ) - assert settings.count_model == TokenCountModel(model="global.anthropic.claude-opus-5-5", region="us-east-1") + assert settings.count_model == TokenCountModel(model=model, region="us-east-1") def test_shipped_australian_count_region_stays_with_an_australian_deployment() -> None: @@ -157,21 +167,24 @@ def test_count_source_fields_name_the_counting_model_only_when_token_count_model assert count_source_fields(_counted(source)) == expected -def test_every_shipped_token_count_model_is_another_mapped_model_of_the_same_provider_in_a_named_region() -> None: +def test_every_shipped_token_count_model_resolves_to_a_mapped_aws_model() -> None: cost_map: Final = json.loads(_COST_MAP.read_text()) + aws_models: Final = frozenset( + get_bedrock_base_model(name) + for name, entry in cost_map.items() + if isinstance(entry, dict) and entry.get("litellm_provider") in ("bedrock", "bedrock_converse") + ) count_models: Final = { name: entry["token_count_model"] for name, entry in cost_map.items() if isinstance(entry, dict) and "token_count_model" in entry } assert count_models - assert { - name: count_model - for name, count_model in count_models.items() - if cost_map.get(count_model["model"], {}).get("litellm_provider") != cost_map[name]["litellm_provider"] - or count_model["model"] == name - or not count_model.get("region") - } == {} + for name, count_model in count_models.items(): + model: Final = count_model["model"] + assert get_bedrock_base_model(model) in aws_models + assert cost_map[name]["litellm_provider"] in ("bedrock", "bedrock_converse") + assert "region" not in count_model or count_model["region"].strip() _UPSTREAM_404: Final = { diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 720739a730b..d17fad7e9b3 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -65,6 +65,7 @@ from litellm.proxy.common_utils.token_count_settings import ( TokenCountModel, TokenCountSettings, cost_map_entry, + resolve_token_count_settings, ) from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request from litellm.proxy.proxy_server import ( @@ -1368,6 +1369,65 @@ async def test_model_without_token_count_model_counts_itself_in_the_deployment_r assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-sonnet-5-5"], ["us-west-2"]) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("provider", "expected_host", "expected_tokens"), + ( + ("bedrock", "bedrock-runtime.us-east-1.amazonaws.com", 193), + ("bedrock_mantle", "bedrock-mantle.us-east-1.api.aws", 271), + ), +) +async def test_token_count_model_provider_selects_the_endpoint_and_preserves_the_public_model( + provider: str, expected_host: str, expected_tokens: int, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + model: Final = "anthropic.claude-opus-5-5" + + def peer(request: httpx.Request) -> httpx.Response: + assert request.url.host == expected_host + if provider == "bedrock_mantle": + assert json.loads(request.content)["model"] == model + return httpx.Response(200, json={"input_tokens": expected_tokens}) + assert request.url.path == f"/model/{model}/count-tokens" + return httpx.Response(200, json={"inputTokens": expected_tokens}) + + settings: Final = resolve_token_count_settings( + configured={}, + default={"token_count_model": {"model": model, "litellm_provider": provider, "region": "us-east-1"}}, + deployment_region=lambda: "us-west-2", + ) + counter: Final = BedrockTokenCounter(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer))) + result: Final = await _provider_token_count( + count_with=partial( + _try_provider_token_count, + counter, + "bedrock", + messages=_COUNT_REQUEST.messages, + contents=_COUNT_REQUEST.contents, + request_model=_COUNT_REQUEST.model, + deployment={ + "litellm_params": { + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + } + }, + tools=_COUNT_REQUEST.tools, + system=_COUNT_REQUEST.system, + ), + model_to_use="us.anthropic.claude-sonnet-5-5", + settings=settings, + local_estimate_allowed=False, + ) + + assert isinstance(result, TokenCountResponse) + assert (result.total_tokens, result.model_used, result.count_source) == ( + expected_tokens, + model, + "token_count_model", + ) + + @pytest.mark.parametrize( ("count_model", "region"), [ @@ -2390,7 +2450,7 @@ async def test_uncountable_model_keeps_the_provider_error_when_local_counting_is @pytest.mark.parametrize( ("model_info", "counted_with"), [ - ({}, ("us.anthropic.claude-opus-5-5", "us-east-1", 447)), + ({}, ("bedrock_mantle/us.anthropic.claude-opus-5-5", "us-east-1", 447)), ( {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}}, ("us.anthropic.claude-opus-5-5", "ap-southeast-4", 447), @@ -2406,6 +2466,7 @@ async def test_configured_token_count_model_overrides_the_shipped_one_and_an_emp { "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=300), "us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447), + "bedrock_mantle/us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447), } ) monkeypatch.setattr( @@ -2444,7 +2505,7 @@ async def test_configured_token_count_model_overrides_the_shipped_one_and_an_emp async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geography( monkeypatch: pytest.MonkeyPatch, aws_region_name: str, counted_in: str ): - counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + counter: Final = _ScriptedCounter({"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) monkeypatch.setattr( litellm.proxy.proxy_server, "llm_router", @@ -2468,7 +2529,7 @@ async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geog request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True ) - assert (counter.calls[0][0], counter.regions[0]) == ("anthropic.claude-opus-5-5", counted_in) + assert (counter.calls[0][0], counter.regions[0]) == ("bedrock_mantle/anthropic.claude-opus-5-5", counted_in) class _SlowVault: @@ -2488,7 +2549,7 @@ async def _longest_stall(done: asyncio.Event) -> float: async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_free(monkeypatch: pytest.MonkeyPatch): from litellm.proxy import proxy_server - counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + counter: Final = _ScriptedCounter({"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) def bedrock_counter(deployment: object, model: str) -> tuple[_ScriptedCounter, str, str]: return counter, model, "bedrock"