diff --git a/litellm/router.py b/litellm/router.py index d01c8443dab..72d0996e513 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5914,17 +5914,35 @@ class Router: ) ## OLD MODEL REGISTRATION ## Kept to prevent breaking changes - _model_name = deployment.litellm_params.model - if deployment.litellm_params.custom_llm_provider is not None: - _model_name = ( - deployment.litellm_params.custom_llm_provider + "/" + _model_name - ) - - litellm.register_model( - model_cost={ - _model_name: _model_info, - } + ## Only register by backend model name when this deployment does NOT + ## have custom pricing in model_info. Otherwise we would overwrite + ## built-in pricing for the backend model and break cost for other + ## deployments sharing the same backend (e.g. two router models + ## both using vertex_ai/gemini-3-pro-preview). + _cost_override_keys = ( + "input_cost_per_token", + "output_cost_per_token", + "input_cost_per_second", + "output_cost_per_second", + "cost_per_query", ) + _has_custom_pricing = any( + _model_info.get(k) is not None for k in _cost_override_keys + ) + if not _has_custom_pricing: + _model_name = deployment.litellm_params.model + if deployment.litellm_params.custom_llm_provider is not None: + _model_name = ( + deployment.litellm_params.custom_llm_provider + + "/" + + _model_name + ) + + litellm.register_model( + model_cost={ + _model_name: _model_info, + } + ) ## Check if LLM Deployment is allowed for this deployment if ( @@ -5953,6 +5971,43 @@ class Router: else: raise e + def _register_backend_model_cost_when_single_custom_pricing(self) -> None: + """ + When exactly one deployment uses a backend and has custom pricing in model_info, + register that backend key in litellm.model_cost so get_model_info(backend_name) + returns the custom pricing. Restores backward compat for single-deployment + custom pricing without overwriting built-in cost when multiple deployments + share the same backend. + """ + _cost_override_keys = ( + "input_cost_per_token", + "output_cost_per_token", + "input_cost_per_second", + "output_cost_per_second", + "cost_per_query", + ) + backend_to_deployments: dict = {} + for m in self.model_list: + lp = m.get("litellm_params") or {} + backend_model = lp.get("model") + if not backend_model: + continue + custom_llm_provider = lp.get("custom_llm_provider") + if custom_llm_provider is not None: + backend_key = f"{custom_llm_provider}/{backend_model}" + else: + backend_key = backend_model + if backend_key not in backend_to_deployments: + backend_to_deployments[backend_key] = [] + backend_to_deployments[backend_key].append(m) + for backend_key, deployments in backend_to_deployments.items(): + if len(deployments) != 1: + continue + model_info = (deployments[0].get("model_info") or {}).copy() + if not any(model_info.get(k) is not None for k in _cost_override_keys): + continue + litellm.register_model(model_cost={backend_key: model_info}) + def _is_auto_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: """ Check if the deployment is an auto-router deployment. @@ -6092,6 +6147,11 @@ class Router: _model_info=_model_info, ) + ## Backward compat: when exactly one deployment uses a backend and has custom + ## pricing, register that backend key so get_model_info(backend_name) returns + ## the custom pricing (e.g. single deployment with 0 or custom rates). + self._register_backend_model_cost_when_single_custom_pricing() + verbose_router_logger.debug( f"\nInitialized Model List {self.get_model_names()}" ) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 08ae804ea80..f075f03a581 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1548,6 +1548,82 @@ def test_get_deployment_model_info_base_model_merge_priority(): print("✓ Base model merge priority test passed!") +def test_two_models_same_backend_custom_pricing_does_not_overwrite_builtin_cost(): + """ + When two router models share the same backend (e.g. vertex_ai/gemini-3-pro-preview), + the first with explicit 0 cost in model_info must not overwrite built-in pricing + for the backend. The second model (no cost in model_info) should still get correct cost. + """ + backend_key = "openai/gpt-4" + builtin_cost = { + "key": "gpt-4", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + "litellm_provider": "openai", + "mode": "chat", + } + # Start with built-in pricing for the shared backend + model_cost_copy = {backend_key: builtin_cost.copy()} + + with patch.object(litellm, "model_cost", model_cost_copy): + router = litellm.Router( + model_list=[ + { + "model_name": "zero-cost-gpt4", + "litellm_params": {"model": "gpt-4"}, + "model_info": { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + { + "model_name": "normal-gpt4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"access_groups": ["default"]}, + }, + ], + ) + # First deployment has custom pricing -> must NOT have registered by backend key + # So built-in cost for backend must be unchanged (or merged for second only) + assert backend_key in litellm.model_cost + assert litellm.model_cost[backend_key].get("input_cost_per_token") != 0.0, ( + "Backend model cost was overwritten by first deployment's 0 cost; " + "second model would see 0 cost incorrectly." + ) + + +def test_single_model_custom_pricing_registers_backend_name(): + """ + Backward compat: when exactly one deployment uses a backend and has custom + pricing, get_model_info(backend_name) should return that custom pricing. + """ + backend_key = "openai/gpt-4" + custom_input = 0.001 + custom_output = 0.002 + model_cost_copy = {} + + with patch.object(litellm, "model_cost", model_cost_copy): + litellm.Router( + model_list=[ + { + "model_name": "my-gpt4", + "litellm_params": { + "model": "gpt-4", + "custom_llm_provider": "openai", + }, + "model_info": { + "input_cost_per_token": custom_input, + "output_cost_per_token": custom_output, + }, + }, + ], + ) + # Single deployment with custom pricing -> backend key should be registered + assert backend_key in litellm.model_cost + assert litellm.model_cost[backend_key].get("input_cost_per_token") == custom_input + assert litellm.model_cost[backend_key].get("output_cost_per_token") == custom_output + + def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): """ Test that _add_deployment_model_to_endpoint_for_llm_passthrough_route correctly strips bedrock provider prefix