fix(proxy): count candidateCount in the budget hold only for models that generate the candidates

A non-Gemini model reached through /v1beta generateContent never receives candidateCount, so a huge count held the whole budget for the length of the call while the model produced one answer

litellm/proxy/spend_tracking/budget_reservation.py

@@ -51,6 +51,7 @@ from litellm.proxy.utils import PrismaClient, ProxyLogging
 from litellm.router import Router
 from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
 from litellm.types.router import DeploymentTypedDict
+from litellm.utils import ProviderConfigManager
 
 
 @dataclass
@@ -1398,12 +1399,16 @@ def _estimate_request_max_cost_for_model(
     llm_router: Router | None,
     input_tokens: int | None = None,
 ) -> float | None:
+    output_multiplier: Final = _get_output_multiplier(
+        _requested_output_counts(request_body=request_body, route=route, model=model, llm_router=llm_router)
+    )
     estimates: Final = [
         _max_cost_for_cost_info(
             request_body=request_body,
             route=route,
             model=model,
             model_info=model_info,
+            output_multiplier=output_multiplier,
             input_tokens=input_tokens,
         )
         for model_info in _get_model_cost_infos(model=model, llm_router=llm_router)
@@ -1417,6 +1422,7 @@ def _max_cost_for_cost_info(
     route: str,
     model: str,
     model_info: Mapping[str, object],
+    output_multiplier: int,
     input_tokens: int | None = None,
 ) -> float | None:
     image_cost: Final = _estimate_image_generation_cost(
@@ -1441,7 +1447,6 @@ 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, 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)
@@ -1698,13 +1703,50 @@ def _requested_output_tokens(request_body: Mapping[str, object], route: str) ->
     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], route: str) -> int:
-    candidates: Final = (
-        _gemini_candidate_counts(request_body)
-        if RouteChecks.is_generate_content_route(route)
-        else (request_body.get("n"), request_body.get("best_of"))
+def _requested_output_counts(
+    request_body: Mapping[str, object],
+    route: str,
+    model: str,
+    llm_router: Router | None,
+) -> tuple[object, ...]:
+    if not RouteChecks.is_generate_content_route(route):
+        return (request_body.get("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 ()
+
+
+def _takes_candidate_count(model: str, llm_router: Router | None) -> bool:
+    deployments: Final = (llm_router.get_model_list(model_name=model) if llm_router is not None else None) or ()
+    deployment_params: Final = tuple(
+        params for deployment in deployments if (params := as_str_mapping(_get_value(deployment, "litellm_params")))
     )
-    return max((1, *(count for count in map(_to_int, candidates) if count is not None)))
+    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:
+    model: Final = litellm_params.get("model")
+    custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
+    if not isinstance(model, str):
+        return True
+    try:
+        backend_model, provider, _, _ = litellm.get_llm_provider(
+            model=model,
+            custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None,
+        )
+        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
+    return native_config is not None
+
+
+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)))
 
 
 def _output_cost(output_tokens: int, output_multiplier: int, output_rate: float) -> float:

tests/unit/proxy/spend_tracking/test_budget_reservation.py

@@ -357,6 +357,43 @@ def test_an_n_too_large_to_price_holds_one_choice(count: object) -> None:
     assert _chat_hold({**GEMINI_CHAT_BODY, "n": count}) == one_choice_hold
 
 
+def _group_generate_content_hold(router: Router, generation_config: dict[str, object]) -> float | None:
+    return estimate_request_max_cost(
+        request_body={"contents": [GEMINI_TEXT_CONTENT], "generationConfig": generation_config},
+        route="/v1beta/models/assistant:generateContent",
+        llm_router=router,
+        input_token_counts={"assistant": 10},
+    )
+
+
+@pytest.mark.parametrize(
+    ("backend_models", "candidates_are_generated"),
+    (
+        (("vertex_ai/claude-haiku-5-5",), False),
+        (("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),
+    ),
+)
+def test_a_candidate_count_is_held_only_when_a_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
+        ]
+    )
+    one_candidate_hold: Final = _group_generate_content_hold(router, {"maxOutputTokens": 100})
+    three_candidate_hold: Final = _group_generate_content_hold(router, {"maxOutputTokens": 100, "candidateCount": 3})
+
+    assert one_candidate_hold is not None
+    assert three_candidate_hold is not None
+    assert (three_candidate_hold > one_candidate_hold) is candidates_are_generated
+    assert three_candidate_hold >= one_candidate_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}),