diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 55e59c0f8b0..566fbc886ad 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -44,7 +44,6 @@ from litellm.proxy.common_utils.user_api_key_cache import ( from litellm.proxy.spend_tracking.input_tokens import ( count_input_tokens, count_input_tokens_for_model, - is_gemini_contents_body, tokenizer_models, ) from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope @@ -1441,7 +1440,7 @@ def _max_cost_for_cost_info( if estimated_input_tokens is None or output_tokens is None: return None - output_multiplier: Final = _get_output_multiplier(request_body=request_body) + output_multiplier: Final = _get_output_multiplier(request_body=request_body, route=route) tiered_pricing: Final = model_info.get("tiered_pricing") if isinstance(tiered_pricing, list) and tiered_pricing: tier: Final = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=estimated_input_tokens) @@ -1611,6 +1610,7 @@ async def count_request_input_tokens( request_body=request_body, raw_body=raw_body, models=tuple(dict.fromkeys(chain.from_iterable(counted_as.values()))), + gemini_route=RouteChecks.is_generate_content_route(route), ) return MappingProxyType( { @@ -1631,7 +1631,9 @@ def _estimate_input_tokens( counted: Final = ( input_tokens if input_tokens is not None - else count_input_tokens_for_model(request_body=request_body, model=model) + else count_input_tokens_for_model( + request_body=request_body, model=model, gemini_route=RouteChecks.is_generate_content_route(route) + ) ) if counted is not None: return counted @@ -1654,7 +1656,7 @@ def _estimate_output_tokens( if _is_input_only_route(route=route): return 0 - requested: Final = _requested_output_tokens(request_body) + requested: Final = _requested_output_tokens(request_body, route) # Clamp at min(requested-or-default, model_max-or-default). Two purposes: # (1) Without an explicit cap we still need a finite reservation so the @@ -1678,33 +1680,28 @@ def _nested_field(request_body: Mapping[str, object], config_field: str, field: def _gemini_candidate_counts(request_body: Mapping[str, object]) -> tuple[object, ...]: - if not is_gemini_contents_body(request_body): - return () return tuple( _nested_field(request_body, config_field, field) for config_field, field in product(_GEMINI_CONFIG_FIELDS, ("candidateCount", "candidate_count")) ) -def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: - gemini_output_cap: Final = ( - _nested_field(request_body, "generationConfig", "maxOutputTokens") - if is_gemini_contents_body(request_body) - else None - ) +def _requested_output_tokens(request_body: Mapping[str, object], route: str) -> int | None: + if RouteChecks.is_generate_content_route(route): + gemini_cap: Final = _to_int(_nested_field(request_body, "generationConfig", "maxOutputTokens")) + return gemini_cap if gemini_cap is not None and gemini_cap > 0 else None candidates: Final = ( *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS), _nested_field(request_body, "inferenceConfig", "maxTokens"), - gemini_output_cap, ) - return next((tokens for tokens in map(_to_int, candidates) if tokens is not None and tokens > 0), None) + return next((tokens for tokens in map(_to_int, candidates) if tokens is not None and tokens >= 0), None) -def _get_output_multiplier(request_body: Mapping[str, object]) -> int: +def _get_output_multiplier(request_body: Mapping[str, object], route: str) -> int: candidates: Final = ( - request_body.get("n"), - request_body.get("best_of"), - *_gemini_candidate_counts(request_body), + _gemini_candidate_counts(request_body) + if RouteChecks.is_generate_content_route(route) + else (request_body.get("n"), request_body.get("best_of")) ) return max((1, *(count for count in map(_to_int, candidates) if count is not None))) diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 7f33cda7592..30de3e6f596 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -44,7 +44,6 @@ if TYPE_CHECKING: TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 -_NON_GEMINI_INPUT_FIELDS: Final = ("messages", "prompt", "input", "query", "documents") _GEMINI_SYSTEM_INSTRUCTION_FIELDS: Final = ("systemInstruction", "system_instruction") _GEMINI_CACHED_CONTENT_FIELDS: Final = ("cachedContent", "cached_content") _GEMINI_FUNCTION_RESPONSE_FIELDS: Final = ("functionResponse", "function_response") @@ -103,6 +102,7 @@ async def count_input_tokens( request_body: dict[str, object], raw_body: bytes | None, models: Sequence[str], + gemini_route: bool = False, ) -> Mapping[str, int]: """Input-token count per model, sharing one native count across models that select the same tokenizer.""" @@ -110,13 +110,14 @@ async def count_input_tokens( (model, rust_tokenizer(model)) for model in models ) groups: Final[tuple[RustTokenizer | None, ...]] = tuple(dict.fromkeys(tokenizer for _, tokenizer in tokenizers)) - counted_raw_body: Final = None if holds_compaction_summary(request_body.get("input")) else raw_body + counted_raw_body: Final = None if gemini_route or holds_compaction_summary(request_body.get("input")) else raw_body group_counts: Final = [ await _count_group( request_body=request_body, raw_body=counted_raw_body, tokenizer=tokenizer, models=tuple(model for model, selected in tokenizers if selected == tokenizer), + gemini_route=gemini_route, ) for tokenizer in groups ] @@ -130,14 +131,16 @@ async def _count_group( raw_body: bytes | None, tokenizer: RustTokenizer | None, models: tuple[str, ...], + gemini_route: bool, ) -> Mapping[str, int]: async def python() -> Mapping[str, int]: if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: - return _count_input_tokens_for_models(request_body=request_body, models=models) + return _count_input_tokens_for_models(request_body=request_body, models=models, gemini_route=gemini_route) return await asyncio.to_thread( _count_input_tokens_for_models, request_body=request_body, models=models, + gemini_route=gemini_route, ) if tokenizer is None or raw_body is None: @@ -172,18 +175,30 @@ async def _native_counts( def _count_input_tokens_for_models( request_body: dict[str, object], models: Sequence[str], + gemini_route: bool, ) -> Mapping[str, int]: return MappingProxyType( { model: tokens for model in models - if (tokens := count_input_tokens_for_model(request_body=request_body, model=model)) is not None + if ( + tokens := count_input_tokens_for_model( + request_body=request_body, model=model, gemini_route=gemini_route + ) + ) + is not None } ) -def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: +def count_input_tokens_for_model(request_body: dict, model: str, gemini_route: bool = False) -> int | None: try: + if gemini_route: + return ( + _count_gemini_input_tokens(request_body=request_body, model=model) + if "contents" in request_body + else None + ) if "messages" in request_body: try: return litellm.token_counter( @@ -194,8 +209,6 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: ) except ValueError: return _count_text_tokens(model=model, text=request_body.get("messages")) - if is_gemini_contents_body(request_body): - return _count_gemini_input_tokens(request_body=request_body, model=model) if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: @@ -214,10 +227,6 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: return None -def is_gemini_contents_body(request_body: Mapping[str, object]) -> bool: - return "contents" in request_body and not any(field in request_body for field in _NON_GEMINI_INPUT_FIELDS) - - def _count_gemini_input_tokens(request_body: Mapping[str, object], model: str) -> int | None: system_instruction: Final = next( (request_body[field] for field in _GEMINI_SYSTEM_INSTRUCTION_FIELDS if field in request_body), None diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index f5e6fa9af48..12d509d5198 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import math from datetime import datetime, timedelta, timezone from typing import Final @@ -25,6 +26,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.spend_tracking.budget_reservation import ( _get_team_member_budget_counter, + estimate_request_input_cost, estimate_request_max_cost, get_budget_window_start, release_unbound_budget_reservation, @@ -324,15 +326,67 @@ def test_gemini_output_fields_on_a_chat_body_leave_its_hold_alone(gemini_output_ @pytest.mark.parametrize("cap", (0, -1000)) -def test_an_output_cap_of_zero_or_below_holds_the_default_output(cap: int) -> None: - uncapped_gemini_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}) - uncapped_chat_hold: Final = _chat_hold(GEMINI_CHAT_BODY) - assert uncapped_gemini_hold is not None and uncapped_chat_hold is not None +def test_a_gemini_output_cap_of_zero_or_below_holds_the_default_output(cap: int) -> None: + uncapped_hold: Final = _gemini_hold({"contents": [GEMINI_TEXT_CONTENT]}) + assert uncapped_hold is not None assert ( _gemini_hold({"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"maxOutputTokens": cap}}) - == uncapped_gemini_hold + == uncapped_hold ) - assert _chat_hold({**GEMINI_CHAT_BODY, "max_tokens": cap}) == uncapped_chat_hold + + +def test_a_negative_max_tokens_holds_the_default_output() -> None: + uncapped_hold: Final = _chat_hold(GEMINI_CHAT_BODY) + assert uncapped_hold is not None + assert _chat_hold({**GEMINI_CHAT_BODY, "max_tokens": -1000}) == uncapped_hold + + +def test_a_text_completion_with_max_tokens_zero_holds_only_its_prompt() -> None: + body: Final = {"model": "gemini-3.8-flash", "prompt": GEMINI_PROMPT * 50, "max_tokens": 0} + prompt_cost: Final = estimate_request_input_cost(request_body=body, route="/v1/completions", llm_router=None) + assert prompt_cost is not None and prompt_cost > 0 + assert estimate_request_max_cost(request_body=body, route="/v1/completions", llm_router=None) == prompt_cost + + +GEMINI_LONG_CONTENTS: Final = {"contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT * 500}]}]} + + +@pytest.mark.parametrize( + "stray_fields", + ( + {"prompt": ""}, + {"messages": [], "n": 5}, + {"input": "", "max_tokens": 1}, + {"query": "", "documents": [], "max_completion_tokens": 1}, + {"inferenceConfig": {"maxTokens": 1}}, + ), +) +def test_fields_of_other_apis_on_a_gemini_route_leave_its_hold_alone(stray_fields: dict[str, object]) -> None: + body: Final = {**GEMINI_LONG_CONTENTS, "generationConfig": {"maxOutputTokens": 4000, "candidateCount": 2}} + gemini_hold: Final = _gemini_hold(body) + assert gemini_hold is not None + assert _gemini_hold({**body, **stray_fields}) == gemini_hold + + +@pytest.mark.parametrize("stray_input", ({"prompt": ""}, {"messages": []}, {"input": ""})) +@pytest.mark.parametrize("model", ("gpt-4", "claude-sonnet-5-5")) +@pytest.mark.asyncio +async def test_a_stray_input_field_on_a_gemini_route_never_hides_the_counted_contents( + model: str, stray_input: dict[str, object] +) -> None: + async def counted(body: dict[str, object]) -> dict[str, int]: + return dict( + await count_request_input_tokens( + request_body=body, + route=f"/v1beta/models/{model}:generateContent", + llm_router=None, + raw_body=json.dumps(body).encode(), + ) + ) + + contents_counts: Final = await counted(GEMINI_LONG_CONTENTS) + assert contents_counts[model] > 1000 + assert await counted({**GEMINI_LONG_CONTENTS, **stray_input}) == contents_counts def _tiered_deployment(input_cost_per_token: float) -> Deployment: