scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(budget_reservation): exempt vertex and bedrock count-tokens routes from budget reservation
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Details
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Details
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Details

Browse source
This commit is contained in:
mateo-berri 2026-09-07 18:44:04 -07:00
parent fb7d06da4b
commit 9c980b96d6
2 changed files with 17 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -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(

9
tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
View file

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