diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index c711f8b5a62..f6074762d5f 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/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: diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 744d20ab1c7..f5f7b1c6c70 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/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}),