mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
feat(keys): filter /key/list by tag
Adds a tag query parameter that returns only keys whose metadata.tags holds that exact tag name. It is ANDed on top of the caller's visibility, so non-admins still only see keys they could already see
This commit is contained in:
parent
615ed7900f
commit
d7ca77e255
4 changed files with 81 additions and 3 deletions
|
|
@ -158,6 +158,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
|||
CustomKeyPolicyRequest,
|
||||
FailedKeyUpdate,
|
||||
KeySearchWhere,
|
||||
KeyTagWhere,
|
||||
SuccessfulKeyUpdate,
|
||||
)
|
||||
from litellm.types.router import Deployment
|
||||
|
|
@ -6525,6 +6526,7 @@ async def list_keys(
|
|||
project_id: str | None = Query(None, description="Filter keys by project ID"),
|
||||
access_group_id: str | None = Query(None, description="Filter keys by access group ID"),
|
||||
agent_id: str | None = Query(None, description="Filter keys by agent ID"),
|
||||
tag: str | None = Query(None, description="Filter keys by tag name. Exact match on an entry of metadata.tags"),
|
||||
substring_matching: bool = Query(
|
||||
False,
|
||||
description="If true, match key_alias (any caller) and user_id (proxy admins only) as case-insensitive substrings instead of exact values. Defaults to false: /key/list matched these exactly before substring search was added, and an exact user_id filter must never return another user's keys.",
|
||||
|
|
@ -6655,6 +6657,7 @@ async def list_keys(
|
|||
project_id=project_id,
|
||||
access_group_id=access_group_id,
|
||||
agent_id=agent_id,
|
||||
tag=tag,
|
||||
use_substring_matching=use_substring_matching,
|
||||
use_key_alias_substring_matching=use_key_alias_substring_matching,
|
||||
expires_filter=expires if isinstance(expires, str) else None,
|
||||
|
|
@ -6900,6 +6903,11 @@ def _build_key_search_where(search: str) -> KeySearchWhere:
|
|||
return search_where
|
||||
|
||||
|
||||
def _build_key_tag_where(tag: str) -> KeyTagWhere:
|
||||
tag_where: Final[KeyTagWhere] = {"metadata": {"path": ("tags",), "array_contains": json.dumps((tag,))}}
|
||||
return tag_where
|
||||
|
||||
|
||||
def _build_key_filter_conditions(
|
||||
user_id: str | None,
|
||||
team_id: str | None,
|
||||
|
|
@ -6913,6 +6921,7 @@ def _build_key_filter_conditions(
|
|||
project_id: str | None = None,
|
||||
access_group_id: str | None = None,
|
||||
agent_id: str | None = None,
|
||||
tag: str | None = None,
|
||||
use_substring_matching: bool = False,
|
||||
use_key_alias_substring_matching: bool = False,
|
||||
expires_filter: str | None = None,
|
||||
|
|
@ -7005,7 +7014,7 @@ def _build_key_filter_conditions(
|
|||
elif len(or_conditions) == 1:
|
||||
where.update(or_conditions[0])
|
||||
|
||||
# Apply team_id, project_id and access_group_id as global AND filters so they
|
||||
# Apply team_id, project_id, access_group_id and tag as global AND filters so they
|
||||
# narrow results across all visibility conditions (own keys, team keys, etc.)
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
status_where: Final = _build_status_where_clause(status_filter, now)
|
||||
|
|
@ -7025,6 +7034,7 @@ def _build_key_filter_conditions(
|
|||
*(({"project_id": project_id},) if project_id else ()),
|
||||
*(({"access_group_ids": {"hasSome": [access_group_id]}},) if access_group_id else ()),
|
||||
*(({"agent_id": agent_id},) if agent_id and isinstance(agent_id, str) else ()),
|
||||
*((_build_key_tag_where(tag),) if tag and isinstance(tag, str) else ()),
|
||||
*(
|
||||
(_build_expires_where_clause(expires_filter, now),)
|
||||
if expires_filter is not None and expires_filter in VALID_EXPIRES_FILTER_VALUES
|
||||
|
|
@ -7059,6 +7069,7 @@ async def _list_key_helper(
|
|||
project_id: str | None = None,
|
||||
access_group_id: str | None = None,
|
||||
agent_id: str | None = None,
|
||||
tag: str | None = None,
|
||||
use_substring_matching: bool = False,
|
||||
use_key_alias_substring_matching: bool = False,
|
||||
expires_filter: str | None = None,
|
||||
|
|
@ -7099,6 +7110,7 @@ async def _list_key_helper(
|
|||
project_id=project_id,
|
||||
access_group_id=access_group_id,
|
||||
agent_id=agent_id,
|
||||
tag=tag,
|
||||
use_substring_matching=use_substring_matching,
|
||||
use_key_alias_substring_matching=use_key_alias_substring_matching,
|
||||
expires_filter=expires_filter,
|
||||
|
|
|
|||
|
|
@ -29,6 +29,17 @@ class KeySearchWhere(TypedDict):
|
|||
OR: ReadOnly[tuple[KeyTokenWhere, KeyAliasContainsWhere]]
|
||||
|
||||
|
||||
class JsonPathArrayContains(TypedDict):
|
||||
path: ReadOnly[tuple[str, ...]]
|
||||
array_contains: ReadOnly[str]
|
||||
|
||||
|
||||
class KeyTagWhere(TypedDict):
|
||||
"""Prisma filter behind `/key/list?tag=`: keys whose metadata.tags array holds the exact tag."""
|
||||
|
||||
metadata: ReadOnly[JsonPathArrayContains]
|
||||
|
||||
|
||||
class BulkUpdateKeyRequestItem(BaseModel):
|
||||
"""One /key/bulk_update item; only the fields it carries are written."""
|
||||
|
||||
|
|
|
|||
|
|
@ -17182,6 +17182,12 @@ def _prisma_where_matches(row, where):
|
|||
elif field == "OR":
|
||||
if not any(_prisma_where_matches(row, child) for child in expected):
|
||||
return False
|
||||
elif isinstance(expected, dict) and "array_contains" in expected:
|
||||
value = getattr(row, field)
|
||||
for key in expected["path"]:
|
||||
value = value.get(key) if isinstance(value, dict) else None
|
||||
if not isinstance(value, list) or not all(item in value for item in json.loads(expected["array_contains"])):
|
||||
return False
|
||||
elif isinstance(expected, dict):
|
||||
value = getattr(row, field)
|
||||
if "in" in expected and value not in expected["in"]:
|
||||
|
|
@ -17222,14 +17228,14 @@ _TEAM_A_KEYS = (
|
|||
)
|
||||
|
||||
|
||||
def _list_team_a_keys_as(user_role, members_with_roles, query):
|
||||
def _list_team_a_keys_as(user_role, members_with_roles, query, rows=_TEAM_A_KEYS):
|
||||
from fastapi import FastAPI
|
||||
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import router
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken = _InMemoryVerificationTokenTable(_TEAM_A_KEYS)
|
||||
mock_prisma_client.db.litellm_verificationtoken = _InMemoryVerificationTokenTable(rows)
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=LiteLLM_UserTable(user_id="alice", teams=["team-a"], organization_memberships=[])
|
||||
)
|
||||
|
|
@ -17276,6 +17282,53 @@ def test_list_keys_key_alias_stays_exact_without_substring_matching():
|
|||
) == ["tok-alice-first"]
|
||||
|
||||
|
||||
def _tagged_team_key(token, key_alias, user_id, metadata):
|
||||
return LiteLLM_VerificationToken(
|
||||
token=token, key_alias=key_alias, user_id=user_id, team_id="team-a", metadata=metadata
|
||||
)
|
||||
|
||||
|
||||
_TAGGED_TEAM_A_KEYS = (
|
||||
_tagged_team_key("tok-alice-batch", "alice-batch", "alice", {"tags": ["prod-batch", "nightly"]}),
|
||||
_tagged_team_key("tok-bob-batch", "bob-batch", "bob", {"tags": ["prod-batch"]}),
|
||||
_tagged_team_key("tok-svc-batch", "svc-batch", None, {"tags": ["prod-batch"]}),
|
||||
_tagged_team_key("tok-alice-similar", "alice-similar", "alice", {"tags": ["prod-batch-2"]}),
|
||||
_tagged_team_key("tok-alice-scalar", "alice-scalar", "alice", {"tags": "prod-batch"}),
|
||||
_tagged_team_key("tok-alice-untagged", "alice-untagged", "alice", {}),
|
||||
)
|
||||
|
||||
|
||||
def test_list_keys_tag_returns_only_keys_carrying_that_exact_tag():
|
||||
keys = _list_team_a_keys_as(
|
||||
LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_ADMIN, "tag=prod-batch", rows=_TAGGED_TEAM_A_KEYS
|
||||
)
|
||||
assert keys == ["tok-alice-batch", "tok-bob-batch", "tok-svc-batch"]
|
||||
|
||||
|
||||
def test_list_keys_tag_composes_with_other_filters():
|
||||
keys = _list_team_a_keys_as(
|
||||
LitellmUserRoles.INTERNAL_USER,
|
||||
_ALICE_TEAM_ADMIN,
|
||||
"tag=prod-batch&key_alias=alice&substring_matching=true",
|
||||
rows=_TAGGED_TEAM_A_KEYS,
|
||||
)
|
||||
assert keys == ["tok-alice-batch"]
|
||||
|
||||
|
||||
def test_list_keys_tag_stays_within_caller_visibility():
|
||||
keys = _list_team_a_keys_as(
|
||||
LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_MEMBER, "tag=prod-batch", rows=_TAGGED_TEAM_A_KEYS
|
||||
)
|
||||
assert keys == ["tok-alice-batch", "tok-svc-batch"]
|
||||
|
||||
|
||||
def test_list_keys_tag_without_keys_returns_none():
|
||||
keys = _list_team_a_keys_as(
|
||||
LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_ADMIN, "tag=unused-tag", rows=_TAGGED_TEAM_A_KEYS
|
||||
)
|
||||
assert keys == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_keys_search_is_honored_for_non_admin():
|
||||
"""LIT-4741: unlike substring_matching, `search` is not admin-gated. A non-admin's
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -61186,6 +61186,8 @@ export interface operations {
|
|||
access_group_id?: string | null;
|
||||
/** @description Filter keys by agent ID */
|
||||
agent_id?: string | null;
|
||||
/** @description Filter keys by tag name. Exact match on an entry of metadata.tags */
|
||||
tag?: string | null;
|
||||
/** @description If true, match key_alias (any caller) and user_id (proxy admins only) as case-insensitive substrings instead of exact values. Defaults to false: /key/list matched these exactly before substring search was added, and an exact user_id filter must never return another user's keys. */
|
||||
substring_matching?: boolean;
|
||||
/** @description Filter keys by expiration. 'expired' returns keys whose expires is in the past; 'active' returns keys that never expire or expire in the future. Omit to return keys regardless of expiration. */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue