fix(proxy): keep same-model wildcard row prices, skip deleted team aliases, use the caller's team

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)

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]:

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
 

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]],