From fa4b26fc9fd9357730e7b04b1e3f9b7ce094c464 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 18:03:32 -0700 Subject: [PATCH] fix(proxy): keep the listed price on catch-all wildcard rows and match the edit path to the read path --- litellm/proxy/auth/model_checks.py | 35 ++++- .../model_management_endpoints.py | 45 ++++-- litellm/proxy/proxy_server.py | 31 +--- .../proxy_server/test_routes_model_info.py | 146 ++++++++++++++++++ 4 files changed, 216 insertions(+), 41 deletions(-) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 9c3e0f4b67c..afe9b3e7229 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -13,6 +13,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth +from litellm.proxy.common_utils.token_count_settings import cost_map_entry from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.router import Router from litellm.router_api import RouterRegistry @@ -578,9 +579,37 @@ def wildcard_row_priced_model(row: Mapping[str, object], llm_router: Router | No source_model: Final = _litellm_model(source) if not isinstance(source_pattern, str) or not isinstance(source_model, str) or _litellm_model(row) == source_model: return None - if model_id not in registry.names(listed_name, kind="route_candidates", team_id=_deployment_team_id(row)): - return source_model - return wildcard_filled_model(source_pattern, source_model, listed_name) or source_model + return wildcard_listed_priced_model( + source_pattern, + source_model, + listed_name, + routes_through_wildcard=routes_through_deployment(llm_router, model_id, listed_name, _deployment_team_id(row)), + ) + + +def wildcard_listed_priced_model( + model_name_pattern: str, wildcard_model: str, listed_name: str, routes_through_wildcard: bool +) -> str | None: + if not routes_through_wildcard: + return wildcard_model + filled: Final = wildcard_filled_model(model_name_pattern, wildcard_model, listed_name) + if filled is None: + return wildcard_model + return None if _repeats_listed_provider_prefix(filled, listed_name) else filled + + +def _repeats_listed_provider_prefix(filled_model: str, listed_name: str) -> bool: + listed_provider, separator, _ = listed_name.partition("/") + return separator != "" and filled_model == f"{listed_provider}/{listed_name}" + + +def routes_through_deployment(llm_router: Router, model_id: str, listed_name: str, team_id: str | None) -> bool: + return model_id in llm_router.registry().names(listed_name, kind="route_candidates", team_id=team_id) + + +def model_info_cost_map_entry(model: str) -> Mapping[str, object]: + provider, _, _ = model.partition("/") + return cost_map_entry(model) or cost_map_entry(model.rsplit("/", 1)[-1], custom_llm_provider=provider) def _get_wildcard_models( diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ac4123038ee..152c14ef038 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -15,7 +15,7 @@ import datetime import itertools import json from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence -from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress +from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass from fnmatch import fnmatchcase from functools import partial @@ -61,7 +61,11 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY -from litellm.proxy.auth.model_checks import wildcard_filled_model +from litellm.proxy.auth.model_checks import ( + model_info_cost_map_entry, + routes_through_deployment, + wildcard_listed_priced_model, +) from litellm.proxy.auth.team_grants import team_model_aliases from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.config_sync_pubsub import ( @@ -1266,23 +1270,33 @@ def _ptu_priced_deployment(model_params: Deployment) -> Deployment: ) -def _live_cost_map_model(db_model: Deployment, listed_name: str | None) -> str | None: +def _live_cost_map_model(db_model: Deployment, listed_name: str | None, llm_router: Router | None) -> str | None: stored_model: Final = _decrypted_model(db_model.litellm_params.model) if stored_model is None or listed_name is None: return stored_model - return wildcard_filled_model(db_model.model_name, stored_model, listed_name) or stored_model + model_id: Final = db_model.model_info.id + priced_as: Final = wildcard_listed_priced_model( + db_model.model_name, + stored_model, + listed_name, + routes_through_wildcard=llm_router is None + or model_id is None + or routes_through_deployment(llm_router, model_id, listed_name, db_model.model_info.team_id), + ) + return listed_name if priced_as is None else priced_as def _cost_map_entry( - db_model: Deployment, incoming_model_info: Mapping[str, object], listed_name: str | None + db_model: Deployment, + incoming_model_info: Mapping[str, object], + listed_name: str | None, + llm_router: Router | None, ) -> Mapping[str, object]: base_model: Final = incoming_model_info.get("base_model") - lookup: Final = base_model if isinstance(base_model, str) else _live_cost_map_model(db_model, listed_name) - if lookup is None: - return MappingProxyType({}) - with suppress(Exception): - return MappingProxyType(dict(litellm.get_model_info(model=lookup))) - return MappingProxyType({}) + lookup: Final = ( + base_model if isinstance(base_model, str) else _live_cost_map_model(db_model, listed_name, llm_router) + ) + return MappingProxyType({}) if lookup is None else model_info_cost_map_entry(lookup) LoadedCatalog: TypeAlias = Callable[[], Mapping[str, Mapping[str, object]]] @@ -1305,6 +1319,7 @@ def update_db_model( db_model: Deployment, updated_patch: updateDeployment, loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map, + llm_router: Router | None = None, ) -> PrismaCompatibleUpdateDBModel: if updated_patch.model_info is not None: _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) @@ -1336,7 +1351,7 @@ def update_db_model( incoming_model_info: Final = updated_patch.model_info.model_dump(exclude_none=True) echoed_fields: Final = echoed_cost_map_fields( incoming_model_info, - _cost_map_entry(db_model, incoming_model_info, updated_patch.model_name), + _cost_map_entry(db_model, incoming_model_info, updated_patch.model_name, llm_router), _loaded_catalog_entry(incoming_model_info, loaded_catalog), ) stored_model_info_fields: Final[Mapping[str, object]] = without_server_derived_pricing(incoming_model_info) @@ -1689,6 +1704,7 @@ async def patch_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, write_row=write_row, + llm_router=llm_router, ) if updated_model is None: @@ -2044,6 +2060,7 @@ async def _update_team_model_in_db( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, write_row: Callable[[PrismaCompatibleUpdateDBModel], Awaitable[_RowT]], + llm_router: Router | None = None, ) -> _RowT: """ Handle team model updates with proper alias management. @@ -2088,7 +2105,7 @@ async def _update_team_model_in_db( # No team_id in patch, proceed with standard update if patch_team_id is None: - return await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) + return await write_row(update_db_model(db_model=db_model, updated_patch=patch_data, llm_router=llm_router)) # Determine public model name public_model_name: Final = _get_public_model_name( @@ -2109,7 +2126,7 @@ async def _update_team_model_in_db( ), } ) - row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=team_patch)) + row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=team_patch, llm_router=llm_router)) if is_new_team_assignment: await _setup_new_team_model_assignment( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d73dd268b93..5f545a558b9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -443,6 +443,7 @@ from litellm.proxy.auth.model_checks import ( get_key_models, get_mcp_server_ids, get_team_models, + model_info_cost_map_entry, wildcard_row_priced_model, ) from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy @@ -16149,30 +16150,12 @@ def _enrich_model_info_with_litellm_data( priced_model: Final = ( model if priced_as is None else {**model, "litellm_params": {**model["litellm_params"], "model": priced_as}} ) - litellm_model_info = get_litellm_model_info(model=priced_model) - - # 2nd pass on the model, try seeing if we can find model in litellm model_cost map - if litellm_model_info == {}: - # use litellm_param model_name to get model_info - litellm_params = priced_model.get("litellm_params", {}) - litellm_model = litellm_params.get("model", None) - try: - litellm_model_info = litellm.get_model_info(model=litellm_model) - except Exception: - litellm_model_info = {} - # 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map - if litellm_model_info == {}: - # use litellm_param model_name to get model_info - litellm_params = priced_model.get("litellm_params", {}) - litellm_model = litellm_params.get("model", None) - if litellm_model: - split_model: Final = litellm_model.split("/") - if len(split_model) > 0: - litellm_model = split_model[-1] - try: - litellm_model_info = litellm.get_model_info(model=litellm_model, custom_llm_provider=split_model[0]) - except Exception: - litellm_model_info = {} + priced_litellm_model: Final = priced_model.get("litellm_params", {}).get("model") + litellm_model_info: Final = get_litellm_model_info(model=priced_model) or ( + model_info_cost_map_entry(priced_litellm_model) + if isinstance(priced_litellm_model, str) + else MappingProxyType({}) + ) discovered_model_info: Final = ( llm_router.registry().discovered_model_info(model_info.get("id")) if llm_router is not None diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index 08947a54a0c..35b7df7eda8 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -888,6 +888,152 @@ def test_model_info_prices_a_team_wildcard_row_by_the_route_its_team_takes( _invalidate_model_cost_lowercase_map() +def _catch_all_and_shadowed_wildcard_router(monkeypatch: pytest.MonkeyPatch) -> litellm.Router: + from litellm.proxy.auth import model_checks + + catalog: Final = { + "gpt-5.6-acme": { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "openai", + "mode": "chat", + }, + "fireworks_ai/accounts/fireworks/models/acme-pro": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 2.5e-07, + "litellm_provider": "fireworks_ai", + "mode": "chat", + }, + "deepseek/acme-lite": { + "input_cost_per_token": 7e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepseek", + "mode": "chat", + }, + } + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) + _invalidate_model_cost_lowercase_map() + monkeypatch.setattr( + model_checks, + "get_provider_models", + lambda provider, litellm_params=None: { + "openai": ["gpt-5.6-acme"], + "fireworks_ai": ["fireworks_ai/accounts/fireworks/models/acme-pro"], + }.get(provider, []), + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "*", + "litellm_params": {"model": "openai/*", "api_key": "x"}, + "model_info": {"id": "openai-catch-all"}, + }, + { + "model_name": "my-fw/*", + "litellm_params": {"model": "fireworks_ai/accounts/fireworks/models/*", "api_key": "x"}, + "model_info": {"id": "fireworks-wildcard"}, + }, + { + "model_name": "acme-lite", + "litellm_params": {"model": "deepseek/acme-lite", "api_key": "x"}, + "model_info": {"id": "deepseek-acme-lite"}, + }, + ], + model_group_alias={"my-fw/acme-pro": "acme-lite"}, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + return router + + +def _expanded_wildcard_rows(client: TestClient) -> dict[str, dict[str, object]]: + response: Final = client.get("/model/info") + assert response.status_code == 200, response.text + return { + row["model_name"]: row + for row in response.json()["data"] + if row["model_info"]["id"] != "deepseek-acme-lite" and "*" not in row["model_name"] + } + + +def test_model_info_prices_a_catch_all_row_named_with_its_provider_as_the_listed_model( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """A catch-all over ``openai/*`` lists ``openai/gpt-5.6-acme``. Filling the wildcard with that name repeats + the provider, so the row must keep the price of the model it names, the one calls to ``gpt-5.6-acme`` are + billed at.""" + router: Final = _catch_all_and_shadowed_wildcard_router(monkeypatch) + + with auth_as(): + rows: Final = _expanded_wildcard_rows(client) + + shown: Final = rows["openai/gpt-5.6-acme"]["model_info"] + assert (shown.get("input_cost_per_token"), shown.get("output_cost_per_token")) == (3e-07, 4e-07) + call: Final = router.completion( + model="gpt-5.6-acme", messages=[{"role": "user", "content": "hi"}], mock_response="hello there" + ) + assert call._hidden_params["response_cost"] == pytest.approx( + shown["input_cost_per_token"] * call.usage.prompt_tokens + + shown["output_cost_per_token"] * call.usage.completion_tokens + ) + _invalidate_model_cost_lowercase_map() + + +def test_saving_a_replayed_model_info_wildcard_row_stores_none_of_its_shown_metadata( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """An API client that saves a ``/model/info`` row back unchanged, apart from one edit, must not pin the + cost-map metadata the row showed onto the wildcard, whether the row is a catch-all name carrying its + provider or a name an alias sends to another deployment.""" + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment + + router: Final = _catch_all_and_shadowed_wildcard_router(monkeypatch) + with auth_as(): + rows: Final = _expanded_wildcard_rows(client) + + assert sorted(rows) == ["my-fw/acme-pro", "openai/gpt-5.6-acme"] + sources: Final = { + deployment["model_info"]["id"]: deployment for deployment in router.get_model_list() or [] + } + stored: Final = { + name: json.loads( + update_db_model( + db_model=Deployment( + model_name=sources[row["model_info"]["id"]]["model_name"], + litellm_params=LiteLLM_Params(model=sources[row["model_info"]["id"]]["litellm_params"]["model"]), + model_info=ModelInfo(id=row["model_info"]["id"]), + ), + updated_patch=updateDeployment( + model_name=name, model_info=ModelInfo(**{**row["model_info"], "access_groups": ["prod"]}) + ), + llm_router=router, + )["model_info"] + ) + for name, row in rows.items() + } + shown_metadata: Final = { + name: set(row["model_info"]) - set(sources[row["model_info"]["id"]]["model_info"]) for name, row in rows.items() + } + assert all(shown_metadata.values()), shown_metadata + assert {name: info["access_groups"] for name, info in stored.items()} == { + "my-fw/acme-pro": ["prod"], + "openai/gpt-5.6-acme": ["prod"], + } + assert {name: shown_metadata[name] & set(info) for name, info in stored.items()} == { + "my-fw/acme-pro": set(), + "openai/gpt-5.6-acme": set(), + } + _invalidate_model_cost_lowercase_map() + + # --------------------------------------------------------------------------- # GET /model/info — team BYOK scoping (issue #30983) # ---------------------------------------------------------------------------