mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(auth): resolve bare model names against wildcard deployments in model access groups (#37492)
* fix(auth): resolve bare model names against wildcard deployments in model access groups * test(e2e): cover model access group permission checks on keys and teams
This commit is contained in:
ryan-crabbe-berri
GitHub
parent
6ca48efc8b
commit
74b279bc44
7 changed files with 553 additions and 14 deletions
|
|
@ -10253,11 +10253,13 @@ class Router:
|
|||
returned_models.extend(self.get_model_list_from_routing_groups(model_name=model_name))
|
||||
|
||||
if len(returned_models) == 0: # check if wildcard route
|
||||
potential_wildcard_models: Final = self.pattern_router.route(model_name) or []
|
||||
potential_wildcard_models: Final = self.pattern_router.get_deployments_by_pattern(model=model_name or "")
|
||||
|
||||
## check for team-specific wildcard models
|
||||
if team_id is not None and team_id in self.team_pattern_routers:
|
||||
potential_team_only_wildcard_models: Final = self.team_pattern_routers[team_id].route(model_name) or []
|
||||
potential_team_only_wildcard_models: Final = self.team_pattern_routers[
|
||||
team_id
|
||||
].get_deployments_by_pattern(model=model_name or "")
|
||||
potential_wildcard_models.extend(potential_team_only_wildcard_models)
|
||||
|
||||
if model_name is not None and potential_wildcard_models is not None:
|
||||
|
|
|
|||
|
|
@ -2,12 +2,13 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import StreamingResponse
|
||||
from e2e_http import NoBody, StreamingResponse, is_ok, unwrap
|
||||
from models import (
|
||||
ChatBody,
|
||||
ChatMessage,
|
||||
|
|
@ -15,9 +16,16 @@ from models import (
|
|||
LiteLLMParamsBody,
|
||||
ModelInfoBody,
|
||||
ModelNewBody,
|
||||
TeamDeleteBody,
|
||||
TeamInfoParams,
|
||||
TeamInfoResponse,
|
||||
TeamNewBody,
|
||||
TeamNewResponse,
|
||||
TeamUpdateBody,
|
||||
)
|
||||
|
||||
MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied"
|
||||
TEAM_MODEL_ACCESS_DENIED_MARKER = "team_model_access_denied"
|
||||
ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route"
|
||||
|
||||
|
||||
|
|
@ -31,6 +39,14 @@ class ApiErrorEnvelope(BaseModel):
|
|||
error: ApiErrorDetail
|
||||
|
||||
|
||||
class AccessGroupInfoResponse(BaseModel):
|
||||
"""GET /access_group/{name}/info: the deployments a model access group grants."""
|
||||
|
||||
access_group: str
|
||||
model_names: list[str]
|
||||
deployment_count: int
|
||||
|
||||
|
||||
def error_envelope(body: str) -> ApiErrorEnvelope | None:
|
||||
"""The OpenAI-shaped `{"error": {...}}` a client parses, or None if absent."""
|
||||
try:
|
||||
|
|
@ -51,15 +67,75 @@ class AccessControlClient:
|
|||
def delete_key(self, key: str) -> None:
|
||||
self.proxy.delete_key(key)
|
||||
|
||||
def chat_status(self, key: str, model: str, content: str) -> StreamingResponse:
|
||||
def chat_status(
|
||||
self, key: str, model: str, content: str, max_completion_tokens: int | None = None
|
||||
) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/chat/completions",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=ChatBody(
|
||||
model=model, messages=[ChatMessage(role="user", content=content)]
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=content)],
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
def create_team(self, team_alias: str, models: list[str]) -> str:
|
||||
team_id = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/team/new",
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamNewBody(team_alias=team_alias, models=models),
|
||||
response_type=TeamNewResponse,
|
||||
)
|
||||
).team_id
|
||||
self._await_team(team_id)
|
||||
return team_id
|
||||
|
||||
def set_team_models(self, team_id: str, team_alias: str, models: list[str]) -> None:
|
||||
"""Replace the team's allow-list. /model/new appends a team-scoped deployment's
|
||||
public name to it, so a test that means to grant only an access group has to
|
||||
put the allow-list back afterwards."""
|
||||
_ = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/team/update",
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamUpdateBody(team_id=team_id, team_alias=team_alias, models=models),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
|
||||
def delete_team(self, team_id: str) -> None:
|
||||
_ = self.proxy.transport.post(
|
||||
"/team/delete",
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamDeleteBody(team_ids=[team_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def access_group_info(self, access_group: str) -> AccessGroupInfoResponse | None:
|
||||
result = self.proxy.transport.get(
|
||||
f"/access_group/{access_group}/info",
|
||||
headers=self.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=AccessGroupInfoResponse,
|
||||
)
|
||||
return unwrap(result) if is_ok(result) else None
|
||||
|
||||
def _await_team(self, team_id: str) -> None:
|
||||
deadline = time.monotonic() + self.proxy.poll_timeout
|
||||
while time.monotonic() < deadline:
|
||||
result = self.proxy.transport.get(
|
||||
"/team/info",
|
||||
headers=self.proxy.transport.master,
|
||||
params=TeamInfoParams(team_id=team_id),
|
||||
response_type=TeamInfoResponse,
|
||||
)
|
||||
if is_ok(result):
|
||||
return
|
||||
time.sleep(self.proxy.poll_interval)
|
||||
raise AssertionError(f"/team/info never resolved team {team_id!r} created by /team/new")
|
||||
|
||||
def create_model_status(self, key: str, model_name: str) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/model/new",
|
||||
|
|
|
|||
277
tests/e2e/access_control/test_model_access_group_e2e.py
Normal file
277
tests/e2e/access_control/test_model_access_group_e2e.py
Normal file
|
|
@ -0,0 +1,277 @@
|
|||
"""Live e2e: a model access group as the grant on a key and on a team.
|
||||
|
||||
Whoever holds the group can call every deployment in it and nothing else, whether
|
||||
the request names a deployment exactly, names a model that a wildcard deployment
|
||||
in the group covers, or spells that model with its provider prefix. The bare-name
|
||||
spelling is the LIT-5813 regression: the group-membership lookup skipped the
|
||||
provider-prefix retry every other model-resolution path performs, so a group
|
||||
holding `openai/gpt-5.4*` denied `gpt-5.4-nano` while allowing `openai/gpt-5.4-nano`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Callable, Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from access_control_client import (
|
||||
AccessControlClient,
|
||||
MODEL_ACCESS_DENIED_MARKER,
|
||||
TEAM_MODEL_ACCESS_DENIED_MARKER,
|
||||
)
|
||||
from e2e_config import unique_marker
|
||||
from lifecycle import ResourceManager
|
||||
from models import (
|
||||
ChatResponse,
|
||||
KeyGenerateBody,
|
||||
LiteLLMParamsBody,
|
||||
ModelInfoBody,
|
||||
ModelNewBody,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
WILDCARD_PATTERN: Final = "openai/gpt-5.4*"
|
||||
WILDCARD_BARE_MODEL: Final = "gpt-5.4-nano"
|
||||
WILDCARD_PREFIXED_MODEL: Final = "openai/gpt-5.4-nano"
|
||||
GROUP_BACKEND: Final = "openai/gpt-5.4-nano"
|
||||
UNCOVERED_OPENAI_MODEL: Final = "gpt-5.2"
|
||||
|
||||
TEAM_WILDCARD_PATTERN: Final = "openai/gpt-5.6*"
|
||||
TEAM_WILDCARD_BARE_MODEL: Final = "gpt-5.6-luna"
|
||||
|
||||
MAX_COMPLETION_TOKENS: Final = 16
|
||||
PROMPT: Final = "Reply with exactly: OK"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GroupedDeployments:
|
||||
"""A wildcard deployment and an exactly-named one inside `access_group`, plus a
|
||||
deployment left out of it."""
|
||||
|
||||
access_group: str
|
||||
member_model: str
|
||||
outsider_model: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TeamGrant:
|
||||
"""A team whose whole allow-list is `access_group`, holding one team-scoped
|
||||
wildcard deployment, and a key that belongs to it."""
|
||||
|
||||
access_group: str
|
||||
team_id: str
|
||||
key: str
|
||||
|
||||
|
||||
ModelSelector = Callable[[GroupedDeployments], str]
|
||||
|
||||
ALLOWED: Final[tuple[tuple[str, ModelSelector], ...]] = (
|
||||
("bare name the group's wildcard covers", lambda grouped: WILDCARD_BARE_MODEL),
|
||||
("provider-prefixed name the group's wildcard covers", lambda grouped: WILDCARD_PREFIXED_MODEL),
|
||||
("exactly-named deployment in the group", lambda grouped: grouped.member_model),
|
||||
)
|
||||
|
||||
DENIED: Final[tuple[tuple[str, ModelSelector], ...]] = (
|
||||
("deployment outside the group", lambda grouped: grouped.outsider_model),
|
||||
("provider model outside the group's wildcard", lambda grouped: UNCOVERED_OPENAI_MODEL),
|
||||
("name no provider claims", lambda grouped: f"e2e-ag-unknown-{unique_marker()}"),
|
||||
)
|
||||
|
||||
|
||||
def _provider_key(env_var: str) -> str:
|
||||
return os.environ.get(env_var) or f"os.environ/{env_var}"
|
||||
|
||||
|
||||
def _grouped_model(model_name: str, backend: str, access_groups: list[str] | None) -> ModelNewBody:
|
||||
return ModelNewBody(
|
||||
model_name=model_name,
|
||||
litellm_params=LiteLLMParamsBody(model=backend, api_key=_provider_key("OPENAI_API_KEY")),
|
||||
model_info=ModelInfoBody(access_groups=access_groups),
|
||||
)
|
||||
|
||||
|
||||
def _await_group_members(client: AccessControlClient, access_group: str, expected: frozenset[str]) -> None:
|
||||
"""The grant under test is the group's membership, so prove the proxy recorded it
|
||||
before asserting on what the group lets through."""
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
listed: list[str] = []
|
||||
while time.monotonic() < deadline:
|
||||
info = client.access_group_info(access_group)
|
||||
listed = info.model_names if info is not None else []
|
||||
if expected.issubset(listed):
|
||||
return
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
pytest.fail(
|
||||
f"/access_group/{access_group}/info never listed {sorted(expected)} as members; last read {listed}"
|
||||
)
|
||||
|
||||
|
||||
def _await_team_allowlist(client: AccessControlClient, grant_key: str, access_group: str) -> None:
|
||||
"""Registering a team-scoped deployment appends its public name to the team's
|
||||
allow-list, and a wildcard sitting there directly would grant the model under test
|
||||
on its own. Poll a denial until the message enumerates the allow-list the test
|
||||
means to exercise: the group, and nothing else."""
|
||||
allowlist: Final = f"models=['{access_group}']"
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
body = ""
|
||||
while time.monotonic() < deadline:
|
||||
body = client.chat_status(
|
||||
grant_key, UNCOVERED_OPENAI_MODEL, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS
|
||||
).body
|
||||
if allowlist in body:
|
||||
return
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
pytest.fail(f"the team's allow-list never settled to {allowlist}; last denial read {body[:300]}")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def grouped(client: AccessControlClient) -> Iterator[GroupedDeployments]:
|
||||
marker: Final = unique_marker()
|
||||
deployments: Final = GroupedDeployments(
|
||||
access_group=f"e2e-ag-{marker}",
|
||||
member_model=f"e2e-ag-member-{marker}",
|
||||
outsider_model=f"e2e-ag-outsider-{marker}",
|
||||
)
|
||||
registrations: Final = (
|
||||
_grouped_model(WILDCARD_PATTERN, WILDCARD_PATTERN, [deployments.access_group]),
|
||||
_grouped_model(deployments.member_model, GROUP_BACKEND, [deployments.access_group]),
|
||||
_grouped_model(deployments.outsider_model, GROUP_BACKEND, None),
|
||||
)
|
||||
created: Final = tuple(client.proxy.register_model(body) for body in registrations)
|
||||
try:
|
||||
_await_group_members(
|
||||
client,
|
||||
deployments.access_group,
|
||||
frozenset({WILDCARD_PATTERN, deployments.member_model}),
|
||||
)
|
||||
yield deployments
|
||||
finally:
|
||||
for model_id in created:
|
||||
client.proxy.delete_model(model_id)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def team_grant(client: AccessControlClient) -> Iterator[TeamGrant]:
|
||||
marker: Final = unique_marker()
|
||||
access_group: Final = f"e2e-agt-{marker}"
|
||||
team_alias: Final = f"e2e-ag-team-{marker}"
|
||||
team_id: Final = client.create_team(team_alias, [access_group])
|
||||
key: Final = client.proxy.generate_key(KeyGenerateBody(models=[], team_id=team_id))
|
||||
model_id: Final = client.proxy.register_model(
|
||||
ModelNewBody(
|
||||
model_name=TEAM_WILDCARD_PATTERN,
|
||||
litellm_params=LiteLLMParamsBody(
|
||||
model=TEAM_WILDCARD_PATTERN, api_key=_provider_key("OPENAI_API_KEY")
|
||||
),
|
||||
model_info=ModelInfoBody(team_id=team_id, access_groups=[access_group]),
|
||||
),
|
||||
listed_for=key,
|
||||
)
|
||||
client.set_team_models(team_id, team_alias, [access_group])
|
||||
try:
|
||||
_await_team_allowlist(client, key, access_group)
|
||||
yield TeamGrant(access_group=access_group, team_id=team_id, key=key)
|
||||
finally:
|
||||
client.proxy.delete_model(model_id)
|
||||
client.proxy.delete_key(key)
|
||||
client.delete_team(team_id)
|
||||
|
||||
|
||||
class TestKeyScopedToAccessGroup:
|
||||
@pytest.mark.covers(
|
||||
"other.auth.model_access_group.wildcard_bare_name_allowed",
|
||||
"other.auth.model_access_group.member_allowed",
|
||||
)
|
||||
@pytest.mark.parametrize(("case", "select_model"), ALLOWED)
|
||||
def test_group_grants_every_deployment_in_it(
|
||||
self,
|
||||
case: str,
|
||||
select_model: ModelSelector,
|
||||
client: AccessControlClient,
|
||||
resources: ResourceManager,
|
||||
grouped: GroupedDeployments,
|
||||
) -> None:
|
||||
key = resources.key(models=[grouped.access_group])
|
||||
model = select_model(grouped)
|
||||
|
||||
result = client.chat_status(
|
||||
key, model, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS
|
||||
)
|
||||
|
||||
assert result.status_code == 200, (
|
||||
f"a key holding access group {grouped.access_group!r} must be able to call "
|
||||
f"{model!r} ({case}), got {result.status_code}: {result.body[:300]}"
|
||||
)
|
||||
assert ChatResponse.model_validate_json(result.body).choices, (
|
||||
f"200 must carry a real completion, not an error envelope: {result.body[:300]}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("other.auth.model_access_group.non_member_denied")
|
||||
@pytest.mark.parametrize(("case", "select_model"), DENIED)
|
||||
def test_group_grants_nothing_outside_it(
|
||||
self,
|
||||
case: str,
|
||||
select_model: ModelSelector,
|
||||
client: AccessControlClient,
|
||||
resources: ResourceManager,
|
||||
grouped: GroupedDeployments,
|
||||
) -> None:
|
||||
key = resources.key(models=[grouped.access_group])
|
||||
model = select_model(grouped)
|
||||
|
||||
result = client.chat_status(
|
||||
key, model, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS
|
||||
)
|
||||
|
||||
assert result.status_code == 403, (
|
||||
f"a key holding only access group {grouped.access_group!r} must be denied 403 on "
|
||||
f"{model!r} ({case}), got {result.status_code}: {result.body[:300]}"
|
||||
)
|
||||
assert MODEL_ACCESS_DENIED_MARKER in result.body, (
|
||||
f"403 body must be a key model-access denial, got: {result.body[:300]}"
|
||||
)
|
||||
|
||||
|
||||
class TestTeamScopedToAccessGroup:
|
||||
@pytest.mark.covers("other.auth.model_access_group.team_wildcard_bare_name_allowed")
|
||||
def test_group_grants_the_teams_own_wildcard(
|
||||
self, client: AccessControlClient, team_grant: TeamGrant
|
||||
) -> None:
|
||||
result = client.chat_status(
|
||||
team_grant.key,
|
||||
TEAM_WILDCARD_BARE_MODEL,
|
||||
f"{PROMPT} {unique_marker()}",
|
||||
MAX_COMPLETION_TOKENS,
|
||||
)
|
||||
|
||||
assert result.status_code == 200, (
|
||||
f"a team whose allow-list is access group {team_grant.access_group!r} must be able to "
|
||||
f"call {TEAM_WILDCARD_BARE_MODEL!r} through its team-scoped {TEAM_WILDCARD_PATTERN!r} "
|
||||
f"deployment, got {result.status_code}: {result.body[:300]}"
|
||||
)
|
||||
assert ChatResponse.model_validate_json(result.body).choices, (
|
||||
f"200 must carry a real completion, not an error envelope: {result.body[:300]}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("other.auth.model_access_group.team_non_member_denied")
|
||||
def test_group_grants_the_team_nothing_outside_it(
|
||||
self, client: AccessControlClient, team_grant: TeamGrant
|
||||
) -> None:
|
||||
model = f"e2e-ag-unknown-{unique_marker()}"
|
||||
|
||||
result = client.chat_status(
|
||||
team_grant.key, model, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS
|
||||
)
|
||||
|
||||
assert result.status_code == 403, (
|
||||
f"a team holding only access group {team_grant.access_group!r} must be denied 403 on "
|
||||
f"{model!r}, got {result.status_code}: {result.body[:300]}"
|
||||
)
|
||||
assert TEAM_MODEL_ACCESS_DENIED_MARKER in result.body, (
|
||||
f"403 body must be a team model-access denial, got: {result.body[:300]}"
|
||||
)
|
||||
|
|
@ -12,6 +12,11 @@
|
|||
- {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:77-150", rationale: "Valid JWT with correct issuer + claims grants access"}
|
||||
- {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:125-135", rationale: "Expired JWT rejected even with valid signature"}
|
||||
- {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:145-150", rationale: "Bad/missing signature fails verification"}
|
||||
- {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"}
|
||||
- {id: other.auth.model_access_group.team_wildcard_bare_name_allowed, module: other, tier: P1, area: auth, assertions: [team_wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "The same bare-name grant holds when the wildcard deployment is team-scoped and the team's allow-list is the group"}
|
||||
- {id: other.auth.model_access_group.team_non_member_denied, module: other, tier: P1, area: auth, assertions: [team_non_member_denied], source: "auth_checks.py:3232", rationale: "A team-level group grant reaches nothing outside the group"}
|
||||
- {id: other.auth.virtual_key.route_permission_enforced, module: other, tier: P0, area: auth, assertions: [route_permission_enforced], source: "route_checks.py:89-151", rationale: "allowed_routes whitelist denies disallowed routes"}
|
||||
- {id: other.auth.virtual_key.route_group_allowed, module: other, tier: P1, area: auth, assertions: [route_group_allowed], source: "route_checks.py:106-128", rationale: "allowed_routes=[llm_api_routes] grants all LLM endpoints"}
|
||||
- {id: other.auth.passthrough.model_allowlist_enforced, module: other, tier: P1, area: auth, assertions: [model_allowlist_enforced], source: "route_checks.py:135-151", rationale: "Passthrough enforces per-key model allow-lists"}
|
||||
|
|
|
|||
|
|
@ -766,6 +766,8 @@ class ModelInfoBody(BaseModel):
|
|||
# constraint when a prior run's teardown had not removed the row.
|
||||
id: str | None = None
|
||||
mode: ModelMode | None = None
|
||||
access_groups: list[str] | None = None
|
||||
team_id: str | None = None
|
||||
|
||||
|
||||
class ModelNewBody(BaseModel):
|
||||
|
|
@ -861,6 +863,7 @@ class TeamNewResponse(BaseModel):
|
|||
class TeamUpdateBody(BaseModel):
|
||||
team_id: str
|
||||
team_alias: str
|
||||
models: list[str] | None = None
|
||||
|
||||
|
||||
class TeamInfoParams(BaseModel):
|
||||
|
|
|
|||
|
|
@ -278,7 +278,21 @@ class ProxyClient:
|
|||
mode: ModelMode | None = None,
|
||||
) -> str:
|
||||
"""Register a deployment under `model_name` and return its proxy-assigned
|
||||
model_id, once the model is actually servable on the data plane.
|
||||
model_id, once the model is actually servable on the data plane."""
|
||||
return self.register_model(
|
||||
ModelNewBody(
|
||||
model_name=model_name,
|
||||
litellm_params=litellm_params,
|
||||
model_info=ModelInfoBody(mode=mode),
|
||||
)
|
||||
)
|
||||
|
||||
def register_model(self, body: ModelNewBody, listed_for: str | None = None) -> str:
|
||||
"""`create_model` for deployments that carry more than a mode: access groups,
|
||||
team scoping, a pinned id. `listed_for` is the virtual key whose /v1/models
|
||||
view must list the deployment before it counts as servable, because a
|
||||
team-scoped deployment is listed to its own team and to nobody else, master
|
||||
key included; leave it unset for a proxy-wide model.
|
||||
|
||||
/model/new is a control-plane route; the data plane (which serves /chat,
|
||||
/ocr, ...) only picks the new model up on its next DB reload, so a call
|
||||
|
|
@ -296,25 +310,22 @@ class ProxyClient:
|
|||
self.transport.post(
|
||||
"/model/new",
|
||||
headers=self.transport.master,
|
||||
json=ModelNewBody(
|
||||
model_name=model_name,
|
||||
litellm_params=litellm_params,
|
||||
model_info=ModelInfoBody(mode=mode),
|
||||
),
|
||||
json=body,
|
||||
response_type=ModelNewResponse,
|
||||
)
|
||||
).model_id
|
||||
written_at = time.monotonic()
|
||||
self._await_model_servable(model_name)
|
||||
self._await_model_servable(body.model_name, listed_for)
|
||||
settle_propagation(written_at)
|
||||
return model_id
|
||||
|
||||
def _await_model_servable(self, model_name: str) -> None:
|
||||
def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None:
|
||||
"""Block until the data plane lists `model_name`, or fail at model_servable_timeout."""
|
||||
headers = self.transport.master if listed_for is None else self.transport.bearer(listed_for)
|
||||
outcome = await_servable(
|
||||
lambda poll_timeout: self.transport.get(
|
||||
"/v1/models",
|
||||
headers=self.transport.master,
|
||||
headers=headers,
|
||||
params=NoBody(),
|
||||
response_type=ModelsListResponse,
|
||||
timeout=poll_timeout,
|
||||
|
|
|
|||
|
|
@ -6392,3 +6392,168 @@ def test_is_user_proxy_admin_rejects_view_only_admin():
|
|||
assert _is_user_proxy_admin(user_obj=viewer) is False
|
||||
assert _is_user_proxy_admin(user_obj=admin) is True
|
||||
assert _is_user_proxy_admin(user_obj=None) is False
|
||||
|
||||
|
||||
def _make_wildcard_access_group_router():
|
||||
"""
|
||||
`openai/*` tagged into an access group, plus an untagged `azure/*`, mirroring a
|
||||
proxy that fronts a whole provider behind one wildcard deployment.
|
||||
"""
|
||||
from litellm import Router
|
||||
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "fake"},
|
||||
"model_info": {
|
||||
"id": "wildcard-openai",
|
||||
"access_groups": ["default-models"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "azure/*",
|
||||
"litellm_params": {"model": "azure/*", "api_key": "fake"},
|
||||
"model_info": {"id": "wildcard-azure"},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_can_object_call_model_access_group_wildcard_accepts_bare_model_name():
|
||||
"""
|
||||
Regression: a key holding only the access group name was denied for `gpt-4o`
|
||||
while `openai/gpt-4o` was allowed, because group membership resolved through the
|
||||
pattern router's raw regex and skipped the `{provider}/{model}` retry that both
|
||||
routing and the direct-wildcard grant already perform.
|
||||
"""
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model
|
||||
|
||||
router = _make_wildcard_access_group_router()
|
||||
|
||||
assert (
|
||||
_can_object_call_model(
|
||||
model="gpt-4o",
|
||||
llm_router=router,
|
||||
models=["default-models"],
|
||||
object_type="key",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_can_object_call_model_access_group_wildcard_accepts_prefixed_model_name():
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model
|
||||
|
||||
router = _make_wildcard_access_group_router()
|
||||
|
||||
assert (
|
||||
_can_object_call_model(
|
||||
model="openai/gpt-4o",
|
||||
llm_router=router,
|
||||
models=["default-models"],
|
||||
object_type="key",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"totally-made-up-model-zzz", # no provider can be inferred
|
||||
"azure/some-deployment", # wildcard exists but carries no access group
|
||||
],
|
||||
)
|
||||
def test_can_object_call_model_access_group_wildcard_does_not_over_grant(model):
|
||||
"""The bare-name retry must not turn an access group into a blanket grant."""
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model
|
||||
|
||||
router = _make_wildcard_access_group_router()
|
||||
|
||||
with pytest.raises(ProxyException):
|
||||
_can_object_call_model(
|
||||
model=model,
|
||||
llm_router=router,
|
||||
models=["default-models"],
|
||||
object_type="key",
|
||||
)
|
||||
|
||||
|
||||
def test_can_object_call_model_access_group_rejects_unconsumed_namespace():
|
||||
"""
|
||||
`bedrockz/...` infers provider `bedrock` from a fragment of the name, so
|
||||
re-prefixing would smuggle an unrecognized namespace through a `bedrock/*` group.
|
||||
"""
|
||||
from litellm import Router
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock/*",
|
||||
"litellm_params": {"model": "bedrock/*"},
|
||||
"model_info": {
|
||||
"id": "wildcard-bedrock",
|
||||
"access_groups": ["bedrock-models"],
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
_can_object_call_model(
|
||||
model="anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
llm_router=router,
|
||||
models=["bedrock-models"],
|
||||
object_type="key",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException):
|
||||
_can_object_call_model(
|
||||
model="bedrockz/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
llm_router=router,
|
||||
models=["bedrock-models"],
|
||||
object_type="key",
|
||||
)
|
||||
|
||||
|
||||
def test_can_object_call_model_team_scoped_wildcard_accepts_bare_model_name():
|
||||
"""
|
||||
Same regression as the proxy-wide wildcard, but for a team-scoped deployment
|
||||
whose public name is a wildcard: those live in a separate per-team pattern
|
||||
index that needed the same `{provider}/{model}` retry.
|
||||
"""
|
||||
from litellm import Router
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/*_team-a_abc",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "fake"},
|
||||
"model_info": {
|
||||
"id": "team-byok-wildcard",
|
||||
"team_id": "team-a",
|
||||
"team_public_model_name": "openai/*",
|
||||
"access_groups": ["team-models"],
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
for model in ("gpt-4o", "openai/gpt-4o"):
|
||||
assert (
|
||||
_can_object_call_model(
|
||||
model=model,
|
||||
llm_router=router,
|
||||
models=["team-models"],
|
||||
object_type="team",
|
||||
team_id="team-a",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue