diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 31d7d236657..d985075464f 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -183,11 +183,17 @@ _UNBILLED_ROUTES: Final[frozenset[str]] = frozenset( "/openai/v1/responses/input_tokens", } ) -_UNBILLED_ROUTE_SUFFIXES: Final[tuple[str, ...]] = ("/v1/messages/count_tokens", ":countTokens") +_TOKEN_COUNTING_SEGMENTS: Final[frozenset[str]] = frozenset({"count_tokens", "count-tokens"}) +_TOKEN_COUNTING_ACTION: Final = "countTokens" + + +def _is_token_counting_route(route: str) -> bool: + resource, _, action = route.rsplit("/", 1)[-1].partition(":") + return resource in _TOKEN_COUNTING_SEGMENTS or action == _TOKEN_COUNTING_ACTION def _is_unbilled_route(route: str) -> bool: - return route in _UNBILLED_ROUTES or route.endswith(_UNBILLED_ROUTE_SUFFIXES) + return route in _UNBILLED_ROUTES or _is_token_counting_route(route) async def reserve_budget_for_request( diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 5e88268c283..de6c7c2a40a 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -17,6 +17,10 @@ TOKEN_COUNTING_ROUTES: Final = ( "/v1/messages/count_tokens", "/v1beta/models/gemini-3.8-flash:countTokens", "/models/gemini-3.8-flash:countTokens", + "/bedrock/v1/messages/count-tokens", + "/bedrock/model/us.anthropic.claude-sonnet-4-6/count-tokens", + "/vertex_ai/v1/projects/p/locations/us-east5/publishers/anthropic/models/count-tokens:rawPredict", + "/vertex-ai/v1/projects/p/locations/us-east5/publishers/anthropic/models/count-tokens:rawPredict", ) @@ -56,6 +60,11 @@ ANTHROPIC_MESSAGES: Final = [{"role": "user", "content": "hello!!!"}] COUNT_TOKENS_REQUESTS: Final[tuple[tuple[str, dict[str, object]], ...]] = ( ("/v1/messages/count_tokens", {"model": "claude-sonnet-5", "messages": ANTHROPIC_MESSAGES}), ("/v1beta/models/gemini-3.8-flash:countTokens", {"contents": [{"role": "user", "parts": [{"text": "hello!!!"}]}]}), + ( + "/vertex_ai/v1/projects/p/locations/us-east5/publishers/anthropic/models/count-tokens:rawPredict", + {"model": "claude-sonnet-5", "messages": ANTHROPIC_MESSAGES}, + ), + ("/bedrock/v1/messages/count-tokens", {"model": "claude-sonnet-5", "messages": ANTHROPIC_MESSAGES}), ) TINY_BUDGET_KEY_TOKEN: Final = "hashed-count-tokens-key"