scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

feat(access-groups): grant vector stores through access groups
Some checks are pending
Cost map pins recheck / rerun-cost-map-guard (push) Waiting to run
Details
LiteLLM Rust DB / db-queries (push) Waiting to run
Details
LiteLLM Rust / rust-lint (push) Waiting to run
Details
LiteLLM Rust / rust-test (push) Waiting to run
Details
LiteLLM Rust / rust-wheel (push) Waiting to run
Details

Browse source 
An access group can now list vector stores. A key or team in the group can use them,
and a non-empty combined list restricts it to its own stores plus its groups' stores.
This commit is contained in:
shrey kharbanda 2026-11-05 14:09:41 -08:00
parent a6b61a2c3b
commit 3f3a63ad97
24 changed files with 1397 additions and 63 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm-proxy-extras/litellm_proxy_extras/migrations/20261110120000_add_access_group_vector_store_ids/migration.sql Normal file
View file

@ -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[];

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1595,6 +1595,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
litellm/models/access_group.py
View file

@ -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

4
litellm/models/team.py
View file

@ -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."""

36
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -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": [
{

257
litellm/proxy/auth/auth_checks.py
View file

@ -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,
@ -5788,6 +5789,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,
@ -5816,9 +5828,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:
@ -5837,9 +5846,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))
@ -8076,8 +8083,215 @@ def _legacy_requested_vector_store_ids(
)
def _restricts_vector_stores(grants: Sequence[tuple[GrantLayer, LiteLLM_ObjectPermissionTable | 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
) -> Mapping[GrantLayer, tuple[str, ...]]:
"""Stores the key's and its team's access groups grant, per layer.
A deleted group grants nothing, and an unreadable one raises."""
key_groups: Final = await get_access_group_tables(
(valid_token.access_group_ids or ()) if valid_token is not None else (), unreadable="raise"
)
team_groups: Final = await get_access_group_tables(
(team_object.access_group_ids or ()) if team_object is not None else (), unreadable="raise"
)
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)),
}
)
@dataclass(frozen=True, slots=True)
class _CombinedVectorStores:
vector_stores: Sequence[str]
def _with_group_stores(
grant: LiteLLM_ObjectPermissionTable | None, group_stores: Sequence[str]
) -> _CombinedVectorStores:
own_stores: Final = (grant.vector_stores or ()) if grant is not None else ()
return _CombinedVectorStores((*own_stores, *group_stores))
async def _grants_with_group_stores(
group_stores: Mapping[GrantLayer, tuple[str, ...]],
valid_token: UserAPIKeyAuth | None,
team_object: LiteLLM_TeamTable | None,
user_object: LiteLLM_UserTable | None,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
) -> tuple[tuple[GrantLayer, _CombinedVectorStores], ...]:
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
)
)
async def _team_for_group_check(
valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None, team_is_strict: bool
) -> LiteLLM_TeamTable | None:
"""A team we cannot read must deny, not widen: auth's rebuilt stand-in has none of the team's groups.
A strict team layer can only lose grants without them, so it keeps the stand-in."""
team_id: Final = valid_token.team_id if valid_token is not None else None
if (
team_is_strict
or 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,
)
async def _combined_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,
team_is_strict: bool = False,
) -> tuple[tuple[GrantLayer, _CombinedVectorStores], ...]:
try:
team: Final = await _team_for_group_check(valid_token, team_object, team_is_strict)
group_stores: Final = await access_group_vector_stores(valid_token, team)
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 access groups to check vector store access",
type=ProxyErrorTypes.internal_server_error,
param="vector_store",
code=status.HTTP_503_SERVICE_UNAVAILABLE,
) from e
return await _grants_with_group_stores(
group_stores, valid_token, team, user_object, prisma_client, user_api_key_cache
)
def _require_every_layer_allows(
grants: Sequence[tuple[GrantLayer, _CombinedVectorStores]],
strict_layers: frozenset[GrantLayer],
vector_store_ids: Sequence[str],
) -> None:
for layer, grant in grants:
if layer in strict_layers:
_require_vector_store_grant(layer, vector_store_ids, grant)
else:
_can_object_call_vector_stores(
object_type=layer,
vector_store_ids_to_run=vector_store_ids,
object_permissions=grant,
)
def _every_layer_allows(
grants: Sequence[tuple[GrantLayer, _CombinedVectorStores]],
strict_layers: frozenset[GrantLayer],
vector_store_id: str,
) -> bool:
try:
_require_every_layer_allows(grants, strict_layers, (vector_store_id,))
except ProxyException:
return False
return True
async def vector_stores_opened_by_groups(
valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None
) -> frozenset[str]:
"""The stores the caller's groups list and the request check would allow. An unreadable group raises."""
from litellm.proxy.proxy_server import (
general_settings, # pyright: ignore[reportUnknownVariableType] # untyped module-level dict in proxy_server
prisma_client,
user_api_key_cache,
)
group_stores: Final = await access_group_vector_stores(valid_token, team_object)
listed: Final = tuple(chain.from_iterable(group_stores.values()))
if not listed or prisma_client is None:
return frozenset()
grants: Final = await _grants_with_group_stores(
group_stores, valid_token, team_object, None, prisma_client, user_api_key_cache
)
strict_layers: Final = _strict_grant_layers(
_vector_store_deny_by_default(_typed_request_body(general_settings)), valid_token, team_object
)
return frozenset(store for store in listed if _every_layer_allows(grants, strict_layers, store))
async def _team_for_own_grants(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTable | None:
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
if valid_token.team_id is None or valid_token.team_id == UI_TEAM_ID:
return None
try:
return await get_team_object(
team_id=valid_token.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,
)
except Exception as e: # noqa: BLE001 # auth falls back to the key's fields too, read just below
verbose_proxy_logger.debug("Team unreadable for a vector store change, using the key's fields: %s", e)
return None
async def assert_own_grants_allow_vector_store(valid_token: UserAPIKeyAuth, vector_store_id: str) -> None:
"""Changing a store needs the caller's own grants, read the way the request check reads them, as on main.
Access groups only grant reads."""
from litellm.proxy.proxy_server import (
general_settings, # pyright: ignore[reportUnknownVariableType] # untyped module-level dict in proxy_server
prisma_client,
user_api_key_cache,
)
if prisma_client is None:
return
team_object: Final = await _team_for_own_grants(valid_token)
team_grant: Final = (
await _cached_object_permission(
team_object.object_permission_id, team_object.object_permission, prisma_client, user_api_key_cache
)
if team_object is not None
else await _cached_object_permission(
valid_token.team_object_permission_id, None, prisma_client, user_api_key_cache
)
)
key_grant: Final = await _cached_object_permission(
valid_token.object_permission_id, valid_token.object_permission, prisma_client, user_api_key_cache
)
own_grants: Final[tuple[tuple[GrantLayer, LiteLLM_ObjectPermissionTable | None], ...]] = (
("key", key_grant),
("team", team_grant),
)
_require_every_layer_allows(
tuple((layer, _with_group_stores(grant, ())) for layer, grant in own_grants),
_strict_grant_layers(
_vector_store_deny_by_default(_typed_request_body(general_settings)), valid_token, team_object
),
(vector_store_id,),
)
def _restricts_vector_stores(grants: Sequence[tuple[GrantLayer, _CombinedVectorStores]]) -> bool:
return any(bool(grant.vector_stores) for _, grant in grants)
def _request_drop_policies(
@ -8136,13 +8350,13 @@ async def vector_store_access_check(
user_object: LiteLLM_UserTable | None = None,
deny_by_default: bool = False,
llm_router: Router | None = None,
):
) -> Literal[True]:
"""
Checks whether the caller may use every vector store the request names.
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 grant.
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
@ -8190,7 +8404,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 _combined_vector_store_grants(valid_token, team_object, None, prisma_client, user_api_key_cache)
):
raise _malformed_vector_store_ids()
return True
@ -8204,26 +8418,19 @@ 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 _combined_vector_store_grants(
valid_token,
team_object,
user_object if "user" in strict_layers else None,
prisma_client,
user_api_key_cache,
team_is_strict="team" in strict_layers,
)
for layer, grant in grants:
if layer in strict_layers:
_require_vector_store_grant(layer, vector_store_ids_to_run, grant)
elif grant is not None:
_can_object_call_vector_stores(
object_type=layer,
vector_store_ids_to_run=vector_store_ids_to_run,
object_permissions=grant,
)
_require_every_layer_allows(grants, strict_layers, vector_store_ids_to_run)
return True
def _provider_vector_store_id(vector_store_id: str) -> str:
def provider_vector_store_id(vector_store_id: str) -> str:
if not is_base64_encoded_unified_id(vector_store_id):
return vector_store_id
return extract_provider_resource_id_from_unified_id(vector_store_id) or vector_store_id
@ -8233,7 +8440,7 @@ def _can_object_call_vector_stores(
object_type: Literal["key", "team", "org", "user"],
vector_store_ids_to_run: Sequence[str],
object_permissions: _VectorStorePermissionsRow | None,
):
) -> Literal[True]:
"""
Raises ProxyException if the object (key, team, org) cannot access the specific vector store.
"""
@ -8247,9 +8454,9 @@ def _can_object_call_vector_stores(
if len(object_permissions.vector_stores) == 0:
return True
allowed_stores: Final = frozenset(map(_provider_vector_store_id, object_permissions.vector_stores))
allowed_stores: Final = frozenset(map(provider_vector_store_id, object_permissions.vector_stores))
for vector_store_id in vector_store_ids_to_run:
if _provider_vector_store_id(vector_store_id) not in allowed_stores:
if provider_vector_store_id(vector_store_id) not in allowed_stores:
raise ProxyException(
message=f"User not allowed to access vector store. Tried to access {vector_store_id}. Only allowed to access {object_permissions.vector_stores}",
type=ProxyErrorTypes.get_vector_store_access_error_type_for_object(object_type),

5
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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, resolve_served_model
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
@ -3320,7 +3321,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
)
except HTTPException:
token_team_models: Final = _token_team_models(valid_token)
_team_obj = LiteLLM_TeamTableCachedObj(
_team_obj = LiteLLM_TeamRebuiltFromToken( # rebind-ok: pre-existing rebinding on a renamed line
team_id=valid_token.team_id,
max_budget=valid_token.team_max_budget,
soft_budget=valid_token.team_soft_budget,
@ -3472,7 +3473,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached
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,

4
litellm/proxy/common_request_processing.py
View file

@ -913,11 +913,11 @@ async def _authorize_response_file_search_vector_stores(
return
from litellm.proxy.vector_store_endpoints.utils import (
assert_user_can_access_vector_store_id,
assert_user_can_read_vector_store_id,
)
for vector_store_id in sorted(vector_store_ids):
await assert_user_can_access_vector_store_id(
await assert_user_can_read_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)

2
litellm/proxy/management_endpoints/access_group_endpoints.py
View file

@ -572,6 +572,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,
@ -670,6 +671,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] = (

14
litellm/proxy/rag_endpoints/endpoints.py
View file

@ -8,7 +8,7 @@ Provides:
import base64
import json
from collections.abc import Mapping
from collections.abc import Awaitable, Callable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
from urllib.parse import urlsplit
@ -68,6 +68,7 @@ from litellm.proxy.vector_store_endpoints.endpoints import (
)
from litellm.proxy.vector_store_endpoints.utils import (
assert_user_can_access_vector_store_id,
assert_user_can_read_vector_store_id,
with_managed_search_file_ids,
)
from litellm.rag.main import get_ingestion_class
@ -164,19 +165,14 @@ def _collect_vector_store_ids_from_payload(payload: object) -> set[str]:
async def _authorize_nested_vector_store_ids(
payload: object,
user_api_key_dict: UserAPIKeyAuth,
authorize: Callable[[str, UserAPIKeyAuth], Awaitable[LiteLLM_ManagedVectorStore | None]],
) -> Mapping[str, LiteLLM_ManagedVectorStore]:
"""Authorize every nested vector store id and return the managed stores it resolved."""
return MappingProxyType(
{
vector_store_id: store
for vector_store_id in sorted(_collect_vector_store_ids_from_payload(payload))
if (
store := await assert_user_can_access_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)
)
is not None
if (store := await authorize(vector_store_id, user_api_key_dict)) is not None
}
)
@ -773,6 +769,7 @@ async def rag_ingest(
resolved_stores: Final = await _authorize_nested_vector_store_ids(
payload=ingest_options,
user_api_key_dict=user_api_key_dict,
authorize=assert_user_can_access_vector_store_id,
)
request_vector_store_config: Final = ingest_options.get("vector_store", {})
@ -1018,6 +1015,7 @@ async def rag_query(
resolved_stores: Final = await _authorize_nested_vector_store_ids(
payload=retrieval_config,
user_api_key_dict=user_api_key_dict,
authorize=assert_user_can_read_vector_store_id,
)
# Merge litellm-managed vector store params (provider, region, embedding

1
litellm/proxy/schema.prisma
View file

@ -1595,6 +1595,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([])

34
litellm/proxy/vector_store_endpoints/endpoints.py
View file

@ -1,4 +1,4 @@
from collections.abc import Mapping, Sequence
from collections.abc import Awaitable, Callable, Mapping, Sequence
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -22,6 +22,7 @@ from litellm.litellm_core_utils.core_helpers import (
require_attributes,
)
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import assert_own_grants_allow_vector_store
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.http_parsing_utils import read_raw_json_body
@ -30,6 +31,7 @@ from litellm.proxy.utils import jsonify_object
from litellm.proxy.vector_store_endpoints.utils import (
assert_proxy_admin_for_vector_store_index_management,
assert_user_can_access_vector_store,
assert_user_can_read_vector_store,
get_litellm_managed_vector_store,
with_managed_search_file_ids,
)
@ -121,6 +123,8 @@ async def _update_request_data_with_litellm_managed_vector_store_registry(
data: dict,
vector_store_id: str,
user_api_key_dict: UserAPIKeyAuth | None = None,
*,
authorize: Callable[[LiteLLM_ManagedVectorStore, UserAPIKeyAuth], Awaitable[None]],
) -> dict:
"""
Update the request data with the litellm managed vector store registry.
@ -129,6 +133,7 @@ async def _update_request_data_with_litellm_managed_vector_store_registry(
data: Request data to update
vector_store_id: ID of the vector store
user_api_key_dict: User API key authentication info for access control
authorize: Refuses a caller who may not run this request on the vector store
Raises:
HTTPException: If user doesn't have access to the vector store
@ -139,10 +144,7 @@ async def _update_request_data_with_litellm_managed_vector_store_registry(
if vector_store_to_run is None:
return data
if user_api_key_dict is not None:
await assert_user_can_access_vector_store(
vector_store=vector_store_to_run,
user_api_key_dict=user_api_key_dict,
)
await authorize(vector_store_to_run, user_api_key_dict)
return {**data, **build_request_data_from_managed_vector_store(vector_store_to_run)}
@ -188,7 +190,10 @@ async def vector_store_search(
# Check for legacy vector store registry (non-managed vector stores)
data = await _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
data=data,
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
authorize=assert_user_can_read_vector_store,
)
reject_caller_opensearch_index_name(
caller_body=caller_body,
@ -381,7 +386,10 @@ async def vector_store_retrieve(
data = {"vector_store_id": vector_store_id}
data = await _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
data=data,
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
authorize=assert_user_can_read_vector_store,
)
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
@ -519,8 +527,12 @@ async def vector_store_update(
await read_raw_json_body(request), VectorStoreWriteBody.model_validate(data)
)
await assert_own_grants_allow_vector_store(user_api_key_dict, vector_store_id)
data = await _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
data=data,
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
authorize=assert_user_can_access_vector_store,
)
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
@ -582,8 +594,12 @@ async def vector_store_delete(
data = {"vector_store_id": vector_store_id}
await assert_own_grants_allow_vector_store(user_api_key_dict, vector_store_id)
data = await _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
data=data,
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
authorize=assert_user_can_access_vector_store,
)
processor: Final = ProxyBaseLLMRequestProcessing(data=data)

9
litellm/proxy/vector_store_endpoints/management_endpoints.py
View file

@ -10,6 +10,7 @@ All /vector_store management endpoints
import copy
import json
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any, Final
from fastapi import APIRouter, Depends, HTTPException
@ -34,6 +35,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
from litellm.proxy.vector_store_endpoints.utils import (
can_user_access_vector_store,
can_user_read_vector_store,
filter_listable_vector_stores,
)
from litellm.repositories.prisma_protocols import TableActions
@ -139,6 +141,7 @@ async def _fetch_and_authorize_vector_store(
vector_store_id: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: "PrismaClient",
is_authorized: Callable[[LiteLLM_ManagedVectorStore, UserAPIKeyAuth], Awaitable[bool]],
) -> "LiteLLM_ManagedVectorStore":
"""
Look up a vector store by id and confirm the caller can access it.
@ -152,7 +155,7 @@ async def _fetch_and_authorize_vector_store(
detail=f"Vector store with ID {vector_store_id} not found",
)
typed: Final = _row_to_vector_store(row)
if not await _check_vector_store_access(typed, user_api_key_dict):
if not await is_authorized(typed, user_api_key_dict):
raise HTTPException(
status_code=403,
detail="Access denied: You do not have permission to access this vector store",
@ -559,7 +562,7 @@ async def get_vector_store_info(
)
if vector_store is not None:
# Check access control
if not await _check_vector_store_access(vector_store, user_api_key_dict):
if not await can_user_read_vector_store(vector_store, user_api_key_dict):
raise HTTPException(
status_code=403,
detail="Access denied: You do not have permission to access this vector store",
@ -593,6 +596,7 @@ async def get_vector_store_info(
vector_store_id=data.vector_store_id,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
is_authorized=can_user_read_vector_store,
)
vector_store_dict: Final = dict(vector_store_typed)
if "litellm_params" in vector_store_dict:
@ -641,6 +645,7 @@ async def update_vector_store(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
is_authorized=_check_vector_store_access,
)
# Handle metadata serialization

86
litellm/proxy/vector_store_endpoints/utils.py
View file

@ -253,16 +253,67 @@ async def _is_vector_store_granted_to_any(
return False
async def _group_store_ids_for_context(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]:
"""Reads only add access here, so a group or team that cannot be read grants nothing."""
from litellm.proxy.auth.auth_checks import get_team_object, vector_stores_opened_by_groups
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
team_id: Final = user_api_key_dict.team_id
try:
team: Final = (
await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_dict.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if team_id is not None
else None
)
return await vector_stores_opened_by_groups(user_api_key_dict, team)
except Exception as e: # noqa: BLE001 # grants nothing, per the docstring
verbose_proxy_logger.debug("Could not read access groups for vector store access: %s", e)
return frozenset()
async def _group_granted_vector_store_ids(auth_contexts: tuple[UserAPIKeyAuth, ...]) -> frozenset[str]:
return frozenset[str]().union(*[await _group_store_ids_for_context(context) for context in auth_contexts])
def _is_group_granted(vector_store: LiteLLM_ManagedVectorStore, group_store_ids: frozenset[str]) -> bool:
"""Exact ids only: a provider id can belong to several registered stores, on different backends."""
return (vector_store.get("vector_store_id") or "") in group_store_ids
async def can_user_read_vector_store(
vector_store: LiteLLM_ManagedVectorStore,
user_api_key_dict: UserAPIKeyAuth,
) -> bool:
"""Using or reading a store also counts access groups. Managing one does not."""
if await can_user_access_vector_store(vector_store, user_api_key_dict):
return True
auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict)
return _is_group_granted(vector_store, await _group_granted_vector_store_ids(auth_contexts))
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 _group_granted_vector_store_ids(auth_contexts)
return tuple(
[
vs
for vs in vector_stores
if _is_group_granted(vs, group_store_ids) or await _is_vector_store_granted_to_any(vs, auth_contexts)
]
)
async def get_litellm_managed_vector_store(
@ -339,6 +390,17 @@ async def assert_user_can_access_vector_store(
raise HTTPException(status_code=403, detail=detail)
async def assert_user_can_change_vector_store_id(
vector_store_id: str,
user_api_key_dict: UserAPIKeyAuth,
) -> LiteLLM_ManagedVectorStore | None:
"""Group grants never change a store, managed or not, so the caller's own grants decide first."""
from litellm.proxy.auth.auth_checks import assert_own_grants_allow_vector_store
await assert_own_grants_allow_vector_store(user_api_key_dict, vector_store_id)
return await assert_user_can_access_vector_store_id(vector_store_id, user_api_key_dict)
async def assert_user_can_access_vector_store_id(
vector_store_id: str,
user_api_key_dict: UserAPIKeyAuth,
@ -359,6 +421,26 @@ async def assert_user_can_access_vector_store_id(
return vector_store
async def assert_user_can_read_vector_store(
vector_store: LiteLLM_ManagedVectorStore,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
if not await can_user_read_vector_store(vector_store, user_api_key_dict):
raise HTTPException(
status_code=403, detail="Access denied: You do not have permission to access this vector store"
)
async def assert_user_can_read_vector_store_id(
vector_store_id: str,
user_api_key_dict: UserAPIKeyAuth,
) -> LiteLLM_ManagedVectorStore | None:
vector_store: Final = await get_litellm_managed_vector_store(vector_store_id=vector_store_id)
if vector_store is not None:
await assert_user_can_read_vector_store(vector_store, user_api_key_dict)
return vector_store
def _does_endpoint_match(endpoint_path: str, request_path: str) -> bool:
if endpoint_path in request_path:
return True

15
litellm/proxy/vector_store_files_endpoints/endpoints.py
View file

@ -39,7 +39,8 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
)
from litellm.proxy.rag_endpoints.upload_security import safe_download_headers
from litellm.proxy.vector_store_endpoints.utils import (
assert_user_can_access_vector_store_id,
assert_user_can_change_vector_store_id,
assert_user_can_read_vector_store_id,
is_allowed_to_call_vector_store_files_endpoint,
)
from litellm.router_rust.invoke.endpoints import PINNED_DEPLOYMENT_KEY
@ -738,7 +739,7 @@ async def vector_store_file_create(
data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line
data["vector_store_id"] = vector_store_id
managed_vector_store: Final = await assert_user_can_access_vector_store_id(
managed_vector_store: Final = await assert_user_can_change_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)
@ -847,7 +848,7 @@ async def vector_store_file_list(
data: dict[str, str | None] = {"vector_store_id": vector_store_id}
data.update(query_params)
data["vector_store_id"] = vector_store_id
managed_vector_store: Final = await assert_user_can_access_vector_store_id(
managed_vector_store: Final = await assert_user_can_read_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)
@ -955,7 +956,7 @@ async def vector_store_file_retrieve(
"vector_store_id": vector_store_id,
"file_id": file_id,
}
managed_vector_store: Final = await assert_user_can_access_vector_store_id(
managed_vector_store: Final = await assert_user_can_read_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)
@ -1063,7 +1064,7 @@ async def vector_store_file_content(
"vector_store_id": vector_store_id,
"file_id": file_id,
}
managed_vector_store: Final = await assert_user_can_access_vector_store_id(
managed_vector_store: Final = await assert_user_can_read_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)
@ -1174,7 +1175,7 @@ async def vector_store_file_update(
data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line
data["vector_store_id"] = vector_store_id
data["file_id"] = file_id
managed_vector_store: Final = await assert_user_can_access_vector_store_id(
managed_vector_store: Final = await assert_user_can_change_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)
@ -1282,7 +1283,7 @@ async def vector_store_file_delete(
"vector_store_id": vector_store_id,
"file_id": file_id,
}
managed_vector_store: Final = await assert_user_can_access_vector_store_id(
managed_vector_store: Final = await assert_user_can_change_vector_store_id(
vector_store_id=vector_store_id,
user_api_key_dict=user_api_key_dict,
)

3
litellm/types/access_group.py
View file

@ -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, ...]

1
schema.prisma
View file

@ -1595,6 +1595,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([])

323
tests/unit/proxy/auth/test_vector_store_access_groups.py Normal file
View file

@ -0,0 +1,323 @@
"""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
from unittest.mock import AsyncMock
import pytest
from litellm import Router
from litellm.models.team import LiteLLM_TeamRebuiltFromToken
from litellm.proxy._types import (
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamTableCachedObj,
ProxyException,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import assert_own_grants_allow_vector_store, vector_store_access_check
from litellm.proxy.auth.user_api_key_auth import _team_obj_from_token
from tests.unit.proxy.vector_store_access_group_state import (
TEAM_ID,
TOKEN,
ProxyState,
make_key,
make_team,
)
@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
@pytest.mark.asyncio
async def test_managed_store_id_in_a_group_matches_its_provider_id(proxy: ProxyState) -> None:
managed_id: Final = base64.urlsafe_b64encode(
b"litellm_proxy:vector_store;unified_id,abc;target_model_names,gpt-4;resource_id,vs_provider"
).decode()
await proxy.add_group("g1", (managed_id,))
assert await _request_is_allowed(make_key(groups=("g1",)), None, "vs_provider") is True
@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
@pytest.mark.asyncio
@pytest.mark.parametrize(
("key_groups", "team_groups"),
[pytest.param(("g1",), (), id="key-group"), pytest.param((), ("g1",), id="team-group")],
)
async def test_unreadable_group_blocks_the_request_instead_of_lifting_the_restriction(
monkeypatch: pytest.MonkeyPatch, key_groups: Sequence[str], team_groups: Sequence[str]
) -> None:
ProxyState(monkeypatch, group_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
with pytest.raises(ProxyException) as denied:
await vector_store_access_check(
request_body={"vector_store_ids": ["store-b"]},
team_object=make_team(groups=team_groups),
valid_token=make_key(groups=key_groups, team_id=TEAM_ID),
)
assert denied.value.code == "503"
assert denied.value.message == "Could not read the team or access groups to check vector store access"
@pytest.mark.asyncio
async def test_a_team_auth_rebuilt_without_groups_is_read_again_and_its_groups_still_apply(proxy: ProxyState) -> None:
await proxy.add_group("g1", ("store-a",))
await proxy.add_team(make_team(groups=("g1",)))
rebuilt: Final = LiteLLM_TeamRebuiltFromToken(team_id=TEAM_ID)
assert await _request_is_allowed(make_key(team_id=TEAM_ID), rebuilt, "store-a") is True
assert await _request_is_allowed(make_key(team_id=TEAM_ID), rebuilt, "store-b") is False
@pytest.mark.asyncio
async def test_a_team_auth_rebuilt_without_groups_blocks_the_request_when_the_team_is_still_unreadable(
monkeypatch: pytest.MonkeyPatch,
) -> None:
ProxyState(monkeypatch, team_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
with pytest.raises(ProxyException) as denied:
await vector_store_access_check(
request_body={"vector_store_ids": ["store-b"]},
team_object=LiteLLM_TeamRebuiltFromToken(team_id=TEAM_ID),
valid_token=make_key(team_id=TEAM_ID),
)
assert denied.value.code == "503"
@pytest.mark.asyncio
async def test_a_store_change_reads_the_team_grant_as_the_request_check_does_not_the_keys_cached_id(
proxy: ProxyState,
) -> None:
await proxy.cache.async_set_cache(
key="object_permission_id:old-grant",
value=LiteLLM_ObjectPermissionTable(object_permission_id="old-grant", vector_stores=["store-a"]),
)
await proxy.add_team(make_team(None))
key: Final = UserAPIKeyAuth(token=TOKEN, team_id=TEAM_ID, team_object_permission_id="old-grant")
await vector_store_access_check(
request_body={"vector_store_ids": ["store-b"]}, team_object=make_team(None), valid_token=key
)
await assert_own_grants_allow_vector_store(key, "store-b")
@pytest.mark.asyncio
async def test_a_store_change_uses_the_teams_current_grant(proxy: ProxyState) -> None:
await proxy.add_team(make_team(("store-a",)))
key: Final = UserAPIKeyAuth(token=TOKEN, team_id=TEAM_ID, team_object_permission_id=None)
await assert_own_grants_allow_vector_store(key, "store-a")
with pytest.raises(ProxyException) as denied:
await assert_own_grants_allow_vector_store(key, "store-b")
assert denied.value.code == "401"
@pytest.mark.asyncio
async def test_a_team_auth_rebuilt_from_the_key_keeps_the_teams_own_stores_and_groups(proxy: ProxyState) -> None:
await proxy.add_group("g1", ("store-a",))
await proxy.add_team(make_team(("store-own",), groups=("g1",)))
key: Final = make_key(team_id=TEAM_ID)
rebuilt: Final = _team_obj_from_token(key)
assert await _request_is_allowed(key, rebuilt, "store-own") is True
assert await _request_is_allowed(key, rebuilt, "store-a") is True
assert await _request_is_allowed(key, rebuilt, "store-b") 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_malformed_store_ids_with_an_unreadable_group_block_the_request(monkeypatch: pytest.MonkeyPatch) -> None:
ProxyState(monkeypatch, group_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
with pytest.raises(ProxyException) as denied:
await vector_store_access_check(
request_body={"model": "gpt-4o-mini", "vector_store_ids": "store-b"},
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"}}]),
)
assert denied.value.code == "503"
@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

50
tests/unit/proxy/management_endpoints/test_access_group_endpoints.py
View file

@ -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
View 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)

489
tests/unit/proxy/vector_store_endpoints/test_vector_store_access_groups.py Normal file
View file

@ -0,0 +1,489 @@
"""Group grants open the read gate and the list, and never the manage gate."""
import asyncio
import base64
import json
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import Final
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.models.access_group import LiteLLM_AccessGroupTable
from litellm.proxy import proxy_server
from litellm.proxy._types import UI_TEAM_ID, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import cache_access_object
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.proxy_server import app
from litellm.proxy.vector_store_endpoints.utils import (
can_user_read_vector_store,
filter_listable_vector_stores,
)
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
from tests.unit.proxy.vector_store_access_group_state import (
OTHER_TEAM_ID,
TEAM_ID,
TOKEN,
ProxyState,
make_key,
make_team,
team_owned_store,
)
@pytest.fixture
def proxy(monkeypatch: pytest.MonkeyPatch) -> ProxyState:
return ProxyState(monkeypatch)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("store_id", "allowed"),
[pytest.param("store-a", True, id="group-store"), pytest.param("store-b", False, id="other-store")],
)
async def test_search_and_files_gate_counts_team_group_stores(proxy: ProxyState, store_id: str, allowed: bool) -> None:
await proxy.add_group("g1", ("store-a",))
await proxy.add_team(make_team(groups=("g1",)))
assert await can_user_read_vector_store(team_owned_store(store_id), make_key(team_id=TEAM_ID)) is allowed
@pytest.mark.asyncio
async def test_search_and_files_gate_counts_a_personal_keys_group_stores(proxy: ProxyState) -> None:
await proxy.add_group("g1", ("store-a",))
assert await can_user_read_vector_store(team_owned_store("store-a"), make_key(groups=("g1",))) is True
assert await can_user_read_vector_store(team_owned_store("store-b"), make_key(groups=("g1",))) is False
@pytest.mark.asyncio
async def test_models_only_group_grants_nothing_on_search_and_list(proxy: ProxyState) -> None:
await proxy.add_group("models-only", models=("gpt-4",))
await proxy.add_team(make_team(groups=("models-only",)))
assert await can_user_read_vector_store(team_owned_store("store-a"), make_key(team_id=TEAM_ID)) is False
assert await filter_listable_vector_stores((team_owned_store("store-a"),), make_key(team_id=TEAM_ID)) == ()
@pytest.mark.asyncio
async def test_direct_grant_never_reads_teams_or_groups(monkeypatch: pytest.MonkeyPatch) -> None:
team_reads: Final = AsyncMock(side_effect=RuntimeError("database unavailable"))
proxy: Final = ProxyState(monkeypatch, team_reads=team_reads)
assert (
await can_user_read_vector_store(team_owned_store("store-a"), make_key(("store-a",), team_id=TEAM_ID)) is True
)
assert team_reads.await_count == 0
assert proxy.group_reads.await_count == 0
@pytest.mark.asyncio
async def test_list_shows_only_granted_stores(proxy: ProxyState) -> None:
await proxy.add_group("g1", ("store-a",))
await proxy.add_team(make_team(groups=("g1",)))
listed: Final = await filter_listable_vector_stores(
(team_owned_store("store-a"), team_owned_store("store-b")), make_key(team_id=TEAM_ID)
)
assert listed == (team_owned_store("store-a"),)
@pytest.mark.asyncio
async def test_list_reads_the_group_once_however_many_stores_exist(monkeypatch: pytest.MonkeyPatch) -> None:
cache: Final = MagicMock(wraps=UserApiKeyCache())
proxy: Final = ProxyState(monkeypatch, cache=cache)
await proxy.add_group("g1", ("store-a",))
await proxy.add_team(make_team(groups=("g1",)))
await filter_listable_vector_stores(
(team_owned_store("store-a"), team_owned_store("store-b"), team_owned_store("store-c")),
make_key(team_id=TEAM_ID),
)
group_lookups: Final = [
call for call in cache.async_get_cache.call_args_list if call.kwargs["key"] == "access_group_id:g1"
]
assert len(group_lookups) == 1
@pytest.mark.asyncio
async def test_unreadable_group_grants_nothing_on_search_and_list(monkeypatch: pytest.MonkeyPatch) -> None:
proxy: Final = ProxyState(monkeypatch, group_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
await proxy.add_team(make_team(groups=("g1",)))
assert await can_user_read_vector_store(team_owned_store("store-a"), make_key(team_id=TEAM_ID)) is False
assert await filter_listable_vector_stores((team_owned_store("store-a"),), make_key(team_id=TEAM_ID)) == ()
@pytest.mark.asyncio
async def test_one_unreadable_group_drops_the_readable_groups_grants_too(monkeypatch: pytest.MonkeyPatch) -> None:
proxy: Final = ProxyState(monkeypatch, group_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
await proxy.add_group("g1", ("store-a",))
await proxy.add_team(make_team(groups=("g1", "unreadable")))
assert await can_user_read_vector_store(team_owned_store("store-a"), make_key(team_id=TEAM_ID)) is False
@pytest.mark.asyncio
async def test_failed_team_read_grants_nothing_and_keeps_direct_grants(monkeypatch: pytest.MonkeyPatch) -> None:
ProxyState(monkeypatch, team_reads=AsyncMock(side_effect=RuntimeError("database unavailable")))
assert await can_user_read_vector_store(team_owned_store("store-a"), make_key(team_id=TEAM_ID)) is False
assert (
await can_user_read_vector_store(team_owned_store("store-a"), make_key(("store-a",), team_id=TEAM_ID)) is True
)
listed: Final = await filter_listable_vector_stores(
(team_owned_store("store-a"),), make_key(("store-a",), team_id=TEAM_ID)
)
assert listed == (team_owned_store("store-a"),)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("key_stores", "key_groups", "team_groups", "store", "allowed"),
[
pytest.param(None, (), ("team-a",), "store-a", True, id="team-group-and-unrestricted-key"),
pytest.param(("store-x",), (), ("team-a",), "store-a", False, id="key-own-list-narrows-team-group"),
pytest.param(None, ("key-b",), ("team-a",), "store-a", False, id="key-group-narrows-team-group"),
pytest.param(None, ("key-b",), ("team-a",), "store-b", False, id="team-group-narrows-key-group"),
pytest.param(None, ("team-a",), ("team-a",), "store-a", True, id="both-layers-list-the-store"),
pytest.param(("store-a",), (), ("team-a",), "store-a", True, id="key-own-list-agrees-with-team-group"),
],
)
async def test_read_gate_and_list_apply_the_request_checks_rule_to_group_grants(
proxy: ProxyState,
key_stores: Sequence[str] | None,
key_groups: Sequence[str],
team_groups: Sequence[str],
store: str,
allowed: bool,
) -> None:
await proxy.add_group("team-a", ("store-a",))
await proxy.add_group("key-b", ("store-b",))
await proxy.add_team(make_team(groups=team_groups))
key: Final = make_key(key_stores, key_groups, TEAM_ID)
assert await can_user_read_vector_store(team_owned_store(store), key) is allowed
listed: Final = await filter_listable_vector_stores((team_owned_store(store),), key)
assert listed == ((team_owned_store(store),) if allowed else ())
@pytest.mark.asyncio
@pytest.mark.parametrize(
("key_groups", "allowed"),
[pytest.param((), False, id="only-the-team-grants"), pytest.param(("g1",), True, id="key-and-team-grant")],
)
async def test_read_gate_requires_both_layers_under_deny_by_default(
proxy: ProxyState, monkeypatch: pytest.MonkeyPatch, key_groups: Sequence[str], allowed: bool
) -> None:
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"vector_store_deny_by_default": True})
await proxy.add_group("g1", ("store-a",))
await proxy.add_team(make_team(groups=("g1",)))
key: Final = UserAPIKeyAuth(token=TOKEN, api_key=TOKEN, team_id=TEAM_ID, access_group_ids=list(key_groups))
key.via_virtual_key = True
assert await can_user_read_vector_store(team_owned_store("store-a"), key) is allowed
@pytest.mark.asyncio
async def test_read_gate_and_list_match_a_group_store_only_by_its_exact_id(proxy: ProxyState) -> None:
managed_id: Final = base64.urlsafe_b64encode(
b"litellm_proxy:vector_store;unified_id,abc;target_model_names,gpt-4;resource_id,store-a"
).decode()
await proxy.add_group("g1", (managed_id,))
await proxy.add_team(make_team(groups=("g1",)))
key: Final = make_key(team_id=TEAM_ID)
assert await can_user_read_vector_store(team_owned_store(managed_id), key) is True
assert await filter_listable_vector_stores((team_owned_store(managed_id),), key) == (team_owned_store(managed_id),)
assert await can_user_read_vector_store(team_owned_store("store-a"), key) is False
assert await filter_listable_vector_stores((team_owned_store("store-a"),), key) == ()
@pytest.mark.asyncio
async def test_group_grant_for_one_managed_store_does_not_open_another_with_the_same_provider_id(
proxy: ProxyState,
) -> None:
def managed_id(unified_id: str) -> str:
return base64.urlsafe_b64encode(
f"litellm_proxy:vector_store;unified_id,{unified_id};target_model_names,gpt-4;resource_id,vs_shared".encode()
).decode()
await proxy.add_group("g1", (managed_id("granted"),))
await proxy.add_team(make_team(groups=("g1",)))
key: Final = make_key(team_id=TEAM_ID)
assert await can_user_read_vector_store(team_owned_store(managed_id("granted")), key) is True
assert await can_user_read_vector_store(team_owned_store(managed_id("other")), key) is False
assert await filter_listable_vector_stores((team_owned_store(managed_id("other")),), key) == ()
@pytest.mark.asyncio
async def test_dashboard_session_gets_the_group_stores_of_each_of_the_users_teams(proxy: ProxyState) -> None:
await proxy.add_group("g1", ("store-a",), team_ids=("team-1",))
await proxy.add_group("g2", ("store-c",), team_ids=("team-2",))
await proxy.cache.async_set_cache(
key="team_id:team-1", value=LiteLLM_TeamTableCachedObj(team_id="team-1", access_group_ids=["g1"])
)
await proxy.cache.async_set_cache(
key="team_id:team-2", value=LiteLLM_TeamTableCachedObj(team_id="team-2", access_group_ids=["g2"])
)
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_read_vector_store(stores[0], session) is True
assert await can_user_read_vector_store(stores[1], session) is False
assert await can_user_read_vector_store(stores[2], session) is 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
async def group_member(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, httpx_transport: None
) -> AsyncIterator[TestClient]:
"""A client for a team key whose team's group lists store-a. Another team owns both stores, and OpenAI is stubbed."""
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"}}]
),
)
stores: Final = {store_id: team_owned_store(store_id) for store_id in ("store-a", "store-b")}
proxy: Final = ProxyState(
monkeypatch,
store_reads=AsyncMock(
side_effect=lambda where: MagicMock(model_dump=MagicMock(return_value=stores[where["vector_store_id"]]))
),
)
await proxy.add_group("g1", ("store-a",))
await proxy.add_team(make_team(groups=("g1",)))
monkeypatch.setattr(litellm, "vector_store_registry", VectorStoreRegistry(vector_stores=list(stores.values())))
original_overrides: Final = app.dependency_overrides.copy()
app.dependency_overrides[user_api_key_auth] = lambda: make_key(team_id=TEAM_ID)
try:
yield TestClient(app, raise_server_exceptions=False)
finally:
app.dependency_overrides = original_overrides
def _send(client: TestClient, method: str, path: str, body: Mapping[str, object] | None, store_id: str) -> int:
return client.request(method, path.format(store_id), json=_body_for(body, store_id)).status_code
@pytest.mark.parametrize(("method", "path", "body"), READ_ROUTES)
def test_read_routes_serve_a_group_member_only_group_stores(
group_member: TestClient, method: str, path: str, body: Mapping[str, object] | None
) -> None:
assert _send(group_member, method, path, body, "store-a") == 200
assert _send(group_member, method, path, body, "store-b") == 403
@pytest.mark.parametrize(
("method", "path", "body"),
[
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"),
],
)
@pytest.mark.parametrize("store", ["store-a", "provider-native-store"])
def test_a_group_grant_does_not_lift_the_keys_own_list_for_changes(
group_member: TestClient, method: str, path: str, body: Mapping[str, object] | None, store: str
) -> None:
"""The key's own list allows only store-x. Its key group lists the store: one its team owns, or one LiteLLM
does not manage."""
asyncio.run(
cache_access_object(
access_group_id="key-group",
access_group_table=LiteLLM_AccessGroupTable(
access_group_id="key-group",
access_group_name="key-group",
access_vector_store_ids=(store,),
assigned_team_ids=[],
assigned_key_ids=[TOKEN],
),
user_api_key_cache=proxy_server.user_api_key_cache,
)
)
app.dependency_overrides[user_api_key_auth] = lambda: make_key(("store-x",), ("key-group",), OTHER_TEAM_ID)
assert _send(group_member, method, path, body, store) == 401
@pytest.mark.parametrize(
("method", "path", "body"),
[
pytest.param("DELETE", "/v1/vector_stores/{}", None, id="provider-delete"),
pytest.param("POST", "/v1/vector_stores/{}/files", {"file_id": "file-1"}, id="file-create"),
],
)
def test_the_owning_team_still_changes_its_store(
group_member: TestClient, method: str, path: str, body: Mapping[str, object] | None
) -> None:
app.dependency_overrides[user_api_key_auth] = lambda: make_key(team_id=OTHER_TEAM_ID)
assert _send(group_member, method, path, body, "store-a") == 200
@pytest.mark.parametrize(("method", "path", "body"), MANAGE_ROUTES)
def test_manage_routes_refuse_a_group_member(
group_member: TestClient, method: str, path: str, body: Mapping[str, object] | None
) -> None:
assert _send(group_member, method, path, body, "store-a") == 403
def test_info_reads_a_database_store_through_the_read_gate(
group_member: TestClient, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "vector_store_registry", None)
assert group_member.post("/vector_store/info", json={"vector_store_id": "store-a"}).status_code == 200
assert group_member.post("/vector_store/info", json={"vector_store_id": "store-b"}).status_code == 403
def test_info_refuses_a_group_store_the_keys_own_list_excludes(group_member: TestClient) -> None:
app.dependency_overrides[user_api_key_auth] = lambda: make_key(("store-x",), team_id=TEAM_ID)
assert group_member.post("/vector_store/info", json={"vector_store_id": "store-a"}).status_code == 403
def test_rag_ingest_refuses_a_group_member(group_member: TestClient) -> None:
response: Final = group_member.post(
"/v1/rag/ingest",
json={
"ingest_options": {"vector_store": {"custom_llm_provider": "openai", "vector_store_id": "store-a"}},
"file_url": "https://example.com/a.txt",
},
)
assert response.status_code == 403

6
tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py
View file

@ -26,6 +26,7 @@ from litellm.proxy.vector_store_endpoints.management_endpoints import (
new_vector_store,
)
from litellm.proxy.vector_store_endpoints.utils import (
assert_user_can_access_vector_store,
check_vector_store_permission,
is_allowed_to_call_vector_store_endpoint,
is_allowed_to_call_vector_store_files_endpoint,
@ -659,7 +660,7 @@ async def test_update_request_data_with_litellm_managed_vector_store_registry():
# Test with vector store registry
with patch.object(litellm, "vector_store_registry", mock_registry):
result = await _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id
data=data, vector_store_id=vector_store_id, authorize=assert_user_can_access_vector_store
)
# Verify the data was updated correctly
@ -681,7 +682,7 @@ async def test_update_request_data_with_litellm_managed_vector_store_registry():
):
original_data = {"existing_key": "existing_value"}
result = await _update_request_data_with_litellm_managed_vector_store_registry(
data=original_data, vector_store_id=vector_store_id
data=original_data, vector_store_id=vector_store_id, authorize=assert_user_can_access_vector_store
)
# Verify data remains unchanged when no registry
@ -706,6 +707,7 @@ async def test_managed_vector_store_keeps_embedding_reference_and_explicit_confi
result = await _update_request_data_with_litellm_managed_vector_store_registry(
data={},
vector_store_id="test_store",
authorize=assert_user_can_access_vector_store,
)
assert result["litellm_embedding_model"] == "team-embedding-alias"

2
tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py
View file

@ -194,7 +194,7 @@ async def test_vector_store_file_list_resolves_managed_ids_and_cursors():
with (
patch(
"litellm.proxy.vector_store_files_endpoints.endpoints.assert_user_can_access_vector_store_id",
"litellm.proxy.vector_store_files_endpoints.endpoints.assert_user_can_read_vector_store_id",
new=AsyncMock(return_value=None),
),
patch("litellm.proxy.proxy_server.llm_router", llm_router),

9
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -26164,6 +26164,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 */
@ -26218,6 +26220,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 */
@ -26270,6 +26277,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 */