mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
feat(access-groups): grant vector stores through access groups (#50290)
This commit is contained in:
shrey
GitHub
parent
41fa625322
commit
6685525fa2
17 changed files with 1044 additions and 18 deletions
|
|
@ -20,6 +20,7 @@ class LiteLLM_AccessGroupTable(LiteLLMPydanticObjectBase):
|
|||
mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
|
||||
access_agent_ids: list[str] = []
|
||||
access_skill_names: tuple[str, ...] = ()
|
||||
access_vector_store_ids: tuple[str, ...] = ()
|
||||
assigned_team_ids: list[str] = []
|
||||
assigned_key_ids: list[str] = []
|
||||
created_at: datetime | None = None
|
||||
|
|
|
|||
|
|
@ -137,6 +137,10 @@ class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable):
|
|||
last_refreshed_at: float | None = None
|
||||
|
||||
|
||||
class LiteLLM_TeamRebuiltFromToken(LiteLLM_TeamTableCachedObj): # frozen-ok: same mutability as its cached base
|
||||
"""Auth's stand-in, built from the key's fields when the team row was unreadable. It has no access groups."""
|
||||
|
||||
|
||||
class LiteLLM_DeletedTeamTable(LiteLLM_TeamTable):
|
||||
"""Audit record for deleted teams; mirrors the team plus deletion metadata."""
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,4 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "access_vector_store_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
|
||||
|
|
@ -1630,6 +1630,7 @@ model LiteLLM_AccessGroupTable {
|
|||
mcp_tool_permissions Json?
|
||||
access_agent_ids String[] @default([])
|
||||
access_skill_names String[] @default([])
|
||||
access_vector_store_ids String[] @default([])
|
||||
|
||||
assigned_team_ids String[] @default([])
|
||||
assigned_key_ids String[] @default([])
|
||||
|
|
|
|||
|
|
@ -939,6 +939,20 @@
|
|||
],
|
||||
"title": "Access Skill Names"
|
||||
},
|
||||
"access_vector_store_ids": {
|
||||
"anyOf": [
|
||||
{
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Access Vector Store Ids"
|
||||
},
|
||||
"assigned_key_ids": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -1128,6 +1142,14 @@
|
|||
"title": "Access Skill Names",
|
||||
"type": "array"
|
||||
},
|
||||
"access_vector_store_ids": {
|
||||
"default": [],
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"title": "Access Vector Store Ids",
|
||||
"type": "array"
|
||||
},
|
||||
"assigned_key_ids": {
|
||||
"items": {
|
||||
"type": "string"
|
||||
|
|
@ -1304,6 +1326,20 @@
|
|||
],
|
||||
"title": "Access Skill Names"
|
||||
},
|
||||
"access_vector_store_ids": {
|
||||
"anyOf": [
|
||||
{
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Access Vector Store Ids"
|
||||
},
|
||||
"assigned_key_ids": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ from litellm.llms.base_llm.managed_resources.utils import (
|
|||
is_base64_encoded_unified_id,
|
||||
)
|
||||
from litellm.models.project import LiteLLM_ProjectTable
|
||||
from litellm.models.team import LiteLLM_TeamRebuiltFromToken
|
||||
from litellm.proxy._types import (
|
||||
MODEL_LIST_ROUTES,
|
||||
RBAC_ROLES,
|
||||
|
|
@ -5807,6 +5808,17 @@ async def can_team_access_model(
|
|||
raise
|
||||
|
||||
|
||||
def _group_grants_key(
|
||||
group: LiteLLM_AccessGroupTable, valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None
|
||||
) -> bool:
|
||||
key_team_id: Final = valid_token.team_id or (team_object.team_id if team_object is not None else None)
|
||||
return (
|
||||
bool(key_team_id and key_team_id in (group.assigned_team_ids or []))
|
||||
or bool(valid_token.token and valid_token.token in (group.assigned_key_ids or []))
|
||||
or group.access_group_id in valid_token.access_group_rule_grants
|
||||
)
|
||||
|
||||
|
||||
async def get_authorized_resources_from_key_access_groups(
|
||||
valid_token: UserAPIKeyAuth | None,
|
||||
team_object: LiteLLM_TeamTable | None,
|
||||
|
|
@ -5835,9 +5847,6 @@ async def get_authorized_resources_from_key_access_groups(
|
|||
if _prisma_client is None or _user_api_key_cache is None:
|
||||
return []
|
||||
|
||||
key_team_id: Final = valid_token.team_id or (team_object.team_id if team_object is not None else None)
|
||||
key_token: Final = valid_token.token
|
||||
|
||||
authorized_resources: Final[list[str]] = []
|
||||
for ag_id in key_access_group_ids:
|
||||
try:
|
||||
|
|
@ -5856,9 +5865,7 @@ async def get_authorized_resources_from_key_access_groups(
|
|||
if check_db_only or strict_errors:
|
||||
raise
|
||||
continue
|
||||
team_authorized = bool(key_team_id and key_team_id in (ag.assigned_team_ids or []))
|
||||
key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or []))
|
||||
if team_authorized or key_authorized or ag_id in valid_token.access_group_rule_grants:
|
||||
if _group_grants_key(ag, valid_token, team_object):
|
||||
authorized_resources.extend(getattr(ag, resource_field, []) or [])
|
||||
|
||||
return list(set(authorized_resources))
|
||||
|
|
@ -8095,10 +8102,104 @@ def _legacy_requested_vector_store_ids(
|
|||
)
|
||||
|
||||
|
||||
def _restricts_vector_stores(grants: Sequence[tuple[GrantLayer, LiteLLM_ObjectPermissionTable | None]]) -> bool:
|
||||
def _restricts_vector_stores(grants: Sequence[tuple[GrantLayer, _VectorStorePermissionsRow | None]]) -> bool:
|
||||
return any(grant is not None and bool(grant.vector_stores) for _, grant in grants)
|
||||
|
||||
|
||||
async def access_group_vector_stores(
|
||||
valid_token: UserAPIKeyAuth | None,
|
||||
team_object: LiteLLM_TeamTable | None,
|
||||
unreadable: Literal["skip", "raise"] = "skip",
|
||||
) -> Mapping[GrantLayer, tuple[str, ...]]:
|
||||
"""The stores the key's and its team's access groups list, per layer.
|
||||
A deleted group lists nothing. So does a group that cannot be read, unless `unreadable` is "raise"."""
|
||||
key_groups: Final = await get_access_group_tables(
|
||||
(valid_token.access_group_ids or ()) if valid_token is not None else (), unreadable=unreadable
|
||||
)
|
||||
team_groups: Final = await get_access_group_tables(
|
||||
(team_object.access_group_ids or ()) if team_object is not None else (), unreadable=unreadable
|
||||
)
|
||||
listed_key_groups: Final = (
|
||||
tuple(group for group in key_groups if _group_grants_key(group, valid_token, team_object))
|
||||
if valid_token is not None
|
||||
else ()
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
"key": tuple(chain.from_iterable(group.access_vector_store_ids for group in listed_key_groups)),
|
||||
"team": tuple(chain.from_iterable(group.access_vector_store_ids for group in team_groups)),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _team_read_again(
|
||||
valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None
|
||||
) -> LiteLLM_TeamTable | None:
|
||||
"""Auth's stand-in for a team it could not read has none of the team's groups, so the team is read again."""
|
||||
team_id: Final = valid_token.team_id if valid_token is not None else None
|
||||
if (
|
||||
not isinstance(team_object, LiteLLM_TeamRebuiltFromToken)
|
||||
or valid_token is None
|
||||
or team_id is None
|
||||
or team_id == UI_TEAM_ID
|
||||
):
|
||||
return team_object
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
return await get_team_object(
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=valid_token.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _OwnAndGroupVectorStores:
|
||||
vector_stores: Sequence[str]
|
||||
|
||||
|
||||
def _with_group_stores(
|
||||
grant: LiteLLM_ObjectPermissionTable | None, group_stores: Sequence[str]
|
||||
) -> _VectorStorePermissionsRow | None:
|
||||
if not group_stores:
|
||||
return grant
|
||||
own_stores: Final = (grant.vector_stores if grant is not None else None) or ()
|
||||
return _OwnAndGroupVectorStores([*dict.fromkeys((*own_stores, *group_stores))])
|
||||
|
||||
|
||||
async def _vector_store_grants(
|
||||
valid_token: UserAPIKeyAuth | None,
|
||||
team_object: LiteLLM_TeamTable | None,
|
||||
user_object: LiteLLM_UserTable | None,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
strict_layers: frozenset[GrantLayer] = frozenset(),
|
||||
) -> tuple[tuple[GrantLayer, _VectorStorePermissionsRow | None], ...]:
|
||||
"""Each identity's own grant. The stores its access groups list count as if they were on its own list.
|
||||
A strict team layer can only lose grants without its groups, so it keeps auth's stand-in team."""
|
||||
try:
|
||||
team: Final = team_object if "team" in strict_layers else await _team_read_again(valid_token, team_object)
|
||||
group_stores: Final = await access_group_vector_stores(valid_token, team, unreadable="raise")
|
||||
except Exception as e:
|
||||
# Skipping an unreadable team or group could empty a list and lift the restriction.
|
||||
raise ProxyException(
|
||||
message="Could not read the team or its access groups to check vector store access",
|
||||
type=ProxyErrorTypes.internal_server_error,
|
||||
param="vector_store",
|
||||
code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
) from e
|
||||
return tuple(
|
||||
[
|
||||
(layer, _with_group_stores(grant, group_stores.get(layer, ())))
|
||||
for layer, grant in await _identity_grants(
|
||||
valid_token, team_object, user_object, prisma_client, user_api_key_cache
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _request_drop_policies(
|
||||
request_body: Mapping[str, object], llm_router: Router | None, valid_token: UserAPIKeyAuth | None
|
||||
) -> tuple[DropPolicy, ...]:
|
||||
|
|
@ -8161,7 +8262,7 @@ async def vector_store_access_check(
|
|||
|
||||
Requested stores come from `vector_store_ids`, `tools[].vector_store_ids`, the tools of
|
||||
`additional_tools` input items and the RAG `retrieval_config.vector_store_id`. Grants come from each identity's
|
||||
`object_permission.vector_stores`.
|
||||
`object_permission.vector_stores`, plus the stores its access groups list, which count as if they were on it.
|
||||
|
||||
Both modes read the same request fields, with or without a vector store registry, and answer 400
|
||||
when the top-level `vector_store_ids` is not a list of strings, which legacy mode does only when the
|
||||
|
|
@ -8209,7 +8310,7 @@ async def vector_store_access_check(
|
|||
)
|
||||
if isinstance(requested_ids, _UnreadableVectorStoreIds):
|
||||
if _restricts_vector_stores(
|
||||
await _identity_grants(valid_token, team_object, None, prisma_client, user_api_key_cache)
|
||||
await _vector_store_grants(valid_token, team_object, None, prisma_client, user_api_key_cache)
|
||||
):
|
||||
raise _malformed_vector_store_ids()
|
||||
return True
|
||||
|
|
@ -8223,12 +8324,13 @@ async def vector_store_access_check(
|
|||
# Check if the object (key, team, org) has access to the vector store
|
||||
#########################################################
|
||||
strict_layers: Final = _strict_grant_layers(deny_by_default, valid_token, team_object)
|
||||
grants: Final = await _identity_grants(
|
||||
grants: Final = await _vector_store_grants(
|
||||
valid_token,
|
||||
team_object,
|
||||
user_object if "user" in strict_layers else None,
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
strict_layers,
|
||||
)
|
||||
for layer, grant in grants:
|
||||
if layer in strict_layers:
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ from litellm.litellm_core_utils.dd_tracing import tracer
|
|||
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
|
||||
from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias
|
||||
from litellm.models.end_user import LiteLLM_EndUserTable
|
||||
from litellm.models.team import LiteLLM_TeamRebuiltFromToken
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers
|
||||
from litellm.proxy.auth.access_group_rules import access_group_ids_for_identity
|
||||
|
|
@ -3614,13 +3615,13 @@ async def _safe_fetch(label: str, awaitable):
|
|||
return None
|
||||
|
||||
|
||||
def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCachedObj:
|
||||
def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamRebuiltFromToken:
|
||||
"""Reconstruct a cached team object from the fields already on the
|
||||
UserAPIKeyAuth. Only called when valid_token.team_id is known to be
|
||||
non-None (the caller gates on it)."""
|
||||
assert valid_token.team_id is not None
|
||||
token_team_models: Final = _token_team_models(valid_token)
|
||||
return LiteLLM_TeamTableCachedObj(
|
||||
return LiteLLM_TeamRebuiltFromToken(
|
||||
team_id=valid_token.team_id,
|
||||
max_budget=valid_token.team_max_budget,
|
||||
soft_budget=valid_token.team_soft_budget,
|
||||
|
|
|
|||
|
|
@ -578,6 +578,7 @@ async def create_access_group(
|
|||
"mcp_tool_permissions": _tool_permissions_json(data.mcp_tool_permissions),
|
||||
"access_agent_ids": data.access_agent_ids or [],
|
||||
"access_skill_names": list(data.access_skill_names or ()),
|
||||
"access_vector_store_ids": list(data.access_vector_store_ids or ()),
|
||||
"assigned_team_ids": data.assigned_team_ids or [],
|
||||
"assigned_key_ids": data.assigned_key_ids or [],
|
||||
"created_by": user_api_key_dict.user_id,
|
||||
|
|
@ -676,6 +677,7 @@ async def update_access_group(
|
|||
"access_mcp_server_ids",
|
||||
"access_agent_ids",
|
||||
"access_skill_names",
|
||||
"access_vector_store_ids",
|
||||
):
|
||||
value = list(iter_items(value or ()))
|
||||
update_data[field] = (
|
||||
|
|
|
|||
|
|
@ -1630,6 +1630,7 @@ model LiteLLM_AccessGroupTable {
|
|||
mcp_tool_permissions Json?
|
||||
access_agent_ids String[] @default([])
|
||||
access_skill_names String[] @default([])
|
||||
access_vector_store_ids String[] @default([])
|
||||
|
||||
assigned_team_ids String[] @default([])
|
||||
assigned_key_ids String[] @default([])
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import json
|
||||
import re
|
||||
from collections.abc import Iterable, Mapping
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, Protocol
|
||||
|
||||
|
|
@ -13,6 +14,7 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ObjectPermissionTable,
|
||||
LiteLLM_TeamTable,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
|
|
@ -181,6 +183,8 @@ async def can_user_access_vector_store(
|
|||
4. The caller's team-level object_permission.vector_stores explicitly lists
|
||||
this vector store id.
|
||||
5. The caller's team_id matches the vector store's team_id.
|
||||
6. An access group of the caller's key or team lists this vector store id,
|
||||
which counts exactly like rules 3 and 4.
|
||||
|
||||
A dashboard session credential is evaluated against the same effective
|
||||
contexts as listing (its own grants plus each real team of the user).
|
||||
|
|
@ -193,7 +197,9 @@ async def can_user_access_vector_store(
|
|||
return True
|
||||
|
||||
auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict)
|
||||
return await _is_vector_store_granted_to_any(vector_store, auth_contexts)
|
||||
return await _is_vector_store_granted_to_any(vector_store, auth_contexts) or (
|
||||
vector_store.get("vector_store_id") in await _access_group_vector_store_ids(auth_contexts)
|
||||
)
|
||||
|
||||
|
||||
async def _is_vector_store_granted(
|
||||
|
|
@ -266,16 +272,54 @@ async def _is_vector_store_granted_to_any(
|
|||
return False
|
||||
|
||||
|
||||
async def _team_or_none(auth_context: UserAPIKeyAuth) -> LiteLLM_TeamTable | None:
|
||||
"""A team that cannot be read has no access groups here, as it has no own list."""
|
||||
from litellm.proxy.auth.auth_checks import get_team_object
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
if auth_context.team_id is None:
|
||||
return None
|
||||
try:
|
||||
return await get_team_object(
|
||||
team_id=auth_context.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=auth_context.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # no groups, per the docstring
|
||||
verbose_proxy_logger.debug("Could not read team id=%s for its access groups: %s", auth_context.team_id, e)
|
||||
return None
|
||||
|
||||
|
||||
async def _access_group_vector_store_ids(auth_contexts: tuple[UserAPIKeyAuth, ...]) -> frozenset[str]:
|
||||
"""Groups only add access here, so a group that cannot be read lists nothing."""
|
||||
from litellm.proxy.auth.auth_checks import access_group_vector_stores
|
||||
|
||||
listed: Final = [
|
||||
await access_group_vector_stores(auth_context, await _team_or_none(auth_context))
|
||||
for auth_context in auth_contexts
|
||||
]
|
||||
return frozenset(chain.from_iterable(chain.from_iterable(stores.values()) for stores in listed))
|
||||
|
||||
|
||||
async def filter_listable_vector_stores(
|
||||
vector_stores: Iterable[LiteLLM_ManagedVectorStore],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> tuple[LiteLLM_ManagedVectorStore, ...]:
|
||||
"""Non-admins only see stores their key, one of their teams' object_permission, or team ownership grants."""
|
||||
"""Non-admins only see stores their key, a team's object_permission, team ownership or access groups grant."""
|
||||
if _is_proxy_admin(user_api_key_dict):
|
||||
return tuple(vector_stores)
|
||||
|
||||
auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict)
|
||||
return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)])
|
||||
group_store_ids: Final = await _access_group_vector_store_ids(auth_contexts)
|
||||
return tuple(
|
||||
[
|
||||
vs
|
||||
for vs in vector_stores
|
||||
if vs.get("vector_store_id") in group_store_ids or await _is_vector_store_granted_to_any(vs, auth_contexts)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def get_litellm_managed_vector_store(
|
||||
|
|
|
|||
|
|
@ -1630,6 +1630,7 @@ model LiteLLM_AccessGroupTable {
|
|||
mcp_tool_permissions Json?
|
||||
access_agent_ids String[] @default([])
|
||||
access_skill_names String[] @default([])
|
||||
access_vector_store_ids String[] @default([])
|
||||
|
||||
assigned_team_ids String[] @default([])
|
||||
assigned_key_ids String[] @default([])
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ class AccessGroupCreateRequest(LiteLLMBaseModel):
|
|||
mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
|
||||
access_agent_ids: list[str] | None = None
|
||||
access_skill_names: tuple[str, ...] | None = None
|
||||
access_vector_store_ids: tuple[str, ...] | None = None
|
||||
assigned_team_ids: list[str] | None = None
|
||||
assigned_key_ids: list[str] | None = None
|
||||
|
||||
|
|
@ -24,6 +25,7 @@ class AccessGroupUpdateRequest(LiteLLMBaseModel):
|
|||
mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
|
||||
access_agent_ids: list[str] | None = None
|
||||
access_skill_names: tuple[str, ...] | None = None
|
||||
access_vector_store_ids: tuple[str, ...] | None = None
|
||||
assigned_team_ids: list[str] | None = None
|
||||
assigned_key_ids: list[str] | None = None
|
||||
|
||||
|
|
@ -44,6 +46,7 @@ class AccessGroupResponse(LiteLLMBaseModel):
|
|||
mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
|
||||
access_agent_ids: list[str]
|
||||
access_skill_names: tuple[str, ...] = ()
|
||||
access_vector_store_ids: tuple[str, ...] = ()
|
||||
assigned_team_ids: list[str]
|
||||
assigned_key_ids: list[str]
|
||||
access_mcp_servers: tuple[AccessGroupResource, ...]
|
||||
|
|
|
|||
9
litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
9
litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -26299,6 +26299,8 @@ export interface components {
|
|||
access_model_names?: string[] | null;
|
||||
/** Access Skill Names */
|
||||
access_skill_names?: string[] | null;
|
||||
/** Access Vector Store Ids */
|
||||
access_vector_store_ids?: string[] | null;
|
||||
/** Assigned Key Ids */
|
||||
assigned_key_ids?: string[] | null;
|
||||
/** Assigned Team Ids */
|
||||
|
|
@ -26353,6 +26355,11 @@ export interface components {
|
|||
* @default []
|
||||
*/
|
||||
access_skill_names: string[];
|
||||
/**
|
||||
* Access Vector Store Ids
|
||||
* @default []
|
||||
*/
|
||||
access_vector_store_ids: string[];
|
||||
/** Assigned Key Ids */
|
||||
assigned_key_ids: string[];
|
||||
/** Assigned Keys */
|
||||
|
|
@ -26405,6 +26412,8 @@ export interface components {
|
|||
access_model_names?: string[] | null;
|
||||
/** Access Skill Names */
|
||||
access_skill_names?: string[] | null;
|
||||
/** Access Vector Store Ids */
|
||||
access_vector_store_ids?: string[] | null;
|
||||
/** Assigned Key Ids */
|
||||
assigned_key_ids?: string[] | null;
|
||||
/** Assigned Team Ids */
|
||||
|
|
|
|||
359
tests/unit/proxy/auth/test_vector_store_access_groups.py
Normal file
359
tests/unit/proxy/auth/test_vector_store_access_groups.py
Normal file
|
|
@ -0,0 +1,359 @@
|
|||
"""A team or key in an access group that grants stores can use those stores and no others."""
|
||||
|
||||
import base64
|
||||
from collections.abc import Sequence
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
ProxyException,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import vector_store_access_check
|
||||
from litellm.proxy.auth.user_api_key_auth import _team_obj_from_token
|
||||
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
|
||||
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
|
||||
from tests.unit.proxy.vector_store_access_group_state import (
|
||||
TEAM_ID,
|
||||
TOKEN,
|
||||
ProxyState,
|
||||
make_key,
|
||||
make_team,
|
||||
own_permission,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def proxy(monkeypatch: pytest.MonkeyPatch) -> ProxyState:
|
||||
return ProxyState(monkeypatch)
|
||||
|
||||
|
||||
def _virtual_key(groups: Sequence[str]) -> UserAPIKeyAuth:
|
||||
key: Final = UserAPIKeyAuth(token=TOKEN, api_key=TOKEN, team_id=TEAM_ID, access_group_ids=list(groups))
|
||||
key.via_virtual_key = True
|
||||
return key
|
||||
|
||||
|
||||
async def _request_is_allowed(
|
||||
key: UserAPIKeyAuth, team: LiteLLM_TeamTableCachedObj | None, store: str, *, deny_by_default: bool = False
|
||||
) -> bool:
|
||||
try:
|
||||
await vector_store_access_check(
|
||||
request_body={"vector_store_ids": [store]},
|
||||
team_object=team,
|
||||
valid_token=key,
|
||||
deny_by_default=deny_by_default,
|
||||
)
|
||||
except ProxyException as refused:
|
||||
if refused.code != "401":
|
||||
raise
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("key_stores", "key_groups", "team_stores", "team_groups", "store", "allowed"),
|
||||
[
|
||||
pytest.param(None, ("g1",), None, (), "store-a", True, id="key-group-grants"),
|
||||
pytest.param(None, ("g1",), None, (), "store-b", False, id="key-group-restricts"),
|
||||
pytest.param((), ("g1",), None, (), "store-b", False, id="key-group-restricts-with-empty-own-list"),
|
||||
pytest.param(None, (), None, ("g1",), "store-a", True, id="team-group-grants"),
|
||||
pytest.param(None, (), None, ("g1",), "store-b", False, id="team-group-restricts"),
|
||||
pytest.param(None, (), (), ("g1",), "store-b", False, id="team-group-restricts-with-empty-own-list"),
|
||||
pytest.param(("store-x",), ("g1",), None, (), "store-a", True, id="key-group-widens-own-list"),
|
||||
pytest.param(("store-x",), ("g1",), None, (), "store-x", True, id="key-own-list-still-counts"),
|
||||
pytest.param(("store-x",), ("g1",), None, (), "store-b", False, id="key-neither-list-denies"),
|
||||
pytest.param(None, (), ("store-x",), ("g1",), "store-a", True, id="team-group-widens-own-list"),
|
||||
pytest.param(None, (), ("store-x",), ("g1",), "store-x", True, id="team-own-list-still-counts"),
|
||||
pytest.param(("store-x",), (), None, ("g1",), "store-a", False, id="key-own-list-narrows-team-group"),
|
||||
pytest.param(None, ("g1",), ("store-x",), (), "store-a", False, id="team-own-list-narrows-key-group"),
|
||||
pytest.param(None, ("g1", "g2"), None, (), "store-c", True, id="stores-from-a-second-group"),
|
||||
pytest.param(None, ("models-only",), None, (), "store-b", True, id="models-only-group-changes-nothing"),
|
||||
pytest.param(None, (), None, ("models-only",), "store-b", True, id="team-models-only-group-changes-nothing"),
|
||||
pytest.param(("store-x",), ("models-only",), None, (), "store-b", False, id="models-only-keeps-own-list"),
|
||||
pytest.param(None, (), None, (), "store-b", True, id="no-grants-anywhere-changes-nothing"),
|
||||
],
|
||||
)
|
||||
async def test_request_check_combines_own_list_and_group_stores(
|
||||
proxy: ProxyState,
|
||||
key_stores: Sequence[str] | None,
|
||||
key_groups: Sequence[str],
|
||||
team_stores: Sequence[str] | None,
|
||||
team_groups: Sequence[str],
|
||||
store: str,
|
||||
allowed: bool,
|
||||
) -> None:
|
||||
await proxy.add_group("g1", ("store-a",))
|
||||
await proxy.add_group("g2", ("store-c",))
|
||||
await proxy.add_group("models-only", models=("gpt-4",))
|
||||
|
||||
assert (
|
||||
await _request_is_allowed(make_key(key_stores, key_groups, TEAM_ID), make_team(team_stores, team_groups), store)
|
||||
is allowed
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_group_counts_when_it_lists_only_the_keys_team(proxy: ProxyState) -> None:
|
||||
await proxy.add_group("g1", ("store-a",), key_ids=())
|
||||
|
||||
assert await _request_is_allowed(make_key(groups=("g1",), team_id=TEAM_ID), None, "store-a") is True
|
||||
assert await _request_is_allowed(make_key(groups=("g1",), team_id=TEAM_ID), None, "store-b") is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_granted_by_an_identity_rule_counts_for_the_key(proxy: ProxyState) -> None:
|
||||
await proxy.add_group("g1", ("store-a",), team_ids=(), key_ids=())
|
||||
key: Final = make_key(("store-x",))
|
||||
key.apply_access_group_rule_grants(("g1",))
|
||||
|
||||
assert await _request_is_allowed(key, None, "store-a") is True
|
||||
assert await _request_is_allowed(key, None, "store-b") is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_cannot_claim_a_group_that_does_not_list_it(proxy: ProxyState) -> None:
|
||||
await proxy.add_group("g1", ("store-a",), team_ids=(), key_ids=())
|
||||
|
||||
assert await _request_is_allowed(make_key(("store-x",), ("g1",)), None, "store-a") is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("store", "allowed"),
|
||||
[pytest.param("store-a", True, id="granted"), pytest.param("store-b", False, id="not-granted")],
|
||||
)
|
||||
async def test_deny_by_default_counts_group_grants(proxy: ProxyState, store: str, allowed: bool) -> None:
|
||||
await proxy.add_group("g1", ("store-a",))
|
||||
key: Final = _virtual_key(("g1",))
|
||||
|
||||
assert await _request_is_allowed(key, make_team(groups=("g1",)), store, deny_by_default=True) is allowed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deny_by_default_requires_the_key_too_not_only_its_team(proxy: ProxyState) -> None:
|
||||
await proxy.add_group("g1", ("store-a",))
|
||||
key: Final = _virtual_key(())
|
||||
|
||||
assert await _request_is_allowed(key, make_team(groups=("g1",)), "store-a", deny_by_default=True) is False
|
||||
|
||||
|
||||
_MANAGED_A: Final = base64.urlsafe_b64encode(
|
||||
b"litellm_proxy:vector_store;unified_id,abc;target_model_names,gpt-4;resource_id,kb-shared"
|
||||
).decode()
|
||||
_PLAIN_B: Final = LiteLLM_ManagedVectorStore(vector_store_id="kb-shared", custom_llm_provider="bedrock")
|
||||
|
||||
|
||||
async def _all_allowed(key: UserAPIKeyAuth, stores: Sequence[str]) -> bool:
|
||||
try:
|
||||
await vector_store_access_check(
|
||||
request_body={"vector_store_ids": list(stores)}, team_object=None, valid_token=key
|
||||
)
|
||||
except ProxyException as refused:
|
||||
if refused.code != "401":
|
||||
raise
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_own_list_still_matches_by_provider_id(proxy: ProxyState, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry(vector_stores=[_PLAIN_B]))
|
||||
|
||||
assert await _all_allowed(make_key((_MANAGED_A,)), ("kb-shared",)) is True
|
||||
|
||||
|
||||
async def _answer(
|
||||
layer: Literal["key", "team"],
|
||||
own: Sequence[str] | None,
|
||||
groups: Sequence[str],
|
||||
requested: Sequence[str],
|
||||
deny_by_default: bool,
|
||||
) -> tuple[object, ...] | None:
|
||||
"""The refusal a virtual key gets when `layer` holds the `own` list and the `groups`, or None when allowed.
|
||||
Under deny-by-default the other layer lists every requested store, so `layer` is the one that decides."""
|
||||
other: Final = tuple(requested) if deny_by_default else None
|
||||
key_stores: Final = own if layer == "key" else other
|
||||
key: Final = _virtual_key(groups if layer == "key" else ()).model_copy(
|
||||
update={"object_permission_id": "own", "object_permission": own_permission(key_stores)}
|
||||
if key_stores is not None
|
||||
else {}
|
||||
)
|
||||
team: Final = make_team(own if layer == "team" else other, groups if layer == "team" else ())
|
||||
try:
|
||||
await vector_store_access_check(
|
||||
request_body={"vector_store_ids": list(requested)},
|
||||
team_object=team,
|
||||
valid_token=key,
|
||||
deny_by_default=deny_by_default,
|
||||
)
|
||||
except ProxyException as refused:
|
||||
return (refused.code, refused.type, refused.message)
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("deny_by_default", [False, True], ids=["legacy", "deny-by-default"])
|
||||
@pytest.mark.parametrize("layer", ["key", "team"])
|
||||
@pytest.mark.parametrize("own", [None, ("store-x",)], ids=["no-own-list", "own-list"])
|
||||
@pytest.mark.parametrize(
|
||||
"group", [("store-a",), ("kb-shared",), (_MANAGED_A,)], ids=["plain", "provider-id", "managed"]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"requested",
|
||||
[("store-a",), ("store-x",), ("store-b",), ("kb-shared",), (_MANAGED_A,), ("store-a", "store-b")],
|
||||
ids=["store-a", "store-x", "store-b", "provider-id", "managed", "two-stores"],
|
||||
)
|
||||
async def test_a_group_store_is_checked_exactly_like_the_same_store_on_the_own_list(
|
||||
proxy: ProxyState,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
deny_by_default: bool,
|
||||
layer: Literal["key", "team"],
|
||||
own: Sequence[str] | None,
|
||||
group: Sequence[str],
|
||||
requested: Sequence[str],
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry(vector_stores=[_PLAIN_B]))
|
||||
await proxy.add_group("g1", group)
|
||||
|
||||
through_the_group: Final = await _answer(layer, own, ("g1",), requested, deny_by_default)
|
||||
on_the_own_list: Final = await _answer(layer, (*(own or ()), *group), (), requested, deny_by_default)
|
||||
|
||||
assert through_the_group == on_the_own_list
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_refusal_lists_the_allowed_stores_as_main_does(proxy: ProxyState) -> None:
|
||||
with pytest.raises(ProxyException) as denied:
|
||||
await vector_store_access_check(
|
||||
request_body={"vector_store_ids": ["store-b"]}, team_object=None, valid_token=make_key(("store-x",))
|
||||
)
|
||||
|
||||
assert "Only allowed to access ['store-x']" in denied.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keyless_team_member_gets_the_teams_group_stores(proxy: ProxyState) -> None:
|
||||
await proxy.add_group("g1", ("store-a",))
|
||||
member: Final = UserAPIKeyAuth(team_id=TEAM_ID)
|
||||
|
||||
assert await _request_is_allowed(member, make_team(groups=("g1",)), "store-a") is True
|
||||
assert await _request_is_allowed(member, make_team(groups=("g1",)), "store-b") is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deny_by_default_still_denies_a_models_only_group(proxy: ProxyState) -> None:
|
||||
await proxy.add_group("models-only", models=("gpt-4",))
|
||||
key: Final = _virtual_key(("models-only",))
|
||||
|
||||
assert await _request_is_allowed(key, make_team(groups=("models-only",)), "store-a", deny_by_default=True) is False
|
||||
|
||||
|
||||
_UNREADABLE: Final = "Could not read the team or its access groups to check vector store access"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("own", [None, ("store-x",)], ids=["no-own-list", "own-list"])
|
||||
@pytest.mark.parametrize("layer", ["key", "team"])
|
||||
async def test_a_group_that_cannot_be_read_blocks_the_request_instead_of_lifting_the_restriction(
|
||||
monkeypatch: pytest.MonkeyPatch, layer: Literal["key", "team"], own: Sequence[str] | None
|
||||
) -> None:
|
||||
ProxyState(monkeypatch, group_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
|
||||
|
||||
refusal: Final = await _answer(layer, own, ("g1",), ("store-b",), deny_by_default=False)
|
||||
|
||||
assert refusal is not None
|
||||
assert refusal[0] == "503"
|
||||
assert refusal[2] == _UNREADABLE
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auths_stand_in_for_a_team_it_could_not_read_is_read_again_so_the_teams_groups_still_apply(
|
||||
proxy: ProxyState,
|
||||
) -> None:
|
||||
await proxy.add_group("g1", ("store-a",))
|
||||
await proxy.add_team(make_team(groups=("g1",)))
|
||||
key: Final = make_key(team_id=TEAM_ID)
|
||||
|
||||
assert await _request_is_allowed(key, _team_obj_from_token(key), "store-a") is True
|
||||
assert await _request_is_allowed(key, _team_obj_from_token(key), "store-b") is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_stand_in_team_that_still_cannot_be_read_blocks_the_request(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
ProxyState(monkeypatch, team_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
|
||||
key: Final = make_key(team_id=TEAM_ID)
|
||||
|
||||
with pytest.raises(ProxyException) as blocked:
|
||||
await vector_store_access_check(
|
||||
request_body={"vector_store_ids": ["store-b"]}, team_object=_team_obj_from_token(key), valid_token=key
|
||||
)
|
||||
|
||||
assert (blocked.value.code, blocked.value.message) == ("503", _UNREADABLE)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_strict_team_layer_keeps_the_stand_in_team_and_is_refused_as_on_main(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
ProxyState(monkeypatch, team_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
|
||||
key: Final = _virtual_key(())
|
||||
|
||||
assert await _request_is_allowed(key, _team_obj_from_token(key), "store-b", deny_by_default=True) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleted_group_grants_nothing_and_does_not_fail_the_request(proxy: ProxyState) -> None:
|
||||
assert await _request_is_allowed(make_key(team_id=TEAM_ID), make_team(groups=("deleted",)), "store-b") is True
|
||||
assert (
|
||||
await _request_is_allowed(make_key(("store-x",), team_id=TEAM_ID), make_team(groups=("deleted",)), "store-b")
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_cached_before_the_upgrade_keeps_the_callers_own_list(proxy: ProxyState) -> None:
|
||||
await proxy.cache.async_set_cache(
|
||||
key="access_group_id:old",
|
||||
value={"access_group_id": "old", "access_group_name": "old", "assigned_team_ids": [TEAM_ID]},
|
||||
)
|
||||
|
||||
assert (
|
||||
await _request_is_allowed(make_key(("store-x",), team_id=TEAM_ID), make_team(groups=("old",)), "store-b")
|
||||
is False
|
||||
)
|
||||
assert (
|
||||
await _request_is_allowed(make_key(("store-x",), team_id=TEAM_ID), make_team(groups=("old",)), "store-x")
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_malformed_store_ids_are_rejected_when_only_a_group_restricts(proxy: ProxyState) -> None:
|
||||
await proxy.add_group("g1", ("store-a",))
|
||||
|
||||
with pytest.raises(ProxyException, match="vector_store_ids must be a list of strings"):
|
||||
await vector_store_access_check(
|
||||
request_body={"model": "gpt-4o-mini", "vector_store_ids": 123},
|
||||
team_object=make_team(groups=("g1",)),
|
||||
valid_token=make_key(team_id=TEAM_ID),
|
||||
llm_router=Router(model_list=[{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}]),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_naming_no_store_never_reads_groups(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
proxy: Final = ProxyState(monkeypatch, group_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
|
||||
|
||||
result: Final = await vector_store_access_check(
|
||||
request_body={"model": "gpt-4"}, team_object=make_team(groups=("g1",)), valid_token=make_key(team_id=TEAM_ID)
|
||||
)
|
||||
|
||||
assert result is True
|
||||
assert proxy.group_reads.await_count == 0
|
||||
|
|
@ -3,8 +3,10 @@ Tests for access group management endpoints.
|
|||
"""
|
||||
|
||||
import types
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime
|
||||
from typing import Final, TypeAlias
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -22,6 +24,9 @@ from litellm.proxy.proxy_server import app
|
|||
from litellm.types.agents import AgentResponse
|
||||
|
||||
|
||||
_ClientAndMocks: TypeAlias = tuple[TestClient, MagicMock, MagicMock, MagicMock, MagicMock]
|
||||
|
||||
|
||||
def _make_access_group_record(
|
||||
access_group_id: str = "ag-123",
|
||||
access_group_name: str = "test-group",
|
||||
|
|
@ -30,12 +35,13 @@ def _make_access_group_record(
|
|||
access_mcp_server_ids: list | None = None,
|
||||
access_agent_ids: list | None = None,
|
||||
access_skill_names: list | None = None,
|
||||
access_vector_store_ids: Sequence[str] = (),
|
||||
assigned_team_ids: list | None = None,
|
||||
assigned_key_ids: list | None = None,
|
||||
created_by: str | None = "admin-user",
|
||||
updated_by: str | None = "admin-user",
|
||||
created_at: datetime | None = None,
|
||||
):
|
||||
) -> MagicMock:
|
||||
created_at_val = created_at or datetime.now()
|
||||
updated_at_val = datetime.now()
|
||||
data = {
|
||||
|
|
@ -46,6 +52,7 @@ def _make_access_group_record(
|
|||
"access_mcp_server_ids": access_mcp_server_ids or [],
|
||||
"access_agent_ids": access_agent_ids or [],
|
||||
"access_skill_names": access_skill_names or [],
|
||||
"access_vector_store_ids": list(access_vector_store_ids),
|
||||
"assigned_team_ids": assigned_team_ids or [],
|
||||
"assigned_key_ids": assigned_key_ids or [],
|
||||
"created_at": created_at_val,
|
||||
|
|
@ -78,12 +85,12 @@ def _make_key_record(token: str, key_alias: str | None = None):
|
|||
|
||||
|
||||
@pytest.fixture
|
||||
def client_and_mocks(monkeypatch):
|
||||
def client_and_mocks(monkeypatch: pytest.MonkeyPatch) -> Iterator[_ClientAndMocks]:
|
||||
"""Setup mock prisma and admin auth for access group endpoints."""
|
||||
mock_access_group_table = MagicMock()
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
def _create_side_effect(*, data):
|
||||
def _create_side_effect(*, data: Mapping[str, object]) -> MagicMock:
|
||||
return _make_access_group_record(
|
||||
access_group_id="ag-new",
|
||||
access_group_name=data.get("access_group_name", "new"),
|
||||
|
|
@ -92,6 +99,7 @@ def client_and_mocks(monkeypatch):
|
|||
access_mcp_server_ids=data.get("access_mcp_server_ids", []),
|
||||
access_agent_ids=data.get("access_agent_ids", []),
|
||||
access_skill_names=data.get("access_skill_names", []),
|
||||
access_vector_store_ids=data.get("access_vector_store_ids", []),
|
||||
assigned_team_ids=data.get("assigned_team_ids", []),
|
||||
assigned_key_ids=data.get("assigned_key_ids", []),
|
||||
created_by=data.get("created_by"),
|
||||
|
|
@ -110,6 +118,7 @@ def client_and_mocks(monkeypatch):
|
|||
access_mcp_server_ids=data.get("access_mcp_server_ids", []),
|
||||
access_agent_ids=data.get("access_agent_ids", []),
|
||||
access_skill_names=data.get("access_skill_names", []),
|
||||
access_vector_store_ids=data.get("access_vector_store_ids", []),
|
||||
assigned_team_ids=data.get("assigned_team_ids", []),
|
||||
assigned_key_ids=data.get("assigned_key_ids", []),
|
||||
updated_by=data.get("updated_by"),
|
||||
|
|
@ -1714,6 +1723,41 @@ def test_update_access_group_writes_skill_names(client_and_mocks, update_payload
|
|||
assert mock_table.update.await_args.kwargs["data"]["access_skill_names"] == expected_skills
|
||||
|
||||
|
||||
def test_create_access_group_persists_and_returns_vector_store_ids(client_and_mocks: _ClientAndMocks) -> None:
|
||||
client, _, mock_table, *_ = client_and_mocks
|
||||
|
||||
resp: Final = client.post(
|
||||
"/v1/access_group", json={"access_group_name": "stores", "access_vector_store_ids": ["store-a"]}
|
||||
)
|
||||
|
||||
assert resp.status_code == 201
|
||||
assert resp.json()["access_vector_store_ids"] == ["store-a"]
|
||||
assert mock_table.create.await_args.kwargs["data"]["access_vector_store_ids"] == ["store-a"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("update_payload", "expected_stores"),
|
||||
[
|
||||
({"access_vector_store_ids": ["store-a"]}, ["store-a"]),
|
||||
({"access_vector_store_ids": []}, []),
|
||||
({"access_vector_store_ids": None}, []),
|
||||
],
|
||||
)
|
||||
def test_update_access_group_writes_vector_store_ids(
|
||||
client_and_mocks: _ClientAndMocks, update_payload: Mapping[str, object], expected_stores: Sequence[str]
|
||||
) -> None:
|
||||
client, _, mock_table, *_ = client_and_mocks
|
||||
mock_table.find_unique = AsyncMock(
|
||||
return_value=_make_access_group_record(access_group_id="ag-stores", access_vector_store_ids=["old-store"])
|
||||
)
|
||||
|
||||
resp: Final = client.put("/v1/access_group/ag-stores", json=update_payload)
|
||||
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["access_vector_store_ids"] == expected_stores
|
||||
assert mock_table.update.await_args.kwargs["data"]["access_vector_store_ids"] == expected_stores
|
||||
|
||||
|
||||
def test_update_access_group_without_skill_field_leaves_skills_untouched(client_and_mocks):
|
||||
client, _, mock_table, *_ = client_and_mocks
|
||||
mock_table.find_unique = AsyncMock(
|
||||
|
|
|
|||
104
tests/unit/proxy/vector_store_access_group_state.py
Normal file
104
tests/unit/proxy/vector_store_access_group_state.py
Normal file
|
|
@ -0,0 +1,104 @@
|
|||
"""Proxy state and identities shared by the vector store access group tests."""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.models.access_group import LiteLLM_AccessGroupTable
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ObjectPermissionTable,
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import cache_access_object
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
|
||||
|
||||
TOKEN: Final = "hashed-key"
|
||||
TEAM_ID: Final = "team-t"
|
||||
OTHER_TEAM_ID: Final = "owner-team"
|
||||
|
||||
|
||||
class ProxyState:
|
||||
"""The proxy globals the checks read: a real cache in front of a database that holds no groups."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
group_reads: AsyncMock | None = None,
|
||||
team_reads: AsyncMock | None = None,
|
||||
store_reads: AsyncMock | None = None,
|
||||
cache: UserApiKeyCache | MagicMock | None = None,
|
||||
) -> None:
|
||||
self.cache: Final = cache if cache is not None else UserApiKeyCache()
|
||||
self.group_reads: Final = group_reads if group_reads is not None else AsyncMock(return_value=None)
|
||||
self.prisma: Final = MagicMock(
|
||||
db=MagicMock(
|
||||
litellm_accessgrouptable=MagicMock(find_unique=self.group_reads),
|
||||
litellm_teamtable=MagicMock(find_unique=team_reads or AsyncMock(return_value=None)),
|
||||
litellm_managedvectorstorestable=MagicMock(find_unique=store_reads or AsyncMock(return_value=None)),
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", self.prisma)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", self.cache)
|
||||
|
||||
async def add_group(
|
||||
self,
|
||||
group_id: str,
|
||||
stores: Sequence[str] = (),
|
||||
*,
|
||||
models: Sequence[str] = (),
|
||||
team_ids: Sequence[str] = (TEAM_ID,),
|
||||
key_ids: Sequence[str] = (TOKEN,),
|
||||
) -> None:
|
||||
await cache_access_object(
|
||||
access_group_id=group_id,
|
||||
access_group_table=LiteLLM_AccessGroupTable(
|
||||
access_group_id=group_id,
|
||||
access_group_name=group_id,
|
||||
access_model_names=list(models),
|
||||
access_vector_store_ids=tuple(stores),
|
||||
assigned_team_ids=list(team_ids),
|
||||
assigned_key_ids=list(key_ids),
|
||||
),
|
||||
user_api_key_cache=self.cache,
|
||||
)
|
||||
|
||||
async def add_team(self, team: LiteLLM_TeamTableCachedObj) -> None:
|
||||
await self.cache.async_set_cache(key=f"team_id:{team.team_id}", value=team)
|
||||
|
||||
|
||||
def own_permission(stores: Sequence[str] | None) -> LiteLLM_ObjectPermissionTable | None:
|
||||
if stores is None:
|
||||
return None
|
||||
return LiteLLM_ObjectPermissionTable(object_permission_id="own", vector_stores=list(stores))
|
||||
|
||||
|
||||
def make_key(
|
||||
own_stores: Sequence[str] | None = None, groups: Sequence[str] = (), team_id: str | None = None
|
||||
) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
token=TOKEN,
|
||||
team_id=team_id,
|
||||
access_group_ids=list(groups),
|
||||
object_permission_id="own" if own_stores is not None else None,
|
||||
object_permission=own_permission(own_stores),
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
|
||||
def make_team(own_stores: Sequence[str] | None = None, groups: Sequence[str] = ()) -> LiteLLM_TeamTableCachedObj:
|
||||
return LiteLLM_TeamTableCachedObj(
|
||||
team_id=TEAM_ID,
|
||||
access_group_ids=list(groups),
|
||||
object_permission_id="own" if own_stores is not None else None,
|
||||
object_permission=own_permission(own_stores),
|
||||
)
|
||||
|
||||
|
||||
def team_owned_store(store_id: str) -> LiteLLM_ManagedVectorStore:
|
||||
return LiteLLM_ManagedVectorStore(vector_store_id=store_id, custom_llm_provider="openai", team_id=OTHER_TEAM_ID)
|
||||
|
|
@ -0,0 +1,310 @@
|
|||
"""A store an access group lists is answered exactly like the same store on the key's or team's own list."""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import UI_TEAM_ID, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import app
|
||||
from litellm.proxy.utils import hash_token
|
||||
from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store, filter_listable_vector_stores
|
||||
from litellm.types.vector_stores import LiteLLM_ManagedVectorStore
|
||||
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
|
||||
from tests._master_key import MASTER_KEY
|
||||
from tests.unit.proxy.vector_store_access_group_state import (
|
||||
TEAM_ID,
|
||||
TOKEN,
|
||||
ProxyState,
|
||||
make_key,
|
||||
make_team,
|
||||
own_permission,
|
||||
team_owned_store,
|
||||
)
|
||||
|
||||
Layer = Literal["key", "team"]
|
||||
_KEY: Final = "sk-group-store"
|
||||
_STORES: Final = (
|
||||
team_owned_store("store-a"),
|
||||
team_owned_store("store-b"),
|
||||
LiteLLM_ManagedVectorStore(vector_store_id="store-free", custom_llm_provider="openai"),
|
||||
)
|
||||
|
||||
|
||||
async def _caller(proxy: ProxyState, layer: Layer, how: Literal["own-list", "group"]) -> UserAPIKeyAuth:
|
||||
"""A team key granted store-a on `layer`, either on that layer's own list or through its access group."""
|
||||
if how == "own-list":
|
||||
await proxy.add_team(make_team(("store-a",) if layer == "team" else None))
|
||||
key: Final = make_key(("store-a",) if layer == "key" else None, team_id=TEAM_ID)
|
||||
return (
|
||||
key.model_copy(update={"team_object_permission": own_permission(("store-a",))}) if layer == "team" else key
|
||||
)
|
||||
await proxy.add_group("g1", ("store-a",))
|
||||
await proxy.add_team(make_team(groups=("g1",) if layer == "team" else ()))
|
||||
return make_key(groups=("g1",) if layer == "key" else (), team_id=TEAM_ID)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("layer", ["key", "team"])
|
||||
@pytest.mark.parametrize("how", ["own-list", "group"])
|
||||
async def test_the_store_gate_and_the_list_answer_a_group_store_like_an_own_list_store(
|
||||
monkeypatch: pytest.MonkeyPatch, layer: Layer, how: Literal["own-list", "group"]
|
||||
) -> None:
|
||||
caller: Final = await _caller(ProxyState(monkeypatch), layer, how)
|
||||
|
||||
assert [await can_user_access_vector_store(store, caller) for store in _STORES] == [True, False, True]
|
||||
assert await filter_listable_vector_stores(_STORES, caller) == (_STORES[0],)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("team_is_readable", "allowed"),
|
||||
[pytest.param(True, [True, True, True], id="readable"), pytest.param(False, [True, False, True], id="unreadable")],
|
||||
)
|
||||
async def test_a_team_that_cannot_be_read_loses_its_group_stores_and_the_keys_still_count(
|
||||
monkeypatch: pytest.MonkeyPatch, team_is_readable: bool, allowed: Sequence[bool]
|
||||
) -> None:
|
||||
"""The key's group lists store-a and names only the key's team. The team's group lists store-b."""
|
||||
proxy: Final = ProxyState(monkeypatch, team_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
|
||||
await proxy.add_group("key-group", ("store-a",), key_ids=())
|
||||
await proxy.add_group("team-group", ("store-b",))
|
||||
if team_is_readable:
|
||||
await proxy.add_team(make_team(groups=("team-group",)))
|
||||
caller: Final = make_key(groups=("key-group",), team_id=TEAM_ID)
|
||||
|
||||
assert [await can_user_access_vector_store(store, caller) for store in _STORES] == allowed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_caller_without_groups_keeps_its_answers_when_its_team_cannot_be_read(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
ProxyState(monkeypatch, team_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
|
||||
caller: Final = make_key(("store-a",), team_id=TEAM_ID)
|
||||
|
||||
assert [await can_user_access_vector_store(store, caller) for store in _STORES] == [True, False, True]
|
||||
assert await filter_listable_vector_stores(_STORES, caller) == (_STORES[0],)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_group_that_cannot_be_read_grants_nothing_and_own_list_stores_still_pass(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
proxy: Final = ProxyState(monkeypatch, group_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
|
||||
await proxy.add_team(make_team(groups=("g1",)))
|
||||
caller: Final = make_key(("store-b",), ("g1",), TEAM_ID)
|
||||
|
||||
assert [await can_user_access_vector_store(store, caller) for store in _STORES] == [False, True, True]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_dashboard_session_gets_the_group_stores_of_each_of_the_users_teams(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
proxy: Final = ProxyState(monkeypatch)
|
||||
await proxy.add_group("g1", ("store-a",), team_ids=("team-1",))
|
||||
await proxy.add_group("g2", ("store-c",), team_ids=("team-2",))
|
||||
for team_id, group in (("team-1", "g1"), ("team-2", "g2")):
|
||||
await proxy.cache.async_set_cache(
|
||||
key=f"team_id:{team_id}", value=LiteLLM_TeamTableCachedObj(team_id=team_id, access_group_ids=[group])
|
||||
)
|
||||
await proxy.cache.async_set_cache(
|
||||
key="session-user", value=LiteLLM_UserTable(user_id="session-user", teams=["team-1", "team-2"])
|
||||
)
|
||||
session: Final = UserAPIKeyAuth(token=TOKEN, user_id="session-user", team_id=UI_TEAM_ID)
|
||||
stores: Final = tuple(team_owned_store(store_id) for store_id in ("store-a", "store-b", "store-c"))
|
||||
|
||||
assert await filter_listable_vector_stores(stores, session) == (stores[0], stores[2])
|
||||
assert [await can_user_access_vector_store(store, session) for store in stores] == [True, False, True]
|
||||
|
||||
|
||||
READ_ROUTES: Final = (
|
||||
pytest.param("POST", "/v1/vector_stores/{}/search", {"query": "hi"}, id="search"),
|
||||
pytest.param("GET", "/v1/vector_stores/{}", None, id="retrieve"),
|
||||
pytest.param("GET", "/v1/vector_stores/{}/files", None, id="file-list"),
|
||||
pytest.param("GET", "/v1/vector_stores/{}/files/file-1", None, id="file-retrieve"),
|
||||
pytest.param("GET", "/v1/vector_stores/{}/files/file-1/content", None, id="file-content"),
|
||||
pytest.param("POST", "/vector_store/info", {"vector_store_id": "{}"}, id="info"),
|
||||
pytest.param(
|
||||
"POST",
|
||||
"/v1/rag/query",
|
||||
{"model": "m", "messages": [{"role": "user", "content": "hi"}], "retrieval_config": {"vector_store_id": "{}"}},
|
||||
id="rag-query",
|
||||
),
|
||||
pytest.param(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": "m", "input": "hi", "tools": [{"type": "file_search", "vector_store_ids": ["{}"]}]},
|
||||
id="responses",
|
||||
),
|
||||
)
|
||||
|
||||
MANAGE_ROUTES: Final = (
|
||||
pytest.param(
|
||||
"POST", "/vector_store/update", {"vector_store_id": "{}", "vector_store_name": "renamed"}, id="update"
|
||||
),
|
||||
pytest.param("POST", "/vector_store/delete", {"vector_store_id": "{}"}, id="delete"),
|
||||
pytest.param("POST", "/v1/vector_stores/{}", {"name": "renamed"}, id="provider-update"),
|
||||
pytest.param("DELETE", "/v1/vector_stores/{}", None, id="provider-delete"),
|
||||
pytest.param("POST", "/v1/vector_stores/{}/files", {"file_id": "file-1"}, id="file-create"),
|
||||
pytest.param("POST", "/v1/vector_stores/{}/files/file-1", {"attributes": {"topic": "parks"}}, id="file-update"),
|
||||
pytest.param("DELETE", "/v1/vector_stores/{}/files/file-1", None, id="file-delete"),
|
||||
)
|
||||
|
||||
|
||||
def _body_for(body: Mapping[str, object] | None, store_id: str) -> object:
|
||||
return None if body is None else json.loads(json.dumps(body).replace("{}", store_id))
|
||||
|
||||
|
||||
def _openai_answer(request: httpx.Request) -> httpx.Response:
|
||||
path: Final = request.url.path
|
||||
if path.endswith("/responses"):
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "resp_1",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [],
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
)
|
||||
if path.endswith("/chat/completions"):
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "chat_1",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
)
|
||||
if path.endswith("/search"):
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"object": "vector_store.search_results.page",
|
||||
"search_query": "hi",
|
||||
"data": [],
|
||||
"has_more": False,
|
||||
"next_page": None,
|
||||
},
|
||||
)
|
||||
if request.method == "DELETE":
|
||||
return httpx.Response(
|
||||
200, json={"id": path.rsplit("/", 1)[-1], "object": "vector_store.deleted", "deleted": True}
|
||||
)
|
||||
if path.endswith("/files"):
|
||||
return httpx.Response(
|
||||
200, json={"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False}
|
||||
)
|
||||
if path.endswith("/content"):
|
||||
return httpx.Response(200, json={"object": "vector_store.file_content.page", "data": [], "has_more": False})
|
||||
if "/files/" in path:
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "file-1",
|
||||
"object": "vector_store.file",
|
||||
"created_at": 1,
|
||||
"usage_bytes": 0,
|
||||
"vector_store_id": "store-a",
|
||||
"status": "completed",
|
||||
},
|
||||
)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "store-a",
|
||||
"object": "vector_store",
|
||||
"created_at": 1,
|
||||
"name": "store-a",
|
||||
"usage_bytes": 0,
|
||||
"file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
|
||||
"status": "completed",
|
||||
"metadata": {},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def openai_stub(monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, httpx_transport: None) -> None:
|
||||
respx_mock.route(host="api.openai.com").mock(side_effect=_openai_answer)
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.llm_router",
|
||||
Router(
|
||||
model_list=[{"model_name": "m", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}}]
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry(vector_stores=list(_STORES)))
|
||||
|
||||
|
||||
async def _statuses(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
layer: Layer,
|
||||
how: Literal["own-list", "group"],
|
||||
method: str,
|
||||
path: str,
|
||||
body: Mapping[str, object] | None,
|
||||
) -> Sequence[int]:
|
||||
stores: Final = {store["vector_store_id"]: store for store in _STORES}
|
||||
|
||||
def row(where: Mapping[str, str], data: object = None) -> MagicMock:
|
||||
return MagicMock(model_dump=MagicMock(return_value=stores[where["vector_store_id"]]))
|
||||
|
||||
proxy: Final = ProxyState(monkeypatch, store_reads=AsyncMock(side_effect=row))
|
||||
monkeypatch.setattr(proxy.prisma.db.litellm_managedvectorstorestable, "update", AsyncMock(side_effect=row))
|
||||
monkeypatch.setattr(proxy.prisma.db.litellm_managedvectorstorestable, "delete", AsyncMock())
|
||||
token: Final = hash_token(_KEY)
|
||||
caller: Final = (await _caller(proxy, layer, how)).model_copy(
|
||||
update={"token": token, "allowed_routes": [path] if path.startswith("/vector_store/") else []}
|
||||
)
|
||||
await proxy.cache.async_set_cache(key=token, value=caller)
|
||||
monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY)
|
||||
client: Final = TestClient(app, raise_server_exceptions=False)
|
||||
return [
|
||||
client.request(
|
||||
method,
|
||||
path.format(store_id),
|
||||
json=_body_for(body, store_id),
|
||||
headers={"Authorization": f"Bearer {_KEY}"},
|
||||
).status_code
|
||||
for store_id in stores
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("layer", ["key", "team"])
|
||||
@pytest.mark.parametrize(("method", "path", "body"), [*READ_ROUTES, *MANAGE_ROUTES])
|
||||
async def test_every_store_route_answers_a_group_store_like_an_own_list_store(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
openai_stub: None,
|
||||
layer: Layer,
|
||||
method: str,
|
||||
path: str,
|
||||
body: Mapping[str, object] | None,
|
||||
) -> None:
|
||||
"""Real auth runs the request check and the route's store gate. Statuses are for store-a (granted),
|
||||
store-b (another team's) and store-free (no team), in that order."""
|
||||
listed: Final = await _statuses(monkeypatch, layer, "own-list", method, path, body)
|
||||
grouped: Final = await _statuses(monkeypatch, layer, "group", method, path, body)
|
||||
|
||||
assert grouped == listed
|
||||
assert grouped[0] == 200
|
||||
assert grouped[1] in (401, 403)
|
||||
Loading…
Add table
Reference in a new issue