diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index 4e3059cfe26..3ea51c20cef 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -59,6 +59,7 @@ class TokenCountSettings: tool_framing: int | None = None responses_framing: int | None = None responses_tool_framing: int | None = None + use_token_counting_api: bool = True @property def frames_counts(self) -> bool: @@ -197,6 +198,16 @@ def _tokenizer(value: object) -> TokenCountTokenizer | None: return None +def _use_token_counting_api(value: object) -> bool: + match value: + case bool(): + return value + case str(): + return value.strip().lower() != "false" + case _: + return True + + def resolve_token_count_settings( configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str] ) -> TokenCountSettings: @@ -215,6 +226,7 @@ def resolve_token_count_settings( "token_count_responses_framing", "token_count_tool_framing", ), + use_token_counting_api=_use_token_counting_api(configured.get("use_token_counting_api")), ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fddecb3d6eb..735180e5ec1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15333,6 +15333,8 @@ 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.""" + if not settings.use_token_counting_api and local_estimate_allowed: + return None count_model: Final = settings.count_model counting_model: Final = ( model_to_use diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4c47489e08f..e4aa4f8e0e4 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -252,6 +252,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): token_count_responses_tool_framing: ReadOnly[int | None] token_count_model: ReadOnly[TokenCountModelInfo | None] token_count_tokenizer: ReadOnly[TokenCountTokenizer | None] + use_token_counting_api: ReadOnly[bool | None] class SearchContextCostPerQuery(TypedDict, total=False): @@ -4313,6 +4314,7 @@ DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset( "token_count_responses_framing", "token_count_responses_tool_framing", "token_count_model", + "use_token_counting_api", } ) 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 68910f52b02..3265cd63d95 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -238,6 +238,37 @@ def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> ) == TokenCountSettings(framing=None, count_model=None) +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ({}, True), + ({"use_token_counting_api": True}, True), + ({"use_token_counting_api": False}, False), + ({"use_token_counting_api": "false"}, False), + ({"use_token_counting_api": " False "}, False), + ({"use_token_counting_api": "true"}, True), + ({"use_token_counting_api": "no"}, True), + ({"use_token_counting_api": 0}, True), + ({"use_token_counting_api": None}, True), + ], +) +def test_only_an_explicit_false_turns_the_provider_token_count_off( + configured: dict[str, object], expected: bool +) -> None: + settings: Final = resolve_token_count_settings(configured=configured, default={}, deployment_region=_in_us_west_2) + + assert settings.use_token_counting_api is expected + + +def test_the_token_count_opt_out_stays_off_the_shared_model_entry() -> None: + from litellm.types.utils import shared_backend_model_info + + shared: Final = shared_backend_model_info({"use_token_counting_api": False, "max_input_tokens": 4096}) + + assert "use_token_counting_api" not in shared + assert shared["max_input_tokens"] == 4096 + + @pytest.mark.parametrize( ("value", "expected"), [("kimi_k3", "kimi_k3"), ("cl100k_base", None), (None, None), (True, None)], diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 2f85d1c7516..6e1aa840275 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -5359,3 +5359,82 @@ def test_utils_token_counter_line_stays_one_line_when_the_model_carries_a_newlin assert line["request_model"] == forged_model assert "\n" not in message assert message.startswith(f'token count - litellm_call_id="count log" request_model={json.dumps(forged_model)} ') + + +@pytest.mark.asyncio +async def test_a_deployment_that_opts_out_never_calls_the_provider_counter(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=447)}) + + result: Final = await _ladder( + counter, TokenCountSettings(framing=None, count_model=None, use_token_counting_api=False) + ) + + assert result is None + assert counter.calls == [] + + +@pytest.mark.asyncio +async def test_an_opt_out_still_uses_the_provider_when_no_local_estimate_is_allowed(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=447)}) + + result: Final = await _ladder( + counter, + TokenCountSettings(framing=None, count_model=None, use_token_counting_api=False), + local_estimate_allowed=False, + ) + + assert result is not None + assert (result.total_tokens, result.count_source) == (447, "provider") + + +def _nova_router(model_info: dict[str, object]) -> Router: + return Router( + model_list=[ + { + "model_name": "nova-pro", + "litellm_params": {"model": "bedrock/us.amazon.nova-pro-v1:0", "aws_region_name": "us-east-1"}, + "model_info": model_info, + } + ] + ) + + +@pytest.mark.asyncio +async def test_count_tokens_reads_the_opt_out_from_the_routed_deployment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _nova_router({"use_token_counting_api": False})) + counter: Final = _ScriptedCounter({"us.amazon.nova-pro-v1:0": _provider_reply("nova", total_tokens=900)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "_get_provider_token_counter", + lambda deployment, model_to_use: (counter, model_to_use, "bedrock"), + ) + + result: Final = await count_request_tokens( + request=TokenCountRequest(model="nova-pro", messages=[{"role": "user", "content": "hello"}]), + call_endpoint=True, + messages_format="anthropic", + ) + + assert counter.calls == [] + assert result.count_source == "local_estimate" + assert result.total_tokens > 0 + + +@pytest.mark.asyncio +async def test_count_tokens_still_calls_the_provider_without_the_opt_out(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _nova_router({})) + counter: Final = _ScriptedCounter({"us.amazon.nova-pro-v1:0": _provider_reply("nova", total_tokens=900)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "_get_provider_token_counter", + lambda deployment, model_to_use: (counter, model_to_use, "bedrock"), + ) + + result: Final = await count_request_tokens( + request=TokenCountRequest(model="nova-pro", messages=[{"role": "user", "content": "hello"}]), + call_endpoint=True, + messages_format="anthropic", + ) + + assert len(counter.calls) == 1 + assert (result.total_tokens, result.count_source) == (900, "provider")