From ae4f9e72ac55d37bdcd5abdf804a3853039941c0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:32:44 -0700 Subject: [PATCH] fix(proxy): hold a candidate count only when every deployment generates it, and drop n when every deployment drops it --- .../spend_tracking/budget_reservation.py | 52 +++++++++++--- .../spend_tracking/test_budget_reservation.py | 69 ++++++++++++++++--- 2 files changed, 102 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 962b819b68e..76f0d810e6c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -4,7 +4,7 @@ import asyncio import json import math import time -from collections.abc import AsyncIterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone from itertools import chain, product @@ -1730,41 +1730,71 @@ def _requested_output_counts( llm_router: Router | None, ) -> tuple[object, ...]: if not RouteChecks.is_generate_content_route(route): - return (request_body.get("n"), request_body.get("best_of")) + n: Final = request_body.get("n") + dropped: Final = n is not None and _every_deployment(_drops_n, _deployment_params(model, llm_router)) + return (None if dropped else n, request_body.get("best_of")) candidate_counts: Final = _gemini_candidate_counts(request_body) if all(count is None for count in candidate_counts): return () - return candidate_counts if _takes_candidate_count(model=model, llm_router=llm_router) else () + generated: Final = _every_deployment(_serves_generate_content_natively, _deployment_params(model, llm_router)) + return candidate_counts if generated else () -def _takes_candidate_count(model: str, llm_router: Router | None) -> bool: - deployments: Final = (llm_router.registry().deployments(model) if llm_router is not None else None) or () - deployment_params: Final = tuple( +def _deployment_params(model: str, llm_router: Router | None) -> tuple[Mapping[str, object], ...]: + if llm_router is None: + return ({"model": model},) + deployments: Final = llm_router.registry().deployments(model) or () + return tuple( params for deployment in deployments if (params := as_str_mapping(_get_value(deployment, "litellm_params"))) ) - return any(_serves_generate_content_natively(params) for params in deployment_params or ({"model": model},)) -def _serves_generate_content_natively(litellm_params: Mapping[str, object]) -> bool: +def _every_deployment( + check: Callable[[Mapping[str, object]], bool], deployment_params: tuple[Mapping[str, object], ...] +) -> bool: + return bool(deployment_params) and all(map(check, deployment_params)) + + +def _resolved_provider(litellm_params: Mapping[str, object]) -> tuple[str, str] | None: model: Final = litellm_params.get("model") custom_llm_provider: Final = litellm_params.get("custom_llm_provider") if not isinstance(model, str): - return True + return None try: backend_model, provider, _, _ = litellm.get_llm_provider( model=model, custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, ) + except litellm.BadRequestError: + return None + return backend_model, provider + + +def _serves_generate_content_natively(litellm_params: Mapping[str, object]) -> bool: + resolved: Final = _resolved_provider(litellm_params) + if resolved is None: + return False + backend_model, provider = resolved + try: native_config: Final = ProviderConfigManager.get_provider_google_genai_generate_content_config( model=backend_model, provider=litellm.LlmProviders(provider), litellm_params={key: litellm_params[key] for key in ("base_model", "api_base") if key in litellm_params}, ) - except (litellm.BadRequestError, ValueError): - return True + except ValueError: + return False return native_config is not None +def _drops_n(litellm_params: Mapping[str, object]) -> bool: + resolved: Final = _resolved_provider(litellm_params) + if resolved is None: + return False + backend_model, provider = resolved + supported_params: Final = litellm.get_supported_openai_params(model=backend_model, custom_llm_provider=provider) + return supported_params is not None and "n" not in supported_params + + def _get_output_multiplier(requested_counts: tuple[object, ...]) -> int: return max((1, *(count for count in map(_to_int, requested_counts) if count is not None))) diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 10358d7e6a7..c2627f9faf5 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -366,6 +366,15 @@ def _group_generate_content_hold(router: Router, generation_config: dict[str, ob ) +def _assistant_router(*backend_models: str) -> Router: + return Router( + model_list=[ + {"model_name": "assistant", "litellm_params": {"model": model, "api_key": "sk-test"}} + for model in backend_models + ] + ) + + @pytest.mark.parametrize( ("backend_models", "candidates_are_generated"), ( @@ -373,18 +382,15 @@ def _group_generate_content_hold(router: Router, generation_config: dict[str, ob (("anthropic/claude-haiku-5-5",), False), (("gemini/gemini-3.8-flash",), True), (("vertex_ai/gemini-3.8-flash",), True), - (("vertex_ai/claude-haiku-5-5", "gemini/gemini-3.8-flash"), True), + (("gemini/gemini-3.8-flash", "vertex_ai/gemini-3.8-flash"), True), + (("vertex_ai/claude-haiku-5-5", "gemini/gemini-3.8-flash"), False), + (("gemini/gemini-3.8-flash", "anthropic/claude-haiku-5-5"), False), ), ) -def test_a_candidate_count_is_held_only_when_a_deployment_generates_the_candidates( +def test_a_candidate_count_is_held_only_when_every_deployment_generates_the_candidates( backend_models: tuple[str, ...], candidates_are_generated: bool ) -> None: - router: Final = Router( - model_list=[ - {"model_name": "assistant", "litellm_params": {"model": model, "api_key": "sk-test"}} - for model in backend_models - ] - ) + router: Final = _assistant_router(*backend_models) one_candidate_hold: Final = _group_generate_content_hold(router, {"maxOutputTokens": 100}) three_candidate_hold: Final = _group_generate_content_hold(router, {"maxOutputTokens": 100, "candidateCount": 3}) @@ -394,6 +400,53 @@ def test_a_candidate_count_is_held_only_when_a_deployment_generates_the_candidat assert three_candidate_hold >= one_candidate_hold +def test_a_candidate_count_for_a_gemini_name_no_deployment_serves_holds_one_candidate() -> None: + router: Final = _assistant_router("anthropic/claude-haiku-5-5") + + def hold(generation_config: dict[str, object]) -> float | None: + return estimate_request_max_cost( + request_body={"contents": [GEMINI_TEXT_CONTENT], "generationConfig": generation_config}, + route=GEMINI_ROUTE, + llm_router=router, + input_token_counts={"gemini-3.8-flash": 10}, + ) + + one_candidate_hold: Final = hold({"maxOutputTokens": 100}) + assert one_candidate_hold is not None + assert hold({"maxOutputTokens": 100, "candidateCount": 3}) == one_candidate_hold + + +@pytest.mark.parametrize( + ("backend_models", "choices_are_generated"), + ( + (("anthropic/claude-haiku-5-5",), False), + (("vertex_ai/claude-haiku-5-5", "anthropic/claude-haiku-5-5"), False), + (("gemini/gemini-3.8-flash",), True), + (("anthropic/claude-haiku-5-5", "gemini/gemini-3.8-flash"), True), + ), +) +def test_an_n_is_held_unless_every_deployment_drops_it( + backend_models: tuple[str, ...], choices_are_generated: bool +) -> None: + router: Final = _assistant_router(*backend_models) + + def hold(output_fields: dict[str, object]) -> float | None: + return estimate_request_max_cost( + request_body={"model": "assistant", "messages": GEMINI_CHAT_BODY["messages"], **output_fields}, + route="/v1/chat/completions", + llm_router=router, + input_token_counts={"assistant": 10}, + ) + + one_choice_hold: Final = hold({"max_tokens": 100}) + three_choice_hold: Final = hold({"max_tokens": 100, "n": 3}) + + assert one_choice_hold is not None + assert three_choice_hold is not None + assert (three_choice_hold > one_choice_hold) is choices_are_generated + assert three_choice_hold >= one_choice_hold + + HOLDS_FAR_OVER_A_TINY_BUDGET: Final = ( (GEMINI_ROUTE, {"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"candidateCount": 1e300}}), ("/v1/chat/completions", {**GEMINI_CHAT_BODY, "n": 1e300}),