Merge eae50e18eb into c5da69cd13

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

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

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

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)],

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