diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index 22b35bd1001..4e7dc012802 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -13,6 +13,7 @@ from litellm.types.token_count import TokenCountProvider, TokenCountSourceFields from litellm.types.utils import TokenCountResponse CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404}) +REGION_DENIED_STATUSES: Final = frozenset({401, 403}) TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0 TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0 TOKENIZER_LOAD_MAX_RUNNING: Final = 4 @@ -47,6 +48,7 @@ class TokenCountModel: model: str region: str | None litellm_provider: TokenCountProvider | None = None + configured: bool = False @dataclass(frozen=True, slots=True) @@ -66,6 +68,10 @@ class ProviderCountFailure: def cannot_count_model(self) -> bool: return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES + @property + def denied_in_region(self) -> bool: + return not self.rejected_before_provider_call and self.status_code in REGION_DENIED_STATUSES + def with_hint(self, hint: str) -> "ProviderCountFailure": return replace(self, message=_with_hint(self.message, hint)) @@ -111,7 +117,7 @@ def _count_provider(value: object) -> TokenCountProvider | None: return None -def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: +def _count_model(value: Mapping[str, object], configured: bool) -> TokenCountModel | None: model: Final = _non_blank(value.get("model")) return ( None @@ -120,6 +126,7 @@ def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: model=model, region=_non_blank(value.get("region")), litellm_provider=_count_provider(value.get("litellm_provider")), + configured=configured, ) ) @@ -146,12 +153,12 @@ def _resolved_count_model( ) -> TokenCountModel | None: configured_mapping: Final = as_str_mapping(configured) if configured_mapping is not None: - return _count_model(configured_mapping) + return _count_model(configured_mapping, configured=True) default_mapping: Final = as_str_mapping(default) return ( None if default_mapping is None - else _within_deployment_geography(_count_model(default_mapping), deployment_region) + else _within_deployment_geography(_count_model(default_mapping, configured=False), deployment_region) ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7fcb77e2bbf..64fb2791913 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14853,6 +14853,7 @@ async def _try_native_responses_count( async def _provider_token_count( count_with: Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]], model_to_use: str, + request_model: str, settings: TokenCountSettings, local_estimate_allowed: bool, ) -> Union["TokenCountResponse", ProviderCountFailure, None]: @@ -14875,17 +14876,28 @@ async def _provider_token_count( if count_model is None: return counted.model_copy(update={"count_source": "provider"}) return counted.model_copy(update={"count_source": "token_count_model", "model_used": count_model.model}) - failure: Final = ( - counted - if count_model is None or count_model.region is None - else counted.with_hint( - f"This model counts tokens with {count_model.model} in {count_model.region}. " - "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" - ) + pinned_count_model: Final = None if count_model is None or count_model.region is None else count_model + region_hint: Final = ( + None + if pinned_count_model is None + else f"This model counts tokens with {pinned_count_model.model} in {pinned_count_model.region}. " + "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" ) - if count_model is None and failure.cannot_count_model and settings.framing is not None and local_estimate_allowed: - return failure - raise failure.to_proxy_exception() + failure: Final = counted if region_hint is None else counted.with_hint(region_hint) + framing_can_stand_in: Final = (count_model is None or not count_model.configured) and ( + failure.cannot_count_model or (pinned_count_model is not None and failure.denied_in_region) + ) + if not (framing_can_stand_in and settings.framing is not None and local_estimate_allowed): + raise failure.to_proxy_exception() + if pinned_count_model is not None: + verbose_proxy_logger.warning( + "Token count for %s answered %s from %s in %s, so it falls back to the framed local estimate", + request_model, + failure.status_code, + pinned_count_model.model, + pinned_count_model.region, + ) + return failure @dataclass(frozen=True, slots=True) @@ -15455,6 +15467,7 @@ async def _count_request_tokens( chat_params=MappingProxyType({}) if chat_count is None else chat_count.params, ), model_to_use=model_to_use, + request_model=request.model, settings=settings, local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable, ) diff --git a/tests/integration/providers/test_count_tokens_model_wire.py b/tests/integration/providers/test_count_tokens_model_wire.py index b6ace7ddef3..f15545f4374 100644 --- a/tests/integration/providers/test_count_tokens_model_wire.py +++ b/tests/integration/providers/test_count_tokens_model_wire.py @@ -24,7 +24,10 @@ _SONNET: Final = "us.anthropic.claude-sonnet-5-5" _SONNET_BASE: Final = "anthropic.claude-sonnet-5-5" _OPUS: Final = "us.anthropic.claude-opus-5-5" _OPUS_BASE: Final = "anthropic.claude-opus-5-5" +_OPUS_4_7: Final = "us.anthropic.claude-opus-4-7" _NOVA: Final = "amazon.nova-lite-v1:0" +_RUNTIME_ROUTED: Final = "us.anthropic.claude-sonnet-4-6" +_CONFIGURED_COUNT_REGION: Final = "us-east-2" _KIMI: Final = "us.moonshotai.kimi-k3" _KIMI_BASE: Final = "moonshotai.kimi-k3" _PROFILE_ID: Final = "lit9563count" @@ -168,6 +171,7 @@ def _mantle(request: Request) -> Reply: return _json_reply(200, {"input_tokens": _COUNT_MODEL_COUNT}) error_type: Final = { 401: "authentication_error", + 403: "permission_error", 404: "not_found_error", 429: "rate_limit_error", }[status] @@ -206,9 +210,9 @@ def _mantle_count_requests(wire: Wire) -> tuple[tuple[tuple[str, str], dict[str, return tuple((_signer(request), _JSON_OBJECT.validate_json(request.body)) for request in requests) -def _mantle_body(prompt: str) -> dict[str, JsonValue]: +def _mantle_body(prompt: str, model: str = _OPUS_BASE) -> dict[str, JsonValue]: return { - "model": _OPUS_BASE, + "model": model, "messages": [{"role": "user", "content": prompt}], "system": _SYSTEM, "tools": _SDK_TOOLS, @@ -219,13 +223,13 @@ def _url(port: int) -> str: return f"http://127.0.0.1:{port}" -def _litellm_params(model: str, api_base: str) -> dict[str, JsonValue]: +def _litellm_params(model: str, api_base: str, region: str = _REGION) -> dict[str, JsonValue]: return { "model": f"bedrock/{model}", "api_base": api_base, "aws_access_key_id": _ACCESS_KEY, "aws_secret_access_key": _SECRET_KEY, - "aws_region_name": _REGION, + "aws_region_name": region, } @@ -259,9 +263,13 @@ def counting_proxy(tmp_path_factory: pytest.TempPathFactory, mantle_port: int) - def _deployment( - scenario: Scenario, runtime_port: int, model: str = _SONNET, model_info: Mapping[str, JsonValue] | None = None + scenario: Scenario, + runtime_port: int, + model: str = _SONNET, + model_info: Mapping[str, JsonValue] | None = None, + region: str = _REGION, ) -> str: - return scenario.model(model_info=model_info, api_key=None, **_litellm_params(model, _url(runtime_port))) + return scenario.model(model_info=model_info, api_key=None, **_litellm_params(model, _url(runtime_port), region)) def _count(gateway: Gateway, model: str, status: int, count_model_status: int | None = None) -> MessageTokensCount: @@ -382,6 +390,171 @@ def test_framing_on_one_deployment_leaves_a_sibling_of_the_same_model_raising_th assert _mantle_models(mantle) == () +def test_framing_is_added_to_the_local_estimate_when_a_shipped_runtime_count_model_cannot_count( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model") + assert _shipped_token_count_model(_RUNTIME_ROUTED).get("region") is None + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + framed: Final = _count( + gateway, + _deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_framing": _FRAMING}), + 400, + count_model_status=404, + ) + unframed: Final = _count( + gateway, + _deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_framing": 0}), + 400, + count_model_status=404, + ) + assert (framed.model_extra, unframed.model_extra) == (_LOCAL_SOURCE, _LOCAL_SOURCE), (framed, unframed) + assert framed.input_tokens - unframed.input_tokens == _FRAMING, (framed, unframed) + assert _runtime_models(runtime) == (count_model, count_model) + assert _mantle_models(mantle) == (count_model, count_model) + + +_OPUS_4_7_PINNED: Final = ( + _OPUS_4_7, + _REGION, + _shipped_token_count_model_field(_OPUS_4_7, "region"), + "anthropic.claude-opus-4-7", +) + + +@pytest.mark.parametrize( + ("model", "deployment_region", "count_region", "count_model", "count_model_status"), + [ + (_SONNET, _REGION, _COUNT_REGION, _OPUS_BASE, 404), + (_SONNET_BASE, "ca-central-1", "ca-central-1", _OPUS_BASE, 404), + (*_OPUS_4_7_PINNED, 404), + (*_OPUS_4_7_PINNED, 401), + (*_OPUS_4_7_PINNED, 403), + ], + ids=[ + "sonnet_5_5_in_its_pinned_region", + "sonnet_5_5_outside_the_pinned_geography", + "opus_4_7_in_its_pinned_region", + "opus_4_7_pinned_region_denies_the_key", + "opus_4_7_pinned_region_denies_the_account", + ], +) +def test_framing_is_added_to_the_local_estimate_when_a_shipped_mantle_count_model_cannot_count( + counting_proxy: OwnedProxy, + runtime_port: int, + mantle_port: int, + model: str, + deployment_region: str, + count_region: str, + count_model: str, + count_model_status: int, +) -> None: + gateway: Final = counting_proxy.gateway + assert _shipped_token_count_model_field(model, "litellm_provider") == "bedrock_mantle" + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + framed: Final = _count( + gateway, + _deployment(scenario, runtime_port, model, {"token_count_framing": _FRAMING}, deployment_region), + 200, + count_model_status=count_model_status, + ) + unframed: Final = _count( + gateway, + _deployment(scenario, runtime_port, model, {"token_count_framing": 0}, deployment_region), + 200, + count_model_status=count_model_status, + ) + assert (framed.model_extra, unframed.model_extra) == (_LOCAL_SOURCE, _LOCAL_SOURCE), (framed, unframed) + assert framed.input_tokens - unframed.input_tokens == _FRAMING, (framed, unframed) + assert _runtime_counts(runtime) == () + assert ( + _mantle_count_requests(mantle) + == (((_ACCESS_KEY, count_region), _mantle_body(_prompt(200, count_model_status), count_model)),) * 2 + ) + + +def test_framing_keeps_the_403_when_the_shipped_count_model_runs_in_the_deployment_region( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + with pytest.raises(anthropic.PermissionDeniedError, match=re.escape(f"scripted mantle 403 for {_OPUS_BASE}")): + _count( + gateway, + _deployment(scenario, runtime_port, _SONNET_BASE, {"token_count_framing": _FRAMING}, "ca-central-1"), + 200, + count_model_status=403, + ) + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == ( + ((_ACCESS_KEY, "ca-central-1"), _mantle_body(_prompt(200, 403))), + ) + + +def test_configured_token_count_model_error_stands_with_framing( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model") + model_info: Final[dict[str, JsonValue]] = { + "token_count_framing": _FRAMING, + "token_count_model": {"model": count_model}, + } + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + with pytest.raises(anthropic.NotFoundError, match=re.escape(f"scripted mantle 404 for {count_model}")): + _count( + gateway, _deployment(scenario, runtime_port, _RUNTIME_ROUTED, model_info), 400, count_model_status=404 + ) + assert _runtime_models(runtime) == (count_model,) + assert _mantle_models(mantle) == (count_model,) + + +def test_configured_token_count_model_provider_counts_on_mantle_without_trying_the_runtime( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model") + assert _shipped_token_count_model_field(_RUNTIME_ROUTED, "litellm_provider") == "bedrock" + configured: Final[dict[str, JsonValue]] = { + "model": count_model, + "litellm_provider": "bedrock_mantle", + "region": _CONFIGURED_COUNT_REGION, + } + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + counted: Final = _count( + gateway, _deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_model": configured}), 400 + ) + assert (counted.input_tokens, counted.model_extra) == ( + _COUNT_MODEL_COUNT, + {"litellm_count_source": "token_count_model", "litellm_count_model": count_model}, + ), counted + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == ( + ((_ACCESS_KEY, _CONFIGURED_COUNT_REGION), {**_mantle_body(_prompt(400)), "model": count_model}), + ) + + def test_empty_configured_token_count_model_counts_the_model_itself_in_the_deployment_region( counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int ) -> None: 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 366d8b010d0..7ff32f9d4f3 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -32,22 +32,31 @@ def _in_us_west_2() -> str: ({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), ( {"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")), + TokenCountSettings( + framing=5, + count_model=TokenCountModel(model="custom-model", region="ap-southeast-4", configured=True), + ), ), ( {"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"), + count_model=TokenCountModel( + model="custom-model", region=None, litellm_provider="bedrock_mantle", configured=True + ), ), ), ( {"token_count_model": {"model": "custom-model"}}, - TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + TokenCountSettings( + framing=97, count_model=TokenCountModel(model="custom-model", region=None, configured=True) + ), ), ( {"token_count_model": {"model": "custom-model", "region": " "}}, - TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + TokenCountSettings( + framing=97, count_model=TokenCountModel(model="custom-model", region=None, configured=True) + ), ), ({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)), ( @@ -80,26 +89,26 @@ def test_config_overrides_the_cost_map_default_key_by_key( @pytest.mark.parametrize( - ("deployment_region", "configured", "expected_region"), + ("deployment_region", "configured", "expected"), [ - ("us-west-2", {}, "us-east-1"), - ("eu-central-1", {}, None), - ("ca-central-1", {}, None), - ("us-gov-west-1", {}, None), + ("us-west-2", {}, _OPUS_US_EAST_1), + ("eu-central-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)), + ("ca-central-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)), + ("us-gov-west-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)), ( "eu-central-1", {"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}}, - "us-east-1", + TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1", configured=True), ), ], ) def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography( - deployment_region: str, configured: dict[str, object], expected_region: str | None + deployment_region: str, configured: dict[str, object], expected: TokenCountModel ) -> None: settings: Final = resolve_token_count_settings( configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region ) - assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region) + assert settings.count_model == expected @pytest.mark.parametrize("prefix", ("", "bedrock/", "bedrock_mantle/")) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 5484bc38317..0e8fc6823bd 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -12,6 +12,7 @@ import threading import traceback import uuid from collections.abc import Callable, Coroutine, Mapping +from dataclasses import replace from functools import partial from pathlib import Path from types import MappingProxyType @@ -209,7 +210,6 @@ async def test_gpt_token_counting(): assert response.request_model == "gpt-4" - @pytest.mark.asyncio async def test_token_counter_leaves_the_routers_deployment_without_the_credentials_it_resolves(tmp_path, monkeypatch): """token_counter reads the router's own deployment, so loading the named credential into it would leave the @@ -1685,6 +1685,7 @@ async def _ladder( system=_COUNT_REQUEST.system, ), model_to_use="us.anthropic.claude-sonnet-5-5", + request_model=_COUNT_REQUEST.model, settings=settings, local_estimate_allowed=local_estimate_allowed, ) @@ -1748,6 +1749,7 @@ async def test_token_count_model_provider_selects_the_endpoint_and_preserves_the system=_COUNT_REQUEST.system, ), model_to_use="us.anthropic.claude-sonnet-5-5", + request_model=_COUNT_REQUEST.model, settings=settings, local_estimate_allowed=False, ) @@ -1786,15 +1788,22 @@ async def test_token_count_model_counts_the_same_request_in_its_own_region_witho assert _SONNET_DEPLOYMENT["litellm_params"]["aws_region_name"] == "us-west-2" +@pytest.mark.parametrize( + "count_model", + [None, TokenCountModel(model="us.anthropic.claude-sonnet-5-5", region=None)], + ids=["no_token_count_model", "shipped_token_count_model_in_the_deployment_region"], +) @pytest.mark.parametrize("status_code", [401, 403, 429, 500]) @pytest.mark.asyncio -async def test_errors_other_than_cannot_count_raise_even_with_framing(status_code: int): +async def test_errors_other_than_cannot_count_raise_even_with_framing( + count_model: TokenCountModel | None, status_code: int +): counter: Final = _ScriptedCounter( {"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=status_code)} ) with pytest.raises(ProxyException) as raised: - await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + await _ladder(counter, TokenCountSettings(framing=10, count_model=count_model)) assert (raised.value.code, raised.value.message) == (str(status_code), f"scripted {status_code} for sonnet") @@ -1840,28 +1849,97 @@ async def test_uncountable_model_without_framing_raises(): assert raised.value.code == "404" +@pytest.mark.parametrize( + ("count_model", "counted_as"), + [ + (None, "us.anthropic.claude-sonnet-5-5"), + ( + TokenCountModel(model="anthropic.claude-sonnet-5-5", region=None, litellm_provider="bedrock"), + "bedrock/anthropic.claude-sonnet-5-5", + ), + (_OPUS_IN_US_EAST_1, "us.anthropic.claude-opus-5-5"), + ], + ids=[ + "no_token_count_model", + "shipped_token_count_model_in_the_deployment_region", + "shipped_token_count_model_pinned_to_a_region", + ], +) @pytest.mark.asyncio -async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate(): - counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) +async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate( + count_model: TokenCountModel | None, counted_as: str +): + counter: Final = _ScriptedCounter({counted_as: _provider_reply("sonnet", status_code=400)}) - result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=None)) + result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=count_model)) assert isinstance(result, ProviderCountFailure) - assert len(counter.calls) == 1 + assert [call[0] for call in counter.calls] == [counted_as] -@pytest.mark.parametrize("status_code", [400, 404]) +@pytest.mark.parametrize("status_code", [401, 403]) @pytest.mark.asyncio -async def test_token_count_model_error_stands_even_with_framing(status_code: int): +async def test_framing_stands_in_when_a_shipped_token_count_model_region_denies_the_account(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + result: Final = await _ladder(counter, TokenCountSettings(framing=50, count_model=_OPUS_IN_US_EAST_1)) + + assert isinstance(result, ProviderCountFailure) + assert result.status_code == status_code + + +@pytest.mark.parametrize( + ("count_model", "region_hints"), + [ + (replace(_OPUS_IN_US_EAST_1, configured=True), 1), + (replace(_OPUS_IN_US_EAST_1, region=None, configured=True), 0), + ], + ids=["pinned_to_a_region", "in_the_deployment_region"], +) +@pytest.mark.parametrize("status_code", [400, 401, 403, 404]) +@pytest.mark.asyncio +async def test_configured_token_count_model_error_stands_even_with_framing( + count_model: TokenCountModel, region_hints: int, status_code: int +): counter: Final = _ScriptedCounter( {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} ) with pytest.raises(ProxyException) as raised: - await _ladder(counter, TokenCountSettings(framing=78, count_model=_OPUS_IN_US_EAST_1)) + await _ladder(counter, TokenCountSettings(framing=78, count_model=count_model)) assert raised.value.code == str(status_code) - assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message + assert raised.value.message.startswith(f"scripted {status_code} for opus") + assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints + + +@pytest.mark.parametrize( + ("count_model", "region_hints"), + [(_OPUS_IN_US_EAST_1, frozenset({True})), (replace(_OPUS_IN_US_EAST_1, region=None), frozenset())], + ids=["pinned_to_a_region", "in_the_deployment_region"], +) +@pytest.mark.asyncio +async def test_framing_standing_in_for_a_pinned_token_count_model_warns_with_its_model_and_region( + count_model: TokenCountModel, region_hints: frozenset[bool], caplog: pytest.LogCaptureFixture +): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=404)}) + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + result: Final = await _ladder(counter, TokenCountSettings(framing=50, count_model=count_model)) + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + assert isinstance(result, ProviderCountFailure) + stand_in_warnings: Final = [ + r.getMessage() for r in caplog.records if "falls back to the framed local estimate" in r.getMessage() + ] + assert { + f"for {_COUNT_REQUEST.model} answered 404 from us.anthropic.claude-opus-5-5 in us-east-1" in message + for message in stand_in_warnings + } == region_hints @pytest.mark.asyncio @@ -2837,7 +2915,9 @@ 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({"bedrock_mantle/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", @@ -2888,7 +2968,9 @@ class _VaultWaitingForLoopProgress: 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({"bedrock_mantle/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" @@ -3653,7 +3735,11 @@ async def test_responses_count_sends_openai_the_provider_file_id_inference_sends model_list=[ { "model_name": "gpt-files", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://count.test/v1"}, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": "https://count.test/v1", + }, "model_info": {"id": "deployment-a"}, } ] @@ -3833,9 +3919,7 @@ async def test_responses_count_on_claude_rejects_forced_tool_choice_without_drop monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_forced_tool_use": False}) with pytest.raises(ProxyException) as rejected: - await _claude_count( - monkeypatch, {}, {"input": "Weather?", "tools": [_WEATHER_TOOL], "tool_choice": "required"} - ) + await _claude_count(monkeypatch, {}, {"input": "Weather?", "tools": [_WEATHER_TOOL], "tool_choice": "required"}) assert (rejected.value.code, "forced tool use" in rejected.value.message) == ("400", True)