From e5ad3044c76bda30a6ec6c1338dbb732635c7c39 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 01:36:21 -0700 Subject: [PATCH] fix(proxy): keep same-model wildcard row prices, skip deleted team aliases, use the caller's team --- litellm/proxy/auth/model_checks.py | 77 +++++++-- .../proxy/common_utils/model_listing_utils.py | 7 +- .../model_management_endpoints.py | 21 +-- .../proxy_server/test_routes_model_info.py | 146 +++++++++++++++++- 4 files changed, 221 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 331dfc4b19c..ac218507fc5 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -5,7 +5,7 @@ import functools import itertools import re from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from types import MappingProxyType from typing import Final @@ -14,7 +14,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.model_listing_utils import CallerAliases, served_model +from litellm.proxy.common_utils.model_listing_utils import CallerAliases, alias_map, served_model 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 @@ -594,20 +594,69 @@ def wildcard_row_priced_model( 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 - served_name: Final = served_model(listed_name, caller_aliases) - return wildcard_listed_priced_model( - source_pattern, - source_model, - routed_model_name(llm_router, served_name), - routes_through_wildcard=routes_through_deployment(llm_router, model_id, served_name, _deployment_team_id(row)), + return wildcard_deployment_priced_model( + llm_router, + WildcardDeployment(model_id, source_pattern, source_model, _deployment_team_id(row)), + listed_name, + caller_aliases, ) -def wildcard_listed_priced_model( - model_name_pattern: str, wildcard_model: str, routed_name: str, routes_through_wildcard: bool +@dataclass(frozen=True, slots=True) +class WildcardDeployment: + model_id: str + model_name_pattern: str + litellm_model: str + team_id: str | None + + +def wildcard_deployment_priced_model( + llm_router: Router, deployment: WildcardDeployment, listed_name: str, caller_aliases: CallerAliases ) -> str | RoutedElsewhere: - if not routes_through_wildcard: - return ROUTED_ELSEWHERE + registry: Final = llm_router.registry() + served_name: Final = served_model(listed_name, _without_deleted_team_targets(caller_aliases, registry)) + routed_name: Final = routed_model_name(llm_router, served_name) + priced: Final = wildcard_listed_priced_model(deployment.model_name_pattern, deployment.litellm_model, routed_name) + candidates: Final = registry.names( + served_name, kind="route_candidates", team_id=deployment.team_id or caller_aliases.team_id + ) + if deployment.model_id in candidates or _every_candidate_runs(registry, candidates, routed_name, priced): + return priced + return ROUTED_ELSEWHERE + + +def _every_candidate_runs(registry: RouterRegistry, candidates: frozenset[str], routed_name: str, model: str) -> bool: + return bool(candidates) and all( + _model_run_by(_router_deployment(registry, candidate), routed_name) == model for candidate in candidates + ) + + +def _model_run_by(deployment: Mapping[str, object] | None, routed_name: str) -> object: + if deployment is None: + return None + model_name: Final = deployment.get("model_name") + litellm_model: Final = _litellm_model(deployment) + if not isinstance(model_name, str) or not isinstance(litellm_model, str): + return litellm_model + return wildcard_listed_priced_model(model_name, litellm_model, routed_name) + + +def _without_deleted_team_targets(caller_aliases: CallerAliases, registry: RouterRegistry) -> CallerAliases: + deleted_prefix: Final = f"model_name_{caller_aliases.team_id}_" + live_names: Final = registry.index("model_name") + return replace( + caller_aliases, + team=MappingProxyType( + { + alias: target + for alias, target in alias_map(caller_aliases.team).items() + if not target.startswith(deleted_prefix) or target in live_names + } + ), + ) + + +def wildcard_listed_priced_model(model_name_pattern: str, wildcard_model: str, routed_name: str) -> str: filled: Final = wildcard_filled_model(model_name_pattern, wildcard_model, routed_name) if filled is None: return wildcard_model @@ -623,10 +672,6 @@ def routed_model_name(llm_router: Router, listed_name: str) -> str: return llm_router.registry().resolve(listed_name, kind="alias") or 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) diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 2f7d22afdd1..44ee6cd7a0d 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -160,8 +160,8 @@ class ClaudeCodeRoutingNames: @dataclass(frozen=True, slots=True) class CallerAliases: - """The caller's `team` and `key` alias maps, the names `/v1/models` lists for it, and `router`, - its key's or team's `router_settings.model_group_alias`. + """The caller's `team` and `key` alias maps, the names `/v1/models` lists for it, `router`, + its key's or team's `router_settings.model_group_alias`, and `team_id`, the key's team. `rewrite` are the maps `/chat/completions` rewrites its model through after auth's `router` rewrite, in the order it applies them: the team's, the key's in `add_litellm_data_to_request`, then the global `model_alias_map` and the key's again in `common_processing_pre_call_logic`, @@ -170,6 +170,7 @@ class CallerAliases: team: object key: object router: object = None + team_id: str | None = None @property def own(self) -> tuple[object, ...]: @@ -214,7 +215,7 @@ def caller_alias_maps( """Team aliases count only when listing the team the key authenticated as.""" if listed_team_id is not None and listed_team_id != key_team_id: return CallerAliases(None, key_aliases, router_aliases) - return CallerAliases(team_aliases, key_aliases, router_aliases) + return CallerAliases(team_aliases, key_aliases, router_aliases, key_team_id) def alias_map(aliases: object) -> Mapping[str, str]: diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 774f1216084..f9e7e67aece 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -63,9 +63,9 @@ from litellm.proxy._types import ( from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY from litellm.proxy.auth.model_checks import ( GLOBAL_ALIASES_ONLY, + WildcardDeployment, model_info_cost_map_entry, - routed_model_name, - routes_through_deployment, + wildcard_deployment_priced_model, wildcard_listed_priced_model, ) from litellm.proxy.auth.team_grants import team_model_aliases @@ -1280,14 +1280,15 @@ def _live_cost_map_model( if stored_model is None or listed_name is None: return stored_model model_id: Final = db_model.model_info.id - served_name: Final = served_model(listed_name, caller_aliases) - priced: Final = wildcard_listed_priced_model( - db_model.model_name, - stored_model, - served_name if llm_router is None else routed_model_name(llm_router, served_name), - routes_through_wildcard=llm_router is None - or model_id is None - or routes_through_deployment(llm_router, model_id, served_name, db_model.model_info.team_id), + if llm_router is None or model_id is None: + return wildcard_listed_priced_model( + db_model.model_name, stored_model, served_model(listed_name, caller_aliases) + ) + priced: Final = wildcard_deployment_priced_model( + llm_router, + WildcardDeployment(model_id, db_model.model_name, stored_model, db_model.model_info.team_id), + listed_name, + caller_aliases, ) return priced if isinstance(priced, str) else stored_model 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 a20b9f3572e..26ee2cb214f 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -25,7 +25,7 @@ import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy import proxy_server -from litellm.proxy._types import LiteLLM_ProxyModelTable +from litellm.proxy._types import LiteLLM_ProxyModelTable, UserAPIKeyAuth from litellm.utils import _invalidate_model_cost_lowercase_map from .conftest import normalize, route_entry_points_to_stubs # type: ignore[import-not-found] @@ -1076,6 +1076,150 @@ def test_model_info_prices_a_wildcard_row_the_proxy_renames_as_the_model_its_cal _invalidate_model_cost_lowercase_map() +def test_model_info_prices_a_wildcard_row_whose_team_alias_targets_a_deleted_team_deployment( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """Calls skip a team alias whose target team deployment was deleted, so the row that alias names keeps the + price its calls through the wildcard are billed at.""" + from litellm.proxy.litellm_pre_call_utils import _update_model_if_team_alias_exists + + _catch_all_and_shadowed_wildcard_router(monkeypatch) + caller: Final = {"team_id": "team-a", "team_model_aliases": {"my-fw/acme-mini": "model_name_team-a_deleted"}} + + with auth_as(**caller): + rows: Final = _expanded_wildcard_rows(client) + + request: Final = {"model": "my-fw/acme-mini"} + _update_model_if_team_alias_exists(data=request, user_api_key_dict=UserAPIKeyAuth(**caller)) + shown: Final = rows["my-fw/acme-mini"]["model_info"] + assert request["model"] == "my-fw/acme-mini" + assert ( + shown.get("input_cost_per_token"), + shown.get("output_cost_per_token"), + shown.get("litellm_provider"), + ) == (5e-08, 6e-08, "fireworks_ai") + _invalidate_model_cost_lowercase_map() + + +def _fireworks_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["model_info"] + for row in response.json()["data"] + if row["model_info"]["id"] == "fireworks-wildcard" and "*" not in row["model_name"] + } + + +_FIREWORKS_WILDCARD: Final = { + "model_name": "my-fw/*", + "litellm_params": {"model": "fireworks_ai/accounts/fireworks/models/*", "api_key": "x"}, + "model_info": {"id": "fireworks-wildcard"}, +} + + +def test_model_info_keeps_the_price_of_a_wildcard_row_another_deployment_of_the_same_model_serves( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """A deployment pinned to one model a wildcard lists takes that name's calls but runs the model the wildcard + would run, so the wildcard's row keeps the price those calls bill, and saving that row back stores none of it. + A name pinned to another model shows no price.""" + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment + + _catch_all_and_shadowed_wildcard_router(monkeypatch) + router: Final = litellm.Router( + model_list=[ + _FIREWORKS_WILDCARD, + { + "model_name": "my-fw/acme-mini", + "litellm_params": { + "model": "fireworks_ai/accounts/fireworks/models/acme-mini", + "api_key": "x", + "rpm": 10, + }, + "model_info": {"id": "pinned-acme-mini"}, + }, + { + "model_name": "my-fw/acme-pro", + "litellm_params": {"model": "deepseek/acme-lite", "api_key": "x"}, + "model_info": {"id": "pinned-acme-pro"}, + }, + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + + with auth_as(): + rows: Final = _fireworks_wildcard_rows(client) + + shown: Final = rows["my-fw/acme-mini"] + call: Final = router.completion( + model="my-fw/acme-mini", messages=[{"role": "user", "content": "hi"}], mock_response="hello there" + ) + stored: Final = json.loads( + update_db_model( + db_model=Deployment( + model_name="my-fw/*", + litellm_params=LiteLLM_Params(model="fireworks_ai/accounts/fireworks/models/*"), + model_info=ModelInfo(id="fireworks-wildcard"), + ), + updated_patch=updateDeployment( + model_name="my-fw/acme-mini", model_info=ModelInfo(**{**shown, "access_groups": ["prod"]}) + ), + llm_router=router, + )["model_info"] + ) + assert call._hidden_params["model_id"] == "pinned-acme-mini" + 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 + ) + assert stored["access_groups"] == ["prod"] + assert set(stored).isdisjoint(set(shown) - set(router.registry().deployment_dict("fireworks-wildcard")["model_info"])) + assert rows["my-fw/acme-pro"].get("input_cost_per_token") is None + _invalidate_model_cost_lowercase_map() + + +def test_model_info_shows_no_wildcard_price_for_a_name_the_callers_team_publishes( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """Team ``team-a`` publishes ``my-fw/acme-mini`` on its own DeepSeek deployment, so a team-a key's calls to that + name skip the global wildcard and the wildcard's row shows that key no price, while a key outside the team + still sees the Fireworks price its calls bill.""" + _catch_all_and_shadowed_wildcard_router(monkeypatch) + router: Final = litellm.Router( + model_list=[ + _FIREWORKS_WILDCARD, + { + "model_name": "model_name_team-a_acme-mini", + "litellm_params": {"model": "deepseek/acme-lite", "api_key": "x"}, + "model_info": {"id": "team-deepseek", "team_id": "team-a", "team_public_model_name": "my-fw/acme-mini"}, + }, + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + + with auth_as(team_id="team-a"): + team_rows: Final = _fireworks_wildcard_rows(client) + with auth_as(): + other_rows: Final = _fireworks_wildcard_rows(client) + + assert team_rows["my-fw/acme-mini"].get("input_cost_per_token") is None + assert other_rows["my-fw/acme-mini"].get("input_cost_per_token") == 5e-08 + _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]],