diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 7f2aa33246d..07a0b3e9377 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -6,18 +6,19 @@ import copy import functools import itertools import re -from collections.abc import Callable, Collection, Mapping, Sequence +from collections.abc import Callable, Collection, Iterable, Mapping, Sequence from collections.abc import Set as AbstractSet from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, replace from types import MappingProxyType -from typing import Final, TypeVar +from typing import Final, TypeAlias, TypeVar import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model -from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter +from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils 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, alias_map, dispatched_model @@ -429,6 +430,71 @@ def _ids_routing_to_catalog_models( ) +def _names_whole_models(name_pattern: str) -> bool: + return name_pattern.count("*") == 1 and name_pattern.endswith("/*") + + +@dataclass(frozen=True, slots=True) +class _Rivals: + direct: tuple[re.Pattern[str], ...] + under_a_provider: tuple[re.Pattern[str], ...] + + +RivalPatterns: TypeAlias = Callable[[str], _Rivals | None] + + +def _rival_may_take(listed_id: str, rivals: _Rivals) -> bool: + if any(rival.match(listed_id) for rival in rivals.direct): + return True + if not rivals.under_a_provider: + return False + provider_request: Final = f"{inferred_provider(listed_id)}/{listed_id}" + return any(rival.match(provider_request) for rival in rivals.under_a_provider) + + +def _routes_through_alias( + name_pattern: str, + rivals: _Rivals | None, + routes_to_model: Callable[[str, str], bool], +) -> Callable[[str], bool]: + alias_regex: Final = _wildcard_name_regex(name_pattern) + + def routes_through(listed_id: str) -> bool: + alias_match: Final = alias_regex.match(listed_id) + if alias_match is None or alias_match.group(1) == AZURE_MODEL_PLACEHOLDER: + return False + no_rival_matches: Final = rivals is not None and not _rival_may_take(listed_id, rivals) + return no_rival_matches or routes_to_model(listed_id, alias_match.group(1)) + + return routes_through + + +def _reads_as_another_catalog_id(listed_id: str, catalog_id: str, name_scope: str, catalog: AbstractSet[str]) -> bool: + read_as: Final = listed_id.removeprefix(name_scope) + return read_as != catalog_id and read_as in catalog + + +def _ids_routing_through_catch_all( + name_pattern: str, catalog_ids: Sequence[str], routes_through: Callable[[str], bool] +) -> tuple[str, ...]: + if name_pattern.count("*") != 1: + return () + name_scope: Final = "".join(name_pattern.partition("*")[0].rpartition("/")[:2]) + catalog: Final = frozenset(catalog_ids) + named_ids: Final = ((name_pattern.replace("*", catalog_id), catalog_id) for catalog_id in catalog_ids) + listed_ids: Final = ( + listed_id + for listed_id, catalog_id in named_ids + if not _reads_as_another_catalog_id(listed_id, catalog_id, name_scope, catalog) and routes_through(listed_id) + ) + return tuple(dict.fromkeys(listed_ids)) + + +def _catch_all_catalog(litellm_params: LiteLLM_Params | None) -> Sequence[str]: + pinned_provider: Final = litellm_params.custom_llm_provider if litellm_params is not None else None + return get_provider_models(provider=pinned_provider or "*", litellm_params=litellm_params) or () + + def _catalog_ids_a_wider_name_matches( name_pattern: str, deployment_pattern: str, @@ -477,6 +543,71 @@ def _deployed_by_name(llm_router: Router | None) -> Callable[[str, str | None], return lambda model, team_id: model in deployed_names or (team_id, model) in team_deployed_names +def _literal_start(regex: str) -> str: + return regex.partition("(.*)")[0] + + +def _prefixes_overlap(first: str, second: str) -> bool: + return first.startswith(second) or second.startswith(first) + + +def _may_match_directly(own_regex: str, regex: str) -> bool: + return regex != own_regex and _prefixes_overlap(_literal_start(regex), _literal_start(own_regex)) + + +_UNINFERRED_PROVIDER: Final = str(None) + + +def _may_name_a_provider(escaped_provider: str) -> bool: + provider: Final = re.sub(r"\\(.)", r"\1", escaped_provider) + return provider == _UNINFERRED_PROVIDER or inferred_provider(f"{provider}/{provider}") == provider + + +def _may_match_under_a_provider(own_regex: str, regex: str) -> bool: + provider, has_provider, rest = _literal_start(regex).partition("/") + if not has_provider: + return True + return _prefixes_overlap(rest, _literal_start(own_regex)) and _may_name_a_provider(provider) + + +def _can_take_ids_of(own_regex: str, regex: str) -> bool: + return _may_match_directly(own_regex, regex) and PatternUtils.calculate_pattern_specificity( + regex + ) >= PatternUtils.calculate_pattern_specificity(own_regex) + + +def _compiled(regexes: Iterable[str]) -> tuple[re.Pattern[str], ...]: + return tuple(re.compile(regex) for regex in regexes) + + +def _pattern_regexes(pattern_router: PatternMatchRouter | None) -> tuple[str, ...]: + if pattern_router is None: + return () + return tuple(pattern_router.patterns) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped deployment lists + + +def _rival_patterns(llm_router: Router | None, team_id: str | None) -> RivalPatterns: + if llm_router is None: + return lambda _name_pattern: None + registry: Final = llm_router.registry() + router_regexes: Final = _pattern_regexes(registry.patterns()) + team_regexes: Final = _pattern_regexes(registry.patterns(team=True).get(team_id) if team_id is not None else None) + + def rivals(name_pattern: str) -> _Rivals | None: + own_regex: Final = PatternMatchRouter().pattern_to_regex(name_pattern) + if own_regex in team_regexes: + return _Rivals(_compiled(regex for regex in team_regexes if _can_take_ids_of(own_regex, regex)), ()) + if own_regex not in router_regexes: + return None + return _Rivals( + _compiled(regex for regex in router_regexes if _can_take_ids_of(own_regex, regex)) + + _compiled(regex for regex in team_regexes if _may_match_directly(own_regex, regex)), + _compiled(regex for regex in team_regexes if _may_match_under_a_provider(own_regex, regex)), + ) + + return functools.cache(rivals) + + def _routes_to_model( routed_models: Callable[[str, str | None], frozenset[str]] | None, team_id: str | None ) -> Callable[[str, str], bool]: @@ -495,10 +626,11 @@ def get_known_models_from_wildcard( wildcard_model: str, litellm_params: LiteLLM_Params | None = None, routes_to_model: Callable[[str, str], bool] = lambda _listed_id, _model: True, + rival_patterns: RivalPatterns = lambda _name_pattern: None, ) -> list[str]: provider_id_prefix: Final = _custom_llm_provider_prefix(wildcard_model, litellm_params) if litellm_params is None or provider_id_prefix is None: - return list(_known_models_from_wildcard(wildcard_model, litellm_params, routes_to_model)) + return list(_known_models_from_wildcard(wildcard_model, litellm_params, routes_to_model, rival_patterns)) return [ listed_id.removeprefix(provider_id_prefix) for listed_id in _known_models_from_wildcard( @@ -507,6 +639,7 @@ def get_known_models_from_wildcard( lambda listed_id, model: routes_to_model( listed_id.removeprefix(provider_id_prefix), model.removeprefix(provider_id_prefix) ), + rival_patterns, ) ] @@ -532,6 +665,7 @@ def _known_models_from_wildcard( wildcard_model: str, litellm_params: LiteLLM_Params | None, routes_to_model: Callable[[str, str], bool], + rival_patterns: RivalPatterns, ) -> tuple[str, ...]: wildcard_model_to_expand: Final = ( litellm_params.model @@ -562,6 +696,16 @@ def _known_models_from_wildcard( litellm_params = _hydrate_litellm_credential_name(litellm_params) + routes_through_catch_all: Final = ( + _routes_through_alias(wildcard_model, rival_patterns(wildcard_model), routes_to_model) + if litellm_params is not None and litellm_params.model == "*" + else None + ) + if routes_through_catch_all is not None and not _names_whole_models(wildcard_model): + return _ids_routing_through_catch_all( + wildcard_model, _catch_all_catalog(litellm_params), routes_through_catch_all + ) + wildcard_models = get_provider_models(provider=provider, litellm_params=litellm_params) if wildcard_models is None: @@ -588,9 +732,12 @@ def _known_models_from_wildcard( # add model prefix to wildcard models wildcard_models = [f"{model_prefix}{model.removeprefix(provider_id_prefix)}" for model in wildcard_models] - return tuple( + listed_ids: Final = tuple( dict.fromkeys(_under_wildcard_prefix(model, wildcard_provider_prefix, provider) for model in wildcard_models) ) + if routes_through_catch_all is None: + return listed_ids + return tuple(listed_id for listed_id in listed_ids if routes_through_catch_all(listed_id)) def _under_wildcard_prefix(model: str, wildcard_provider_prefix: str, listed_provider: str) -> str: @@ -647,6 +794,7 @@ def expand_wildcard_deployments_for_model_info( """ expanded: Final[list[dict[str, object]]] = [] routed_models: Final = _cached_wildcard_routing(llm_router) + rival_patterns_of_team: Final = functools.cache(functools.partial(_rival_patterns, llm_router)) deployed_by_name: Final = _deployed_by_name(llm_router) for deployment in deployments: model_name = str(deployment.get("model_name") or "") @@ -668,6 +816,7 @@ def expand_wildcard_deployments_for_model_info( wildcard_model=model_name, litellm_params=litellm_params, routes_to_model=_routes_to_model(routed_models, _deployment_team_id(deployment)), + rival_patterns=rival_patterns_of_team(_deployment_team_id(deployment)), ) if not expanded_names: expanded.append(deployment) @@ -826,6 +975,7 @@ def _get_wildcard_models( ) -> list[str]: models_to_remove: Final = set() routes_to_model: Final = _routes_to_model(_cached_wildcard_routing(llm_router), team_id) + rival_patterns: Final = _rival_patterns(llm_router, team_id) all_wildcard_models: Final = [] for model in unique_models: if _check_wildcard_routing(model=model): @@ -844,6 +994,7 @@ def _get_wildcard_models( wildcard_model=model, litellm_params=LiteLLM_Params(**router_model["litellm_params"]), routes_to_model=routes_to_model, + rival_patterns=rival_patterns, ) ) else: diff --git a/litellm/router_rust/surface/indexes.py b/litellm/router_rust/surface/indexes.py index f5aca3e1bad..32f63bc593b 100644 --- a/litellm/router_rust/surface/indexes.py +++ b/litellm/router_rust/surface/indexes.py @@ -308,8 +308,9 @@ class HostedPatternRouter(PatternMatchRouter): if not isinstance(model, str): # pyright: ignore[reportUnnecessaryIsInstance] # a caller's non-string model goes to the base class return super().get_models_by_pattern(model, custom_llm_provider) - provider: Final = custom_llm_provider or inferred_provider(model) - return self._owner.models_by_pattern(self._team, [model, f"{provider}/{model}"]) + return self._owner.models_by_pattern(self._team, [model]) or self._owner.models_by_pattern( + self._team, [f"{custom_llm_provider or inferred_provider(model)}/{model}"] + ) class _TeamRouters(dict[str, HostedPatternRouter]): diff --git a/tests/unit/proxy/auth/test_model_checks.py b/tests/unit/proxy/auth/test_model_checks.py index 1c576b1ee24..83fc81815d8 100644 --- a/tests/unit/proxy/auth/test_model_checks.py +++ b/tests/unit/proxy/auth/test_model_checks.py @@ -2463,27 +2463,379 @@ def test_wildcard_lists_an_id_a_more_specific_sibling_serves_with_the_same_model assert "openai/gpt-partial-route-9*" not in model_info_names -def test_prefixed_alias_over_a_bare_catch_all_lists_ids_that_route_through_it(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize( + ("model_name", "expected_id"), + [ + pytest.param("team-partial-route/*", f"team-partial-route/{_ROUTE_OPENAI_ID}", id="whole-name"), + pytest.param("team-partial-route/gpt-*", f"team-partial-route/gpt-{_ROUTE_OPENAI_ID}", id="partial-name"), + pytest.param("team-partial-route/*-*", None, id="more-wildcards-than-the-deployment"), + ], +) +def test_prefixed_alias_over_a_bare_catch_all_lists_ids_that_route_to_catalog_models( + monkeypatch: pytest.MonkeyPatch, model_name: str, expected_id: str | None +) -> None: + from litellm import Router + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info, get_complete_model_list + from litellm.utils import get_valid_models + + monkeypatch.setenv("OPENAI_API_KEY", "sk-partial-route-test") + monkeypatch.setenv("AZURE_API_KEY", "azure-partial-route-test") + router: Final = Router(model_list=[{"model_name": model_name, "litellm_params": {"model": "*"}}]) + + def routed_model(listed_id: str) -> object: + routed: Final = router.registry().deployments(listed_id, source="wildcard") + return routed[0]["litellm_params"]["model"] if routed else None + + with _registered_chat_models({_ROUTE_OPENAI_ID: "openai"}): + catalog_ids: Final = frozenset(get_valid_models()) - {AZURE_MODEL_PLACEHOLDER} + v1_models: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=[model_name], + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + model_info_names: Final = [ + str(row["model_name"]) + for row in expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.registry().deployments() or []], router + ) + if row["model_name"] != model_name + ] + + listed_ids: Final = [*v1_models, *model_info_names] + assert [listed_id for listed_id in listed_ids if routed_model(listed_id) not in catalog_ids] == [] + if expected_id is None: + assert listed_ids == [] + else: + assert expected_id in v1_models + assert expected_id in model_info_names + assert routed_model(expected_id) == _ROUTE_OPENAI_ID + + +_ROUTE_OPENAI_CHAT_ID: Final = "partial-route-9-chat" + + +@pytest.mark.parametrize( + ("alias_names", "stolen_id", "kept_ids"), + [ + pytest.param( + ("team-partial-route/gpt-*", "team-partial-route/gpt-gpt-partial-route-*"), + f"team-partial-route/gpt-{_ROUTE_OPENAI_ID}", + ( + f"team-partial-route/gpt-gpt-partial-route-{_ROUTE_OPENAI_ID}", + f"team-partial-route/gpt-{_ROUTE_OPENAI_CHAT_ID}", + ), + id="partial-name-next-to-a-longer-partial-name", + ), + pytest.param( + ("team-partial-route/*", "team-partial-route/gpt-*"), + f"team-partial-route/{_ROUTE_OPENAI_ID}", + (f"team-partial-route/gpt-{_ROUTE_OPENAI_ID}", f"team-partial-route/{_ROUTE_OPENAI_CHAT_ID}"), + id="whole-name-next-to-a-partial-name", + ), + ], +) +def test_alias_over_a_bare_catch_all_skips_ids_a_more_specific_sibling_routes_elsewhere( + monkeypatch: pytest.MonkeyPatch, alias_names: tuple[str, ...], stolen_id: str, kept_ids: tuple[str, ...] +) -> None: + from litellm import Router + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info, get_complete_model_list + from litellm.utils import get_valid_models + + monkeypatch.setenv("OPENAI_API_KEY", "sk-partial-route-test") + router: Final = Router(model_list=[{"model_name": name, "litellm_params": {"model": "*"}} for name in alias_names]) + + def routed_model(listed_id: str) -> object: + routed: Final = router.registry().deployments(listed_id, source="wildcard") + return routed[0]["litellm_params"]["model"] if routed else None + + with _registered_chat_models({_ROUTE_OPENAI_ID: "openai", _ROUTE_OPENAI_CHAT_ID: "openai"}): + catalog_ids: Final = frozenset(get_valid_models()) + v1_models: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=list(alias_names), + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + model_info_names: Final = [ + str(row["model_name"]) + for row in expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.registry().deployments() or []], router + ) + if row["model_name"] not in alias_names + ] + + listed_ids: Final = [*v1_models, *model_info_names] + assert routed_model(stolen_id) not in catalog_ids + assert [listed_id for listed_id in listed_ids if routed_model(listed_id) not in catalog_ids] == [] + assert [kept_id for kept_id in kept_ids if kept_id not in v1_models or kept_id not in model_info_names] == [] + + +def test_suffix_alias_over_a_bare_catch_all_lists_no_id_that_names_another_catalog_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm import Router + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info, get_complete_model_list + from litellm.utils import get_valid_models + + monkeypatch.setenv("OPENAI_API_KEY", "sk-partial-route-test") + alias: Final = "team-partial-route/*-mini" + full_size_id: Final = _ROUTE_OPENAI_ID.removesuffix("-mini") + router: Final = Router(model_list=[{"model_name": alias, "litellm_params": {"model": "*"}}]) + + def routed_model(listed_id: str) -> object: + routed: Final = router.registry().deployments(listed_id, source="wildcard") + return routed[0]["litellm_params"]["model"] if routed else None + + with _registered_chat_models({_ROUTE_OPENAI_ID: "openai", full_size_id: "openai"}): + catalog_ids: Final = frozenset(get_valid_models()) + v1_models: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=[alias], + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + model_info_names: Final = [ + str(row["model_name"]) + for row in expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.registry().deployments() or []], router + ) + if row["model_name"] != alias + ] + + listed_ids: Final = [*v1_models, *model_info_names] + named_ids: Final = {listed_id: listed_id.removeprefix("team-partial-route/") for listed_id in listed_ids} + assert routed_model(f"team-partial-route/{_ROUTE_OPENAI_ID}") == full_size_id + assert { + listed_id: routed_model(listed_id) + for listed_id, named_id in named_ids.items() + if named_id in catalog_ids and routed_model(listed_id) != named_id + } == {} + assert f"team-partial-route/{_ROUTE_OPENAI_ID}-mini" in v1_models + assert f"team-partial-route/{_ROUTE_OPENAI_ID}-mini" in model_info_names + + +@pytest.mark.parametrize( + ("global_alias", "team_alias", "team_alias_model", "stolen_id", "named_model", "kept_ids"), + [ + pytest.param( + "team-partial-route/*", + "team-partial-route/gpt-*", + "*", + f"team-partial-route/{_ROUTE_OPENAI_ID}", + _ROUTE_OPENAI_ID, + (f"team-partial-route/{_ROUTE_OPENAI_CHAT_ID}",), + id="team-alias-matching-the-id", + ), + pytest.param( + "team-partial-route/*", + "openai/*", + "openai/*", + f"team-partial-route/{_ROUTE_OPENAI_CHAT_ID}", + _ROUTE_OPENAI_CHAT_ID, + (f"team-partial-route/{_ROUTE_OPENAI_ID}",), + id="team-alias-matching-the-id-under-its-provider", + ), + pytest.param( + "team-partial-route/gpt-*", + "team-partial-route/*", + "*", + f"team-partial-route/gpt-{_ROUTE_OPENAI_ID}", + _ROUTE_OPENAI_ID, + (), + id="team-alias-broader-than-the-global-alias", + ), + pytest.param( + "team-partial-route/*", + "o*", + "*", + f"team-partial-route/{_ROUTE_OPENAI_ID}", + _ROUTE_OPENAI_ID, + (), + id="provider-less-team-alias-matching-the-id-under-its-provider", + ), + ], +) +def test_team_key_listing_skips_ids_the_team_wildcard_routes_elsewhere( + monkeypatch: pytest.MonkeyPatch, + global_alias: str, + team_alias: str, + team_alias_model: str, + stolen_id: str, + named_model: str, + kept_ids: tuple[str, ...], +) -> None: from litellm import Router from litellm.proxy.auth.model_checks import get_complete_model_list monkeypatch.setenv("OPENAI_API_KEY", "sk-partial-route-test") - router: Final = Router(model_list=[{"model_name": "team-partial-route/*", "litellm_params": {"model": "*"}}]) - with _registered_chat_models({_ROUTE_OPENAI_ID: "openai"}): - v1_models: Final = get_complete_model_list( + router: Final = Router( + model_list=[ + {"model_name": global_alias, "litellm_params": {"model": "*"}}, + { + "model_name": "model_name_team-1_partial-route", + "litellm_params": {"model": team_alias_model}, + "model_info": {"id": "team-partial-route", "team_id": "team-1", "team_public_model_name": team_alias}, + }, + ] + ) + + def listed_for(team_id: str | None) -> list[str]: + return get_complete_model_list( key_models=[], team_models=[], + proxy_model_list=[global_alias], + user_model=None, + infer_model_from_keys=False, + llm_router=router, + team_id=team_id, + ) + + with _registered_chat_models({_ROUTE_OPENAI_ID: "openai", _ROUTE_OPENAI_CHAT_ID: "openai", stolen_id: "openai"}): + listed_without_team: Final = listed_for(None) + listed_for_team: Final = listed_for("team-1") + stolen_route: Final = router.registry().names(stolen_id, kind="wildcard", team_id="team-1") + + assert named_model not in stolen_route + assert stolen_id in listed_without_team + assert stolen_id not in listed_for_team + assert set(kept_ids) <= set(listed_for_team) + + +def test_wildcard_model_group_alias_skips_ids_its_target_sends_to_another_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm import Router + from litellm.proxy.auth.model_checks import get_complete_model_list + + monkeypatch.setenv("OPENAI_API_KEY", "sk-partial-route-test") + router: Final = Router( + model_list=[{"model_name": "team-partial-route/*", "litellm_params": {"model": "*"}}], + model_group_alias={"team-partial-route/gpt-*": "team-partial-route/*"}, + ) + misnamed_id: Final = f"team-partial-route/gpt-{_ROUTE_OPENAI_ID}" + + with _registered_chat_models({_ROUTE_OPENAI_ID: "openai"}): + listed: Final = get_complete_model_list( + key_models=["team-partial-route/gpt-*"], + team_models=[], proxy_model_list=["team-partial-route/*"], user_model=None, infer_model_from_keys=False, llm_router=router, ) + misnamed_route: Final = router.registry().names(misnamed_id, kind="wildcard") - unrouted_ids: Final = [ - listed_id for listed_id in v1_models if not router.wildcard_route_deployments(listed_id, None) - ] - assert f"team-partial-route/{_ROUTE_OPENAI_ID}" in v1_models - assert unrouted_ids == [] + assert _ROUTE_OPENAI_ID not in misnamed_route + assert misnamed_id not in listed + + +def test_team_copy_of_a_partial_alias_lists_ids_a_global_sibling_would_take(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm import Router + from litellm.proxy.auth.model_checks import get_complete_model_list + + monkeypatch.setenv("OPENAI_API_KEY", "sk-partial-route-test") + alias: Final = "team-partial-route/gpt-*" + sibling_taken_id: Final = f"team-partial-route/gpt-{_ROUTE_OPENAI_CHAT_ID}" + router: Final = Router( + model_list=[ + {"model_name": alias, "litellm_params": {"model": "*"}}, + {"model_name": "team-partial-route/gpt-partial-*", "litellm_params": {"model": "openai/*"}}, + { + "model_name": "model_name_team-1_partial-route", + "litellm_params": {"model": "*"}, + "model_info": {"id": "team-partial-route", "team_id": "team-1", "team_public_model_name": alias}, + }, + ] + ) + + def listed_for(team_id: str | None) -> list[str]: + return get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=[alias], + user_model=None, + infer_model_from_keys=False, + llm_router=router, + team_id=team_id, + ) + + with _registered_chat_models({_ROUTE_OPENAI_ID: "openai", _ROUTE_OPENAI_CHAT_ID: "openai"}): + listed_without_team: Final = listed_for(None) + listed_for_team: Final = listed_for("team-1") + team_route: Final = router.registry().names(sibling_taken_id, kind="wildcard", team_id="team-1") + + assert team_route == {_ROUTE_OPENAI_CHAT_ID} + assert sibling_taken_id not in listed_without_team + assert sibling_taken_id in listed_for_team + assert f"team-partial-route/gpt-{_ROUTE_OPENAI_ID}" in listed_for_team + + +def test_partial_alias_over_a_catch_all_pinned_to_a_provider_lists_only_that_providers_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm import Router + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info, get_complete_model_list + + monkeypatch.setenv("OPENAI_API_KEY", "sk-partial-route-test") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-partial-route-test") + alias: Final = "team-partial-route/gpt-*" + anthropic_id: Final = "claude-partial-route-9" + router: Final = Router( + model_list=[{"model_name": alias, "litellm_params": {"model": "*", "custom_llm_provider": "openai"}}] + ) + + with _registered_chat_models({_ROUTE_OPENAI_ID: "openai", anthropic_id: "anthropic"}): + v1_models: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=[alias], + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + model_info_names: Final = [ + str(row["model_name"]) + for row in expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.registry().deployments() or []], router + ) + ] + + assert f"team-partial-route/gpt-{_ROUTE_OPENAI_ID}" in v1_models + assert f"team-partial-route/gpt-{_ROUTE_OPENAI_ID}" in model_info_names + assert f"team-partial-route/gpt-{anthropic_id}" not in v1_models + assert f"team-partial-route/gpt-{anthropic_id}" not in model_info_names + + +@respx.mock +def test_partial_alias_over_a_pinned_catch_all_asks_the_provider_for_its_models_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.setenv("OPENAI_API_KEY", "sk-pinned") + models_route: Final = respx.get("http://pinned-catch-all.test/v1/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "pinned-route-model"}]}) + ) + + listed: Final = get_known_models_from_wildcard( + "team-pinned-route/gpt-*", + LiteLLM_Params( + model="*", custom_llm_provider="openai", api_base="http://pinned-catch-all.test/v1", api_key="sk-pinned" + ), + ) + + assert "team-pinned-route/gpt-pinned-route-model" in listed + assert models_route.call_count == 1 @pytest.mark.asyncio