scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139

fix(router): preserve custom model_info across model cost map reload and harden pre-call checks

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
shivam 2026-07-21 23:39:03 +00:00
parent 58ff0e32ba
commit 69111a46cb
3 changed files with 212 additions and 92 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/proxy_server.py
View file

@ -6412,6 +6412,8 @@ class ProxyConfig:
model_cost_map_url = litellm.model_cost_map_url
new_model_cost_map = get_model_cost_map(url=model_cost_map_url)
litellm.model_cost = new_model_cost_map
if llm_router is not None:
llm_router.re_register_deployments_in_model_cost()
# Invalidate case-insensitive lookup map since model_cost was replaced
_invalidate_model_cost_lowercase_map()
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that
@ -15590,6 +15592,8 @@ async def reload_model_cost_map(
model_cost_map_url = litellm.model_cost_map_url
new_model_cost_map = get_model_cost_map(url=model_cost_map_url)
litellm.model_cost = new_model_cost_map
if llm_router is not None:
llm_router.re_register_deployments_in_model_cost()
# Invalidate case-insensitive lookup map since model_cost was replaced
_invalidate_model_cost_lowercase_map()
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that

198
litellm/router.py
View file

@ -7447,6 +7447,99 @@ class Router:
if backend_value is not None:
model_info[field] = backend_value
@staticmethod
def _prepare_custom_model_info(deployment: Deployment) -> dict:
"""Build the model_info dict registered into ``litellm.model_cost`` for a
deployment: the deployment's ``model_info`` plus any custom pricing
declared in ``litellm_params``, with missing cache pricing inherited from
the backend model.
"""
model_info = deployment.model_info.model_dump(exclude_none=True)
litellm_params = deployment.litellm_params
for field in CustomPricingLiteLLMParams.model_fields.keys():
field_value = litellm_params.get(field)
if field_value is not None:
model_info[field] = field_value
if model_info.get("input_cost_per_token") is not None:
Router._inherit_builtin_cache_pricing(
model_info=model_info,
backend_model=litellm_params.model,
custom_llm_provider=litellm_params.custom_llm_provider,
)
return model_info
@staticmethod
def _register_deployment_model_cost(
model_id: str | None,
model_name: str,
custom_llm_provider: str | None,
model_info: dict,
) -> None:
"""Register a prepared deployment ``model_info`` into ``litellm.model_cost``.
Registers under two keys: the unique deployment model id (full
``model_info``) and the shared provider/model backend name (cost-map
schema fields only, so one deployment's overrides don't pollute a sibling
sharing the same backend). Invoked when a deployment is added and again
after a model cost map reload replaces ``litellm.model_cost`` wholesale,
so custom pricing and token limits are not lost.
"""
if model_id is not None:
litellm.register_model(model_cost={model_id: model_info})
## OLD MODEL REGISTRATION ## Kept to prevent breaking changes
_model_name = model_name
if custom_llm_provider is not None:
_model_name = custom_llm_provider + "/" + _model_name
_shared_model_info = shared_backend_model_info(model_info)
_existing_shared_mode = (cast("dict | None", litellm.model_cost.get(_model_name, {})) or {}).get("mode")
_deployment_mode = _shared_model_info.get("mode")
# Keep the built-in bridge mode stable for shared backend keys.
# Multiple aliases can point at the same provider/model backend,
# but their deployment-level overrides should not downgrade the
# backend from responses -> chat via last-write-wins registration.
_is_responses_to_chat_downgrade = _existing_shared_mode == "responses" and _deployment_mode == "chat"
_would_clear_existing_mode = _existing_shared_mode is not None and _deployment_mode is None
if _is_responses_to_chat_downgrade or _would_clear_existing_mode:
if _deployment_mode is not None:
verbose_router_logger.warning(
"Router: preserving existing mode=%s for shared backend "
"key %s instead of the deployment-specified mode=%s "
"(prevents alias registration from downgrading the "
"shared backend mode).",
_existing_shared_mode,
_model_name,
_deployment_mode,
)
_shared_model_info["mode"] = _existing_shared_mode
_backend_alias_cost = {_model_name: _shared_model_info}
if "responses/" in _model_name:
_stripped_model_name = _model_name.replace("responses/", "")
_backend_alias_cost[_stripped_model_name] = _shared_model_info
litellm.register_model(model_cost=_backend_alias_cost)
def re_register_deployments_in_model_cost(self) -> None:
"""Re-register every deployment's ``model_info`` into ``litellm.model_cost``.
The proxy price-data reload replaces ``litellm.model_cost`` wholesale,
dropping the per-deployment registrations added at startup. Call this
right after such a reload so custom pricing and token limits (used by
``/model_group/info`` and cost tracking) survive.
"""
for model_id in self.get_model_ids():
deployment = self.get_deployment(model_id=model_id)
if deployment is None:
continue
self._register_deployment_model_cost(
model_id=deployment.model_info.id,
model_name=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
model_info=self._prepare_custom_model_info(deployment),
)
def _create_deployment(
self,
deployment_info: dict,
@ -7483,57 +7576,12 @@ class Router:
)
## REGISTER MODEL INFO IN LITELLM MODEL COST MAP
model_id = deployment.model_info.id
if model_id is not None:
litellm.register_model(
model_cost={
model_id: _model_info,
}
)
## 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
# For the shared backend key, keep only cost-map schema fields
# (minus custom pricing) so that one deployment's pricing overrides
# or custom metadata (id, access_via_team_ids, arbitrary keys)
# don't pollute another deployment sharing the same backend model
# name. Each deployment's full model_info is already stored under
# its unique model_id above.
_shared_model_info = shared_backend_model_info(_model_info)
_existing_shared_mode = (cast(Optional[dict], litellm.model_cost.get(_model_name, {})) or {}).get("mode")
_deployment_mode = _shared_model_info.get("mode")
# Keep the built-in bridge mode stable for shared backend keys.
# Multiple aliases can point at the same provider/model backend,
# but their deployment-level overrides should not downgrade the
# backend from responses -> chat via last-write-wins registration.
# Only preserve in that specific direction so legitimate upgrades
# (e.g. chat -> responses) and unrelated mode changes still apply,
# and so a missing deployment mode does not silently clear the
# existing shared backend mode.
_is_responses_to_chat_downgrade = _existing_shared_mode == "responses" and _deployment_mode == "chat"
_would_clear_existing_mode = _existing_shared_mode is not None and _deployment_mode is None
if _is_responses_to_chat_downgrade or _would_clear_existing_mode:
if _deployment_mode is not None:
verbose_router_logger.warning(
"Router: preserving existing mode=%s for shared backend "
"key %s instead of the deployment-specified mode=%s "
"(prevents alias registration from downgrading the "
"shared backend mode).",
_existing_shared_mode,
_model_name,
_deployment_mode,
)
_shared_model_info["mode"] = _existing_shared_mode
# Always register the (possibly mode-preserved) shared backend info.
_backend_alias_cost = {_model_name: _shared_model_info}
if "responses/" in _model_name:
_stripped_model_name = _model_name.replace("responses/", "")
_backend_alias_cost[_stripped_model_name] = _shared_model_info
litellm.register_model(model_cost=_backend_alias_cost)
self._register_deployment_model_cost(
model_id=deployment.model_info.id,
model_name=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
model_info=_model_info,
)
## Check if LLM Deployment is allowed for this deployment
if self.deployment_is_active_for_environment(deployment=deployment) is not True:
@ -8193,46 +8241,12 @@ class Router:
# initialize client
self._add_deployment(deployment=deployment)
_model_info_dict: dict = deployment.model_info.model_dump(exclude_none=True)
for field in CustomPricingLiteLLMParams.model_fields.keys():
field_value = deployment.litellm_params.get(field)
if field_value is not None:
_model_info_dict[field] = field_value
if _model_info_dict.get("input_cost_per_token") is not None:
Router._inherit_builtin_cache_pricing(
model_info=_model_info_dict,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
# Register custom pricing in litellm.model_cost.
# Mirrors _create_deployment() logic to ensure dynamically-added deployments
# (e.g., loaded from DB) also have their custom pricing registered.
# Without this, _is_model_cost_zero() cannot detect explicitly-configured
# zero-cost models, causing budget checks to block free models.
_model_id = deployment.model_info.id
if _model_id is not None:
litellm.register_model(model_cost={_model_id: _model_info_dict})
## REGISTER MODEL INFO IN LITELLM MODEL COST MAP
## 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
# For the shared backend key, keep only cost-map schema fields
# (minus custom pricing) so that one deployment's pricing overrides
# or custom metadata (id, access_via_team_ids, arbitrary keys)
# don't pollute another deployment sharing the same backend model
# name. Each deployment's full model_info is already stored under
# its unique model_id above (when present).
_shared_model_info = shared_backend_model_info(_model_info_dict)
_backend_alias_cost = {_model_name: _shared_model_info}
if "responses/" in _model_name:
_stripped_model_name = _model_name.replace("responses/", "")
_backend_alias_cost[_stripped_model_name] = _shared_model_info
litellm.register_model(model_cost=_backend_alias_cost)
self._register_deployment_model_cost(
model_id=deployment.model_info.id,
model_name=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
model_info=self._prepare_custom_model_info(deployment),
)
# add to model names
self._add_model_to_list_and_index_map(model=_deployment, model_id=deployment.model_info.id)
@ -10121,8 +10135,8 @@ class Router:
base_model = _model_info.get("base_model", None)
if base_model is None:
base_model = _litellm_params.get("base_model", None)
model_info = self.get_router_model_info(deployment=deployment, received_model_name=model)
_deployment_model = base_model or _litellm_params.get("model", None)
model_info = self.get_router_model_info(deployment=deployment, received_model_name=model)
max_input_tokens = model_info.get("max_input_tokens") if isinstance(model_info, dict) else None
if isinstance(max_input_tokens, int) and has_countable_input:

102
tests/test_litellm/test_router_model_cost_isolation.py
View file

@ -20,6 +20,7 @@ sys.path.insert(
import litellm
from litellm import Router
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.utils import _invalidate_model_cost_lowercase_map
@ -944,3 +945,104 @@ def test_wildcard_zero_cost_request_does_not_poison_named_deployment_pricing():
assert named_cost == pytest.approx(10 * builtin_input_cost)
finally:
_restore_model_cost_entries(model_keys)
def test_re_register_restores_custom_model_info_after_cost_map_reload():
"""LIT-4675 bug 1: the proxy price-data reload replaces litellm.model_cost
wholesale, dropping the router's runtime registration for custom models not
in the built-in cost map. /model_group/info then reports null token limits.
re_register_deployments_in_model_cost() must restore them.
"""
alias = "lit4675-vllm"
backend = "hosted_vllm/lit4675-custom-model"
router = Router(
model_list=[
{
"model_name": alias,
"litellm_params": {
"model": backend,
"api_key": "fake-key",
"api_base": "http://localhost:9999",
},
"model_info": {
"id": "lit4675-vllm-1",
"max_input_tokens": 40000,
"max_output_tokens": 4000,
},
}
],
enable_pre_call_checks=True,
)
info = router.get_model_group_info(alias)
assert info is not None
assert info.max_input_tokens == 40000
assert info.max_output_tokens == 4000
original_model_cost = litellm.model_cost
try:
litellm.model_cost = get_model_cost_map(url=litellm.model_cost_map_url)
_invalidate_model_cost_lowercase_map()
wiped = router.get_model_group_info(alias)
assert wiped is not None
assert wiped.max_input_tokens is None
assert wiped.max_output_tokens is None
router.re_register_deployments_in_model_cost()
_invalidate_model_cost_lowercase_map()
restored = router.get_model_group_info(alias)
assert restored is not None
assert restored.max_input_tokens == 40000
assert restored.max_output_tokens == 4000
finally:
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()
@pytest.mark.asyncio
async def test_pre_call_checks_route_when_deployment_absent_from_cost_map():
"""LIT-4675 bug 2: with enable_pre_call_checks=True, a deployment missing
from litellm.model_cost (e.g. after a price-data reload) must still route.
The invalid-params check has to use the provider-qualified deployment model,
not the bare model group alias, which would raise 'LLM Provider NOT
provided'.
"""
alias = "lit4675-vllm-precall"
backend = "hosted_vllm/lit4675-precall-model"
router = Router(
model_list=[
{
"model_name": alias,
"litellm_params": {
"model": backend,
"api_key": "fake-key",
"api_base": "http://localhost:9999",
},
"model_info": {
"id": "lit4675-precall-1",
"max_input_tokens": 40000,
"max_output_tokens": 4000,
},
}
],
enable_pre_call_checks=True,
)
original_model_cost = litellm.model_cost
try:
litellm.model_cost = get_model_cost_map(url=litellm.model_cost_map_url)
_invalidate_model_cost_lowercase_map()
assert "lit4675-precall-1" not in litellm.model_cost
messages = [{"role": "user", "content": "hi"}]
deployment = await router.async_get_available_deployment(
model=alias,
messages=messages,
request_kwargs={"messages": messages},
)
assert deployment["litellm_params"]["model"] == backend
finally:
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()