mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
fix(proxy): list key and team model aliases in GET /v1/models (#42908)
* fix(proxy): list key and team model aliases in GET /v1/models
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): keep alias listing helpers within the type discipline budget
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(proxy): cover alias rows on GET /v1/models and /v1/models/{id}
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): apply team then key aliases like chat completions and keep the alias as the retrieved id
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): apply key aliases twice like chat completions and skip only malformed alias entries
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): apply the global model_alias_map between the key alias passes like chat completions
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): list only the caller's own aliases and never rewrite a listed model id on retrieval
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* style(proxy): ruff format model_info alias lookup
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): hide undiscoverable names from model retrieval so an alias named like one resolves to its target
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): keep undiscoverable models retrievable by id while excluding them from the alias guard
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(proxy): pass an immutable name sequence into the model_info alias guard
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(proxy): type the model list alias test helpers
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(proxy): annotate the new alias listing test fixtures and helpers
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
c4b56b6ada
commit
8477fe4108
8 changed files with 477 additions and 35 deletions
|
|
@ -13,9 +13,12 @@ from __future__ import annotations
|
|||
import re
|
||||
from collections.abc import Container, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -28,6 +31,8 @@ CLAUDE_CODE_CLIENT: Final = "claude-code"
|
|||
_CLAUDE_CODE_ALIAS_PREFIX: Final = "claude-router-"
|
||||
_ONE_MILLION_SUFFIX: Final = "[1m]"
|
||||
_ONE_MILLION_TOKENS: Final = 1_000_000
|
||||
_ALIAS_ENTRIES: Final = TypeAdapter(Mapping[object, object])
|
||||
_NO_ALIASES: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
def configured_display_names(
|
||||
|
|
@ -152,6 +157,77 @@ class ClaudeCodeRoutingNames:
|
|||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CallerAliases:
|
||||
"""`own` are the caller's key and team alias maps, the names `/v1/models` lists for it.
|
||||
`rewrite` are the maps `/chat/completions` rewrites its model through, 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`."""
|
||||
|
||||
own: tuple[object, ...]
|
||||
rewrite: tuple[object, ...]
|
||||
|
||||
|
||||
def caller_alias_maps(
|
||||
key_aliases: object,
|
||||
team_aliases: object,
|
||||
key_team_id: str | None,
|
||||
listed_team_id: str | None,
|
||||
) -> CallerAliases:
|
||||
"""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((key_aliases,), (key_aliases, litellm.model_alias_map, key_aliases))
|
||||
return CallerAliases((team_aliases, key_aliases), (team_aliases, key_aliases, litellm.model_alias_map, key_aliases))
|
||||
|
||||
|
||||
def _alias_map(aliases: object) -> Mapping[str, str]:
|
||||
try:
|
||||
entries: Final = _ALIAS_ENTRIES.validate_python(aliases, strict=True)
|
||||
except ValidationError:
|
||||
return _NO_ALIASES
|
||||
return MappingProxyType(
|
||||
{alias: target for alias, target in entries.items() if isinstance(alias, str) and isinstance(target, str)}
|
||||
)
|
||||
|
||||
|
||||
def _alias_names(alias_maps: Sequence[Mapping[str, str]]) -> tuple[str, ...]:
|
||||
return tuple(dict.fromkeys(alias for aliases in alias_maps for alias in aliases))
|
||||
|
||||
|
||||
def _rewrite(model_id: str, alias_maps: Sequence[Mapping[str, str]]) -> str | None:
|
||||
target: Final = reduce(lambda name, aliases: aliases.get(name, name), alias_maps, model_id)
|
||||
return None if target == model_id else target
|
||||
|
||||
|
||||
def alias_target(model_id: str, aliases: CallerAliases, listed: Container[str] = frozenset()) -> str | None:
|
||||
"""The model group `/chat/completions` rewrites `model_id` to, else None. A `model_id`
|
||||
already `listed` keeps its own row, so it is never rewritten."""
|
||||
if model_id in listed:
|
||||
return None
|
||||
return _rewrite(model_id, tuple(_alias_map(alias_map) for alias_map in aliases.rewrite))
|
||||
|
||||
|
||||
def alias_listing_entries(
|
||||
entries: Sequence[tuple[str, str]],
|
||||
aliases: CallerAliases,
|
||||
) -> tuple[tuple[str, str], ...]:
|
||||
"""`entries` plus one `(alias, lookup_id)` row per key or team alias whose target is
|
||||
listed. An alias colliding with a listed id keeps the listed entry."""
|
||||
maps: Final = tuple(_alias_map(alias_map) for alias_map in aliases.rewrite)
|
||||
own: Final = tuple(_alias_map(alias_map) for alias_map in aliases.own)
|
||||
lookup_by_response: Final = MappingProxyType(dict(entries))
|
||||
lookup_ids: Final = frozenset(lookup_by_response.values())
|
||||
targets: Final = MappingProxyType(
|
||||
{alias: _rewrite(alias, maps) for alias in _alias_names(own) if alias not in lookup_by_response}
|
||||
)
|
||||
added: Final = tuple(
|
||||
(alias, lookup_by_response.get(target, target))
|
||||
for alias, target in targets.items()
|
||||
if target is not None and (target in lookup_by_response or target in lookup_ids)
|
||||
)
|
||||
return (*entries, *added)
|
||||
|
||||
|
||||
def claude_code_requested_group(
|
||||
requested: str,
|
||||
llm_router: Router,
|
||||
|
|
@ -218,7 +294,7 @@ class TeamModelNameTranslator:
|
|||
|
||||
@staticmethod
|
||||
def _response_to_lookup_map(
|
||||
model_names: list[str],
|
||||
model_names: Sequence[str],
|
||||
internal_to_public: dict[str, str],
|
||||
) -> dict[str, str]:
|
||||
"""Map each public response id to the first internal lookup id seen in
|
||||
|
|
@ -235,7 +311,7 @@ class TeamModelNameTranslator:
|
|||
|
||||
@staticmethod
|
||||
def listing_entries(
|
||||
model_names: list[str],
|
||||
model_names: Sequence[str],
|
||||
llm_router: Router | None,
|
||||
general_settings: Mapping[str, object],
|
||||
) -> list[tuple[str, str]]:
|
||||
|
|
|
|||
|
|
@ -422,6 +422,9 @@ from litellm.proxy.common_utils.model_deprecation import collect_model_deprecati
|
|||
from litellm.proxy.common_utils.model_listing_utils import (
|
||||
ClaudeCodeRoutingNames,
|
||||
TeamModelNameTranslator,
|
||||
alias_listing_entries,
|
||||
alias_target,
|
||||
caller_alias_maps,
|
||||
claude_code_view_ids,
|
||||
configured_display_names,
|
||||
is_claude_code_client,
|
||||
|
|
@ -11172,14 +11175,13 @@ async def model_list(
|
|||
view_aliases: Final = (
|
||||
view_router_settings.get("model_group_alias") if isinstance(view_router_settings, Mapping) else None
|
||||
)
|
||||
caller_aliases: Final = caller_alias_maps(
|
||||
user_api_key_dict.aliases, user_api_key_dict.team_model_aliases, user_api_key_dict.team_id, team_id
|
||||
)
|
||||
routing_names: Final = ClaudeCodeRoutingNames(
|
||||
llm_router,
|
||||
team_id or user_api_key_dict.team_id,
|
||||
(
|
||||
user_api_key_dict.aliases,
|
||||
user_api_key_dict.team_model_aliases,
|
||||
view_aliases,
|
||||
),
|
||||
(*caller_aliases.rewrite, view_aliases),
|
||||
)
|
||||
|
||||
# Validate scope parameter if provided
|
||||
|
|
@ -11307,7 +11309,9 @@ async def model_list(
|
|||
# The internal routing key drives the metadata/fallback lookup, while the
|
||||
# public name is what the client sees as the model id.
|
||||
model_data = []
|
||||
entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings)
|
||||
entries: Final = alias_listing_entries(
|
||||
TeamModelNameTranslator.listing_entries(all_models, llm_router, settings), caller_aliases
|
||||
)
|
||||
for response_id, lookup_id in entries:
|
||||
model_info = create_model_info_response(
|
||||
model_id=lookup_id,
|
||||
|
|
@ -11391,7 +11395,8 @@ async def model_info(
|
|||
)
|
||||
|
||||
# Mirror /v1/models' visibility filter so first-occurrence resolution
|
||||
# cannot land on a deployment the listing had hidden.
|
||||
# cannot land on a deployment the listing had hidden. Undiscoverable
|
||||
# models stay retrievable by id, they only drop out of the alias guard.
|
||||
blocked_names: Final = llm_router.get_fully_blocked_model_names() if llm_router is not None else set()
|
||||
unhealthy_names: Final = await get_hidden_unhealthy_model_names(
|
||||
healthy_only=healthy_only,
|
||||
|
|
@ -11401,10 +11406,25 @@ async def model_info(
|
|||
hidden_names: Final = blocked_names | unhealthy_names
|
||||
if hidden_names:
|
||||
all_models = [m for m in all_models if m not in hidden_names]
|
||||
undiscoverable_names: Final = undiscoverable_model_names(
|
||||
all_models, llm_router, user_api_key_dict, team_id or user_api_key_dict.team_id
|
||||
)
|
||||
|
||||
internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, settings)
|
||||
aliased_model_id: Final = alias_target(
|
||||
model_id,
|
||||
caller_alias_maps(
|
||||
user_api_key_dict.aliases, user_api_key_dict.team_model_aliases, user_api_key_dict.team_id, team_id
|
||||
),
|
||||
frozenset(
|
||||
response_id
|
||||
for response_id, _ in TeamModelNameTranslator.listing_entries(
|
||||
tuple(m for m in all_models if m not in undiscoverable_names), llm_router, settings
|
||||
)
|
||||
),
|
||||
)
|
||||
resolved_model_id: Final = TeamModelNameTranslator.resolve_public_name(
|
||||
model_id=model_id,
|
||||
model_id=aliased_model_id or model_id,
|
||||
available_models=all_models,
|
||||
llm_router=llm_router,
|
||||
general_settings=settings,
|
||||
|
|
@ -11434,7 +11454,8 @@ async def model_info(
|
|||
fallback_type=None,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
return {**response, "id": internal_to_public.get(resolved_model_id, model_id)} # mutable-ok: response id differs
|
||||
response_id: Final = model_id if aliased_model_id else internal_to_public.get(resolved_model_id, model_id)
|
||||
return {**response, "id": response_id} # mutable-ok: response id differs
|
||||
|
||||
|
||||
def _blocked_response_usage(original_response: object | None) -> "litellm.Usage":
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@
|
|||
- {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"}
|
||||
- {id: other.auth.jwt.team_header_alias_binds_team, module: other, tier: P0, area: auth, assertions: [team_header_alias_binds_team], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", fail_before_fix: proven, rationale: "x-litellm-team-id carrying the team alias binds and attributes the same team as the team id, so a managed client can pin a stable alias instead of a uuid"}
|
||||
- {id: other.auth.jwt.team_header_non_member_alias_denied, module: other, tier: P0, area: auth, assertions: [team_header_non_member_alias_denied], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", rationale: "x-litellm-team-id naming the alias of a team the JWT does not grant is denied 403 with the same body as an unknown value, so the response does not reveal whether that team exists"}
|
||||
- {id: other.auth.jwt.team_model_alias_listed_and_routes, module: other, tier: P1, area: auth, assertions: [team_model_alias_listed_and_routes], source: "proxy_server.py model_list / common_utils/model_listing_utils.py alias_listing_entries / LIT-8515", fail_before_fix: proven, rationale: "A team model_aliases name the JWT caller can complete on is also listed by GET /v1/models for that caller, in the OpenAI and the Anthropic (Claude Code) shapes, next to its target, so a managed client can discover the alias it is meant to send"}
|
||||
- {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"}
|
||||
- {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"}
|
||||
- {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"}
|
||||
|
|
|
|||
|
|
@ -1426,6 +1426,7 @@ class TeamNewBody(BaseModel):
|
|||
team_id: str | None = None
|
||||
organization_id: str | None = None
|
||||
metadata: TeamMetadata | None = None
|
||||
model_aliases: dict[str, str] | None = None
|
||||
|
||||
|
||||
class TeamNewResponse(BaseModel):
|
||||
|
|
|
|||
|
|
@ -14,12 +14,15 @@ own endpoints, so no test ever holds a signing key.
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from e2e_http import AuthHeaders, NoBody, ProbeResult, Result
|
||||
from e2e_http import AnthropicHeaders, AuthHeaders, NoBody, ProbeResult, Result
|
||||
from idp import Keycloak, keycloak_from_env
|
||||
from models import (
|
||||
ChatBody,
|
||||
ChatResponse,
|
||||
ModelsListParams,
|
||||
ModelsListResponse,
|
||||
ReadinessDetailsResponse,
|
||||
ReadinessResponse,
|
||||
UserListParams,
|
||||
|
|
@ -88,6 +91,17 @@ class OtherClient:
|
|||
response_type=ChatResponse,
|
||||
)
|
||||
|
||||
def list_models_as(self, token: str, *, anthropic: bool = False) -> Result[ModelsListResponse]:
|
||||
"""GET /v1/models under `token`, in the OpenAI shape or, with `anthropic`, the
|
||||
Anthropic Models API shape Claude Code reads. Both carry `data[].id`."""
|
||||
bearer: Final = self.proxy.transport.bearer(token)
|
||||
return self.proxy.transport.get(
|
||||
"/v1/models",
|
||||
headers=AnthropicHeaders(authorization=bearer.authorization) if anthropic else bearer,
|
||||
params=ModelsListParams(return_wildcard_routes=False),
|
||||
response_type=ModelsListResponse,
|
||||
)
|
||||
|
||||
def list_users_as(self, key: str) -> Result[UserListResponse]:
|
||||
"""GET /user/list under `key`. Admin-only, so it doubles as the master
|
||||
key's authorization proof: the master key (proxy admin) reads it, a
|
||||
|
|
|
|||
|
|
@ -78,9 +78,35 @@ def bound_team(client: OtherClient, resources: ResourceManager) -> BoundTeam:
|
|||
return BoundTeam(identity=provisioned, team_id=provisioned.group, team_alias=team_alias)
|
||||
|
||||
|
||||
def _ping() -> ChatBody:
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AliasedTeam:
|
||||
identity: Identity
|
||||
alias: str
|
||||
target: str
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def aliased_team(client: OtherClient, resources: ResourceManager) -> AliasedTeam:
|
||||
"""An identity whose team carries a model_aliases entry, the name a managed
|
||||
client such as Claude Code sends and the team rewrites to a real model group."""
|
||||
marker: Final = unique_marker()
|
||||
provisioned: Final = _provision(client, resources, marker=marker)
|
||||
alias: Final = f"e2e-jwt-model-alias-{marker}"
|
||||
team_id: Final = client.proxy.create_team(
|
||||
TeamNewBody(
|
||||
team_alias=f"e2e-jwt-aliased-{marker}",
|
||||
team_id=provisioned.group,
|
||||
models=[CHEAP_OPENAI_MODEL],
|
||||
model_aliases={alias: CHEAP_OPENAI_MODEL},
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: client.proxy.delete_team(team_id))
|
||||
return AliasedTeam(identity=provisioned, alias=alias, target=CHEAP_OPENAI_MODEL)
|
||||
|
||||
|
||||
def _ping(model: str = CHEAP_OPENAI_MODEL) -> ChatBody:
|
||||
return ChatBody(
|
||||
model=CHEAP_OPENAI_MODEL,
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")],
|
||||
max_tokens=16,
|
||||
)
|
||||
|
|
@ -221,6 +247,23 @@ class TestJwtTeamHeader:
|
|||
f"{bound_team.team_id!r}, got {by_alias!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("other.auth.jwt.team_model_alias_listed_and_routes")
|
||||
@pytest.mark.parametrize("anthropic", [False, True], ids=["openai_shape", "anthropic_shape"])
|
||||
def test_team_model_alias_is_listed_by_v1_models_under_the_same_token_that_routes_it(
|
||||
self, client: OtherClient, aliased_team: AliasedTeam, anthropic: bool
|
||||
) -> None:
|
||||
token: Final = client.idp.access_token(aliased_team.identity)
|
||||
|
||||
routed: Final = unwrap(client.proxy.chat(token, _ping(model=aliased_team.alias)))
|
||||
assert routed.choices, f"precondition: /chat/completions must route the team alias, got {routed}"
|
||||
|
||||
listed: Final = tuple(entry.id for entry in unwrap(client.list_models_as(token, anthropic=anthropic)).data)
|
||||
assert aliased_team.alias in listed, (
|
||||
f"/v1/models must list team alias {aliased_team.alias!r} that the same token routes on "
|
||||
f"/chat/completions, got {listed}"
|
||||
)
|
||||
assert aliased_team.target in listed, f"the alias target {aliased_team.target!r} must stay listed, got {listed}"
|
||||
|
||||
@pytest.mark.covers("other.auth.jwt.team_header_non_member_alias_denied")
|
||||
def test_team_header_with_the_alias_of_a_team_the_caller_is_not_in_is_rejected_like_an_unknown_value(
|
||||
self, client: OtherClient, resources: ResourceManager, bound_team: BoundTeam
|
||||
|
|
|
|||
|
|
@ -4,9 +4,14 @@ from itertools import combinations
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.proxy.common_utils.model_listing_utils import (
|
||||
CallerAliases,
|
||||
ClaudeCodeRoutingNames,
|
||||
alias_listing_entries,
|
||||
alias_target,
|
||||
caller_alias_maps,
|
||||
claude_code_group_name,
|
||||
claude_code_model_id,
|
||||
claude_code_requested_group,
|
||||
|
|
@ -22,6 +27,10 @@ def _marked(name):
|
|||
return f"{_encoded(name)}[1m]"
|
||||
|
||||
|
||||
def _caller(*maps: object) -> CallerAliases:
|
||||
return CallerAliases(maps, maps)
|
||||
|
||||
|
||||
def _row(name, limit=1000000):
|
||||
return {"id": name, "object": "model", "created": 0, "owned_by": "openai", "max_input_tokens": limit}
|
||||
|
||||
|
|
@ -29,15 +38,16 @@ def _row(name, limit=1000000):
|
|||
def _router(*names, aliases=None):
|
||||
return Router(
|
||||
model_list=[
|
||||
{"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}}
|
||||
for name in names
|
||||
{"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in names
|
||||
],
|
||||
model_group_alias=aliases,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("limit", [None, 999999, 1000000])
|
||||
@pytest.mark.parametrize("name", ["foo", "foo[1m]", "foo[1M]", "a/b: 世界", "claude-router-foo", "claude-opus-5", "claude-opus-5[1m]"])
|
||||
@pytest.mark.parametrize(
|
||||
"name", ["foo", "foo[1m]", "foo[1M]", "a/b: 世界", "claude-router-foo", "claude-opus-5", "claude-opus-5[1m]"]
|
||||
)
|
||||
def test_listing_round_trips_entire_source_name(name, limit):
|
||||
names = frozenset({name})
|
||||
view = claude_code_model_id(name, limit, names)
|
||||
|
|
@ -48,7 +58,15 @@ def test_listing_round_trips_entire_source_name(name, limit):
|
|||
|
||||
|
||||
def test_collision_matrix_round_trips_without_duplicate_ids():
|
||||
universe = ("foo", "foo[1m]", "claude-router-foo", _encoded("foo"), _encoded("foo") + "[1m]", "claude-opus-5", "claude-opus-5[1m]")
|
||||
universe = (
|
||||
"foo",
|
||||
"foo[1m]",
|
||||
"claude-router-foo",
|
||||
_encoded("foo"),
|
||||
_encoded("foo") + "[1m]",
|
||||
"claude-opus-5",
|
||||
"claude-opus-5[1m]",
|
||||
)
|
||||
for pair in combinations(universe, 2):
|
||||
for visible in (pair, pair[:1], pair[1:]):
|
||||
names = frozenset(pair)
|
||||
|
|
@ -57,18 +75,31 @@ def test_collision_matrix_round_trips_without_duplicate_ids():
|
|||
assert all((claude_code_group_name(shown, names) or shown) == source for source, shown in view.items())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spelling", ["claude-router-foo", "claude-router-ff", "claude-router-66 6f6f", "claude-router-666F6F", "claude-router-", _encoded("missing")])
|
||||
@pytest.mark.parametrize(
|
||||
"spelling",
|
||||
[
|
||||
"claude-router-foo",
|
||||
"claude-router-ff",
|
||||
"claude-router-66 6f6f",
|
||||
"claude-router-666F6F",
|
||||
"claude-router-",
|
||||
_encoded("missing"),
|
||||
],
|
||||
)
|
||||
def test_unknown_or_noncanonical_ids_are_never_guessed(spelling):
|
||||
assert claude_code_group_name(spelling, frozenset({"foo"})) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("headers,enabled", [
|
||||
({"user-agent": "claude-code/2.1.267"}, True),
|
||||
({"user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, True),
|
||||
({"x-gateway-client": "Claude-Code"}, True),
|
||||
({"user-agent": "anthropic-sdk-python/0.40"}, False),
|
||||
({}, False),
|
||||
])
|
||||
@pytest.mark.parametrize(
|
||||
"headers,enabled",
|
||||
[
|
||||
({"user-agent": "claude-code/2.1.267"}, True),
|
||||
({"user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, True),
|
||||
({"x-gateway-client": "Claude-Code"}, True),
|
||||
({"user-agent": "anthropic-sdk-python/0.40"}, False),
|
||||
({}, False),
|
||||
],
|
||||
)
|
||||
def test_only_claude_code_gets_the_view(headers, enabled):
|
||||
rows = (_row("foo"), _row("claude-opus-5"))
|
||||
view = claude_code_view_ids(rows, headers, frozenset(row["id"] for row in rows))
|
||||
|
|
@ -77,12 +108,15 @@ def test_only_claude_code_gets_the_view(headers, enabled):
|
|||
|
||||
@pytest.mark.parametrize("layer", ["literal", "global", "router", "key", "team", "wildcard"])
|
||||
def test_configured_names_outrank_generated_ids_even_when_hidden_from_listing(monkeypatch, layer):
|
||||
import litellm
|
||||
|
||||
encoded = _encoded("foo")
|
||||
alias = {encoded: "other"}
|
||||
monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {})
|
||||
router = _router("foo", "other", *( (encoded,) if layer == "literal" else ("*",) if layer == "wildcard" else ()), aliases=alias if layer == "router" else None)
|
||||
router = _router(
|
||||
"foo",
|
||||
"other",
|
||||
*((encoded,) if layer == "literal" else ("*",) if layer == "wildcard" else ()),
|
||||
aliases=alias if layer == "router" else None,
|
||||
)
|
||||
maps = (alias,) if layer in ("key", "team") else ()
|
||||
names = ClaudeCodeRoutingNames(router, None, maps)
|
||||
assert claude_code_requested_group(encoded, router, None, maps) is None
|
||||
|
|
@ -98,12 +132,113 @@ def test_mutation_breaking_the_hex_name_cannot_route_to_the_source(source):
|
|||
assert claude_code_requested_group(_marked(source), router, None) == source
|
||||
|
||||
|
||||
def test_team_alias_is_listed_under_its_target_metadata_and_only_when_the_target_is_accessible() -> None:
|
||||
entries = [("gpt-4.1-mini", "gpt-4.1-mini"), ("team-public", "model_name_team_1_abc")]
|
||||
aliases = (
|
||||
{"gpt-4.1-mini": "team-public"},
|
||||
None,
|
||||
{"claude-sonnet-4-5": "gpt-4.1-mini", "via-public": "team-public", "not-granted": "gpt-4.1"},
|
||||
)
|
||||
assert alias_listing_entries(entries, _caller(*aliases)) == (
|
||||
*entries,
|
||||
("claude-sonnet-4-5", "gpt-4.1-mini"),
|
||||
("via-public", "model_name_team_1_abc"),
|
||||
)
|
||||
assert alias_listing_entries(entries, _caller(None, {})) == tuple(entries)
|
||||
|
||||
|
||||
def test_alias_target_resolves_the_requested_alias_across_key_and_team_maps() -> None:
|
||||
maps = _caller({"o": "gpt-4.1"}, {"claude-sonnet-4-5": "gpt-4.1-mini"})
|
||||
assert alias_target("claude-sonnet-4-5", maps) == "gpt-4.1-mini"
|
||||
assert alias_target("gpt-4.1-mini", _caller(None, {"claude-sonnet-4-5": "gpt-4.1-mini"})) is None
|
||||
|
||||
|
||||
def test_alias_colliding_with_a_listed_id_keeps_the_listed_model_at_list_and_retrieval() -> None:
|
||||
entries = [("fast", "fast"), ("gpt-4.1-mini", "gpt-4.1-mini")]
|
||||
maps = _caller({"fast": "gpt-4.1-mini"})
|
||||
listed = frozenset(response_id for response_id, _ in entries)
|
||||
|
||||
assert alias_listing_entries(entries, maps) == tuple(entries)
|
||||
assert alias_target("fast", maps, listed) is None
|
||||
assert alias_target("fast", maps) == "gpt-4.1-mini"
|
||||
|
||||
|
||||
def test_alias_maps_apply_in_the_order_chat_completions_applies_them() -> None:
|
||||
team_then_key = _caller({"fast": "gpt-4.1-mini", "hop": "mid"}, {"fast": "gpt-4.1", "mid": "gpt-4.1"})
|
||||
entries = [("gpt-4.1-mini", "gpt-4.1-mini"), ("gpt-4.1", "gpt-4.1")]
|
||||
|
||||
assert alias_target("fast", team_then_key) == "gpt-4.1-mini"
|
||||
assert alias_target("hop", team_then_key) == "gpt-4.1"
|
||||
assert alias_listing_entries(entries, team_then_key) == (
|
||||
*entries,
|
||||
("fast", "gpt-4.1-mini"),
|
||||
("hop", "gpt-4.1"),
|
||||
("mid", "gpt-4.1"),
|
||||
)
|
||||
|
||||
|
||||
def test_one_bad_alias_entry_hides_only_itself() -> None:
|
||||
aliases = {"fast": "gpt-4.1-mini", "broken": 5, 7: "gpt-4.1-mini"}
|
||||
entries = [("gpt-4.1-mini", "gpt-4.1-mini")]
|
||||
|
||||
assert alias_listing_entries(entries, _caller(aliases)) == (*entries, ("fast", "gpt-4.1-mini"))
|
||||
assert alias_target("fast", _caller(aliases)) == "gpt-4.1-mini"
|
||||
|
||||
|
||||
def test_chained_key_alias_is_listed_only_when_its_final_target_is_listable() -> None:
|
||||
key_aliases = {"a": "b", "b": "hidden"}
|
||||
entries = [("b", "b")]
|
||||
|
||||
assert alias_listing_entries(entries, caller_alias_maps(key_aliases, None, "team-a", None)) == (*entries,)
|
||||
assert alias_target("a", caller_alias_maps(key_aliases, None, "team-a", None)) == "hidden"
|
||||
|
||||
|
||||
def test_team_aliases_only_apply_when_listing_the_team_the_key_authenticated_as(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
key_aliases, team_aliases, global_aliases = {"k": "gpt-4.1"}, {"t": "gpt-4.1-mini"}, {"g": "gpt-4.1"}
|
||||
monkeypatch.setattr(litellm, "model_alias_map", global_aliases)
|
||||
own_team = CallerAliases((team_aliases, key_aliases), (team_aliases, key_aliases, global_aliases, key_aliases))
|
||||
assert caller_alias_maps(key_aliases, team_aliases, "team-a", None) == own_team
|
||||
assert caller_alias_maps(key_aliases, team_aliases, "team-a", "team-a") == own_team
|
||||
assert caller_alias_maps(key_aliases, team_aliases, "team-a", "team-b") == CallerAliases(
|
||||
(key_aliases,), (key_aliases, global_aliases, key_aliases)
|
||||
)
|
||||
|
||||
|
||||
def test_global_alias_rewrites_between_the_two_key_passes_like_chat_completions(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "model_alias_map", {"b": "d"})
|
||||
key_aliases = {"a": "b", "b": "c"}
|
||||
entries = [("c", "c"), ("d", "d")]
|
||||
maps = caller_alias_maps(key_aliases, None, "team-a", None)
|
||||
|
||||
assert alias_target("a", maps) == "d"
|
||||
assert alias_listing_entries(entries, maps) == (*entries, ("a", "d"), ("b", "c"))
|
||||
|
||||
|
||||
def test_global_aliases_rewrite_but_are_not_listed_as_caller_rows(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "model_alias_map", {"g": "gpt-4.1-mini"})
|
||||
entries = [("gpt-4.1-mini", "gpt-4.1-mini")]
|
||||
maps = caller_alias_maps({"k": "g"}, None, "team-a", None)
|
||||
|
||||
assert alias_listing_entries(entries, maps) == (*entries, ("k", "gpt-4.1-mini"))
|
||||
assert alias_target("g", maps) == "gpt-4.1-mini"
|
||||
|
||||
|
||||
def test_team_public_name_uses_the_same_scope_at_list_and_request():
|
||||
router = Router(model_list=[{
|
||||
"model_name": "model_name_team-a_id",
|
||||
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"},
|
||||
"model_info": {"team_id": "team-a", "team_public_model_name": "shared"},
|
||||
}])
|
||||
shown = claude_code_view_ids((_row("shared"),), {"user-agent": "claude-code/2.1.267"}, ClaudeCodeRoutingNames(router, "team-a"))["shared"]
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "model_name_team-a_id",
|
||||
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"},
|
||||
"model_info": {"team_id": "team-a", "team_public_model_name": "shared"},
|
||||
}
|
||||
]
|
||||
)
|
||||
shown = claude_code_view_ids(
|
||||
(_row("shared"),), {"user-agent": "claude-code/2.1.267"}, ClaudeCodeRoutingNames(router, "team-a")
|
||||
)["shared"]
|
||||
assert claude_code_requested_group(shown, router, "team-a") == "shared"
|
||||
assert claude_code_requested_group(shown, router, "team-b") is None
|
||||
|
|
|
|||
151
tests/test_litellm/proxy/test_model_list_aliases.py
Normal file
151
tests/test_litellm/proxy/test_model_list_aliases.py
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
"""
|
||||
Tests for key and team `model_aliases` on the model listing endpoints: GET /v1/models
|
||||
(`model_list`, OpenAI and Anthropic shapes) and GET /v1/models/{id} (`model_info`).
|
||||
An alias the caller can complete on is listed next to its target and resolves by name.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from starlette.requests import Request
|
||||
|
||||
from litellm import Router
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
|
||||
def _deployment(model_name: str, model: str = "openai/gpt-4.1-mini", **model_info: str | bool) -> dict[str, object]:
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"model": model, "api_key": "sk-fake"},
|
||||
"model_info": {"id": f"{model_name}-id", **model_info},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def router(monkeypatch: pytest.MonkeyPatch) -> Router:
|
||||
router = Router(
|
||||
model_list=[
|
||||
_deployment("gpt-4.1-mini"),
|
||||
_deployment("gpt-4.1", model="openai/gpt-4.1"),
|
||||
_deployment("model_name_team1_abc", team_id="team1", team_public_model_name="team-chat"),
|
||||
_deployment("hidden", model="anthropic/claude-sonnet-4-5", discoverable=False),
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(proxy_server, "llm_model_list", router.model_list)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
monkeypatch.setattr(proxy_server, "user_model", None)
|
||||
return router
|
||||
|
||||
|
||||
def _team_member(
|
||||
team_id: str = "team1", models: list[str] | None = None, **aliases: dict[str, str] | None
|
||||
) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
api_key="sk-test",
|
||||
user_id="u",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
team_id=team_id,
|
||||
team_models=["gpt-4.1-mini", "model_name_team1_abc"],
|
||||
models=models or ["gpt-4.1-mini", "model_name_team1_abc"],
|
||||
**aliases,
|
||||
)
|
||||
|
||||
|
||||
def _anthropic_request(*extra_headers: tuple[bytes, bytes]) -> Request:
|
||||
return Request(
|
||||
scope={
|
||||
"type": "http",
|
||||
"method": "GET",
|
||||
"path": "/v1/models",
|
||||
"query_string": b"",
|
||||
"headers": [(b"anthropic-version", b"2023-06-01"), *extra_headers],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _claude_code_request() -> Request:
|
||||
return _anthropic_request((b"user-agent", b"claude-cli/2.1.267 (external, cli)"))
|
||||
|
||||
|
||||
async def _v1_models(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> list[str]:
|
||||
response = await proxy_server.model_list(user_api_key_dict=user_api_key_dict, request=request)
|
||||
return [m["id"] for m in response["data"]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v1_models_lists_team_alias_next_to_its_target_in_both_shapes(router: Router) -> None:
|
||||
caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"})
|
||||
|
||||
assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "claude-sonnet-4-5"]
|
||||
assert await _v1_models(caller, request=_anthropic_request()) == ["gpt-4.1-mini", "team-chat", "claude-sonnet-4-5"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claude_code_picker_lists_the_alias_under_its_own_name(router: Router) -> None:
|
||||
caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"})
|
||||
|
||||
picker_ids = await _v1_models(caller, request=_claude_code_request())
|
||||
assert any(picker_id.startswith("claude-sonnet-4-5") for picker_id in picker_ids), picker_ids
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v1_models_lists_key_alias_and_hides_alias_to_a_model_the_caller_cannot_list(router: Router) -> None:
|
||||
caller = _team_member(aliases={"mini": "gpt-4.1-mini", "big": "gpt-4.1"})
|
||||
|
||||
assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "mini"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v1_models_resolves_a_team_alias_through_the_key_alias_like_chat_completions_does(router: Router) -> None:
|
||||
caller = _team_member(team_model_aliases={"fast": "mid"}, aliases={"fast": "gpt-4.1", "mid": "gpt-4.1-mini"})
|
||||
|
||||
assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "fast", "mid"]
|
||||
response = await proxy_server.model_info(model_id="fast", user_api_key_dict=caller)
|
||||
assert response["id"] == "fast"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v1_models_skips_only_the_malformed_alias_entries(router: Router) -> None:
|
||||
caller = _team_member(team_model_aliases={"claude-sonnet-4-5": 5, "fast": "gpt-4.1-mini"})
|
||||
|
||||
assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "fast"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v1_models_by_id_resolves_a_team_alias_to_its_target_metadata(router: Router) -> None:
|
||||
caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"})
|
||||
|
||||
response = await proxy_server.model_info(model_id="claude-sonnet-4-5", user_api_key_dict=caller)
|
||||
assert response["id"] == "claude-sonnet-4-5"
|
||||
assert response["owned_by"] == "openai"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v1_models_by_id_retrieves_the_listed_model_when_an_alias_collides_with_its_id(router: Router) -> None:
|
||||
caller = _team_member(aliases={"team-chat": "gpt-4.1"})
|
||||
|
||||
assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat"]
|
||||
response = await proxy_server.model_info(model_id="team-chat", user_api_key_dict=caller)
|
||||
assert response["id"] == "team-chat"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v1_models_by_id_resolves_an_alias_named_like_an_undiscoverable_model_to_the_alias_target(
|
||||
router: Router,
|
||||
) -> None:
|
||||
caller = _team_member(aliases={"hidden": "gpt-4.1-mini"}, models=["gpt-4.1-mini", "model_name_team1_abc", "hidden"])
|
||||
|
||||
assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "hidden"]
|
||||
target = await proxy_server.model_info(model_id="gpt-4.1-mini", user_api_key_dict=caller)
|
||||
response = await proxy_server.model_info(model_id="hidden", user_api_key_dict=caller)
|
||||
assert response == {**target, "id": "hidden"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_v1_models_by_id_keeps_the_alias_as_id_when_it_targets_a_team_scoped_model(router: Router) -> None:
|
||||
caller = _team_member(team_model_aliases={"chat": "team-chat"})
|
||||
|
||||
assert "chat" in await _v1_models(caller)
|
||||
response = await proxy_server.model_info(model_id="chat", user_api_key_dict=caller)
|
||||
assert response["id"] == "chat"
|
||||
Loading…
Add table
Reference in a new issue