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")