fix(proxy): hold a candidate count only when every deployment generates it, and drop n when every deployment drops it

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)))
 

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}),