diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1ed4f25dc71..882e98ec097 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -46,6 +46,11 @@ from litellm.constants import ( from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.litellm_core_utils.tool_fields import DropPolicy +from litellm.llms.base_llm.managed_resources.utils import ( + extract_provider_resource_id_from_unified_id, + is_base64_encoded_unified_id, +) from litellm.models.project import LiteLLM_ProjectTable from litellm.proxy._types import ( MODEL_LIST_ROUTES, @@ -165,6 +170,7 @@ from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.types.router import DeploymentTypedDict from litellm.utils import get_utc_datetime +from litellm.vector_stores.vector_store_registry import listed_vector_store_ids, vector_store_ids_in_tool from .auth_checks_organization import ( add_team_org_context_to_request_body, @@ -1514,6 +1520,7 @@ async def common_checks( valid_token=valid_token, user_object=user_object, deny_by_default=_vector_store_deny_by_default(_typed_request_body(general_settings)), + llm_router=llm_router, ) # 12. [OPTIONAL] Tool allowlist - key/team allowed_tools (no DB in hot path) @@ -7985,6 +7992,57 @@ def _strict_requested_vector_store_ids(request_body: Mapping[str, object]) -> tu ) +@dataclass(frozen=True, slots=True) +class _UnreadableVectorStoreIds: + pass + + +def _top_level_vector_store_ids( + vector_store_ids: object, drop_policies: Callable[[], tuple[DropPolicy, ...]] +) -> tuple[str, ...] | _UnreadableVectorStoreIds: + listed: Final = listed_vector_store_ids(vector_store_ids) + if vector_store_ids is None or vector_store_ids == [*listed]: + return listed + if any(policy.drops("vector_store_ids") for policy in drop_policies()): + return listed + if isinstance(vector_store_ids, str): + return (vector_store_ids,) + return _UnreadableVectorStoreIds() + + +def _legacy_requested_vector_store_ids( + request_body: Mapping[str, object], drop_policies: Callable[[], tuple[DropPolicy, ...]] +) -> tuple[str, ...] | _UnreadableVectorStoreIds: + """ + Reads tools the way VectorStoreRegistry does and leaves a malformed tool to the proxy's tool + validation. A top-level vector_store_ids string names one store, since a query parameter always + arrives as a string. Any other value that is not a list of strings names no readable stores unless a + drop policy the request can resolve to drops it, and then only the string entries of a list name + stores, as in registry retrieval. + """ + top_level: Final = _top_level_vector_store_ids(request_body.get("vector_store_ids"), drop_policies) + if isinstance(top_level, _UnreadableVectorStoreIds): + return top_level + return tuple( + chain( + top_level, + chain.from_iterable(vector_store_ids_in_tool(tool) for tool in _tools(request_body.get("tools"))), + ) + ) + + +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) + + +def _request_drop_policies( + request_body: Mapping[str, object], llm_router: Router | None, valid_token: UserAPIKeyAuth | None +) -> tuple[DropPolicy, ...]: + from litellm.proxy.common_request_processing import counted_drop_policies + + return counted_drop_policies(request_body, llm_router, valid_token.team_id if valid_token is not None else None) + + def _require_vector_store_grant( object_type: GrantLayer, vector_store_ids_to_run: Sequence[str], @@ -8032,6 +8090,7 @@ async def vector_store_access_check( *, user_object: LiteLLM_UserTable | None = None, deny_by_default: bool = False, + llm_router: Router | None = None, ): """ Checks whether the caller may use every vector store the request names. @@ -8040,11 +8099,19 @@ async def vector_store_access_check( `retrieval_config.vector_store_id`. Grants come from each identity's `object_permission.vector_stores`. - With `deny_by_default=False` (legacy), a key or team only restricts access when its list is - nonempty, and stores are only read from the request when a vector store registry is loaded. + 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 + key or team has an allowlist and the value is not a string, which it checks as one store id. A + LiteLLM-managed store id is compared by the provider store id it decodes to, on both sides. Legacy mode skips that value instead when the request, `litellm_settings`, the router defaults or a deployment of the requested + model drops it (`drop_params`, or `additional_drop_params` naming it), the same settings the proxy's + tool validation reads, and reads tools the way the registry does, where strict mode answers 400 on a + malformed one. - With `deny_by_default=True` (`general_settings.vector_store_deny_by_default`), stores are read - even without a registry, and a missing record, `null` or `[]` grants nothing: + With `deny_by_default=False` (legacy), a key or team only restricts access when its list is + nonempty. + + With `deny_by_default=True` (`general_settings.vector_store_deny_by_default`), a missing record, + `null` or `[]` grants nothing: - virtual key: the key must grant every store, and so must its team when it has one, even if the team failed to load @@ -8056,8 +8123,8 @@ async def vector_store_access_check( rescue nor restrict a key or team request. Raises ProxyException (401, `{key,team,user}_vector_store_access_denied`) on the first identity - that does not grant a requested store, and with the flag on, ProxyException (400, - `invalid_request_error`) when a `vector_store_ids` field is not a list of strings. + that does not grant a requested store, and ProxyException (400, `invalid_request_error`) on a + malformed `vector_store_ids` the request does not drop and, in legacy mode, an allowlist applies to. """ from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -8068,21 +8135,23 @@ async def vector_store_access_check( verbose_proxy_logger.debug("Prisma client not found, skipping vector store access check") return True - registry_ids: Final = ( - _strict_requested_vector_store_ids(_typed_request_body(request_body)) + typed_request_body: Final = _typed_request_body(request_body) + requested_ids: Final = ( + _strict_requested_vector_store_ids(typed_request_body) if deny_by_default - else ( - litellm.vector_store_registry.get_vector_store_ids_to_run( - non_default_params=request_body, tools=request_body.get("tools", None) - ) - if litellm.vector_store_registry is not None - else None + else _legacy_requested_vector_store_ids( + typed_request_body, lambda: _request_drop_policies(typed_request_body, llm_router, valid_token) ) - or () ) - rag_vector_store_id: Final = _get_rag_query_vector_store_id(_typed_request_body(request_body)) + if isinstance(requested_ids, _UnreadableVectorStoreIds): + if _restricts_vector_stores( + await _identity_grants(valid_token, team_object, None, prisma_client, user_api_key_cache) + ): + raise _malformed_vector_store_ids() + return True + rag_vector_store_id: Final = _get_rag_query_vector_store_id(typed_request_body) rag_ids: Final = (rag_vector_store_id,) if rag_vector_store_id is not None else () - vector_store_ids_to_run: Final = tuple(dict.fromkeys((*registry_ids, *rag_ids))) + vector_store_ids_to_run: Final = tuple(dict.fromkeys((*requested_ids, *rag_ids))) if not vector_store_ids_to_run: return True @@ -8109,6 +8178,12 @@ async def vector_store_access_check( return True +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 + + def _can_object_call_vector_stores( object_type: Literal["key", "team", "org", "user"], vector_store_ids_to_run: Sequence[str], @@ -8127,8 +8202,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)) for vector_store_id in vector_store_ids_to_run: - if vector_store_id not in object_permissions.vector_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), diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 792efa480ec..6877d6c0118 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -724,7 +724,7 @@ def _router_defaults(llm_router: Router | None) -> Mapping[str, object]: return {name: value for name, value in _REQUEST_FIELDS.validate_python(configured).items() if value is not None} -def _counted_drop_policies( +def counted_drop_policies( data: Mapping[str, object], llm_router: Router | None, team_id: str | None ) -> tuple[DropPolicy, ...]: requested: Final = {**_router_defaults(llm_router), **data} @@ -740,7 +740,7 @@ def _counted_drop_policies( def resolve_counted_tool_fields( data: Mapping[str, object], llm_router: Router | None, team_id: str | None ) -> Mapping[str, object] | ProxyException: - resolved: Final = resolve_tool_fields(data, lambda: _counted_drop_policies(data, llm_router, team_id)) + resolved: Final = resolve_tool_fields(data, lambda: counted_drop_policies(data, llm_router, team_id)) if isinstance(resolved, MalformedField): return ProxyException( message=resolved.message, diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index b93d9874dbc..598b56f4a79 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -36,7 +36,7 @@ else: _DB_ROW: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) _ITEMS: Final = TypeAdapter(Sequence[object], config=ConfigDict(strict=True)) _TOOL: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(strict=True)) -_STORE_ID_ITEMS: Final = TypeAdapter(list[object], config=ConfigDict(strict=True)) +_STORE_ID_ITEMS: Final = TypeAdapter(list[object] | tuple[object, ...], config=ConfigDict(strict=True)) def _named_store_ids(store_ids: object) -> tuple[str, ...]: @@ -49,7 +49,11 @@ def _named_store_ids(store_ids: object) -> tuple[str, ...]: return tuple(store_id for store_id in items if isinstance(store_id, str)) -def _tool_vector_store_ids(tool: object) -> tuple[str, ...]: +def listed_vector_store_ids(store_ids: object) -> tuple[str, ...]: + return _named_store_ids(store_ids) if isinstance(store_ids, (list, tuple)) else () + + +def vector_store_ids_in_tool(tool: object) -> tuple[str, ...]: try: return _named_store_ids(_TOOL.validate_python(tool).get("vector_store_ids")) except ValidationError: @@ -195,7 +199,7 @@ class VectorStoreRegistry: tools_to_remove: Final[list[int]] = [] for i, tool in enumerate(tools): - tool_vector_store_ids = _tool_vector_store_ids(tool) + tool_vector_store_ids = vector_store_ids_in_tool(tool) if not tool_vector_store_ids: continue @@ -305,7 +309,7 @@ class VectorStoreRegistry: List of vector stores with tool parameters merged into litellm_params """ # Pop vector_store_ids from params - vector_store_ids: Final[list[str]] = non_default_params.pop("vector_store_ids", None) or [] + vector_store_ids: Final = [*listed_vector_store_ids(non_default_params.pop("vector_store_ids", None))] # Extract params from tools and collect IDs params_by_id = self.get_and_pop_recognised_vector_store_tools(tools=tools, vector_store_ids=vector_store_ids) @@ -353,7 +357,7 @@ class VectorStoreRegistry: List of vector stores with tool parameters merged into litellm_params """ # Pop vector_store_ids from params - vector_store_ids: Final[list[str]] = non_default_params.pop("vector_store_ids", None) or [] + vector_store_ids: Final = [*listed_vector_store_ids(non_default_params.pop("vector_store_ids", None))] # Extract params from tools and collect IDs params_by_id = self.get_and_pop_recognised_vector_store_tools(tools=tools, vector_store_ids=vector_store_ids) @@ -423,7 +427,7 @@ class VectorStoreRegistry: validated_tools: Final = _ITEMS.validate_python(tools) except ValidationError: return vector_store_ids - return [*vector_store_ids, *chain.from_iterable(_tool_vector_store_ids(tool) for tool in validated_tools)] + return [*vector_store_ids, *chain.from_iterable(vector_store_ids_in_tool(tool) for tool in validated_tools)] def load_vector_stores_from_config(self, vector_stores_config: list[dict]): """ diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index effdd9725c2..0bdc5f1910a 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -101,6 +101,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( tag_registry_cache_key, ) from litellm.utils import get_utc_datetime +from litellm.llms.base_llm.managed_resources.utils import encode_unified_id, generate_unified_id_string from litellm.vector_stores.vector_store_registry import VectorStoreRegistry from tests._master_key import MASTER_KEY @@ -1551,37 +1552,24 @@ async def test_get_team_db_check_does_not_call_new_team_if_exists(mock_new_team, @pytest.mark.asyncio -@pytest.mark.parametrize( - "prisma_client,vector_store_registry,expected_result", - [ - (None, MagicMock(), True), # No prisma client - (MagicMock(), None, True), # No vector store registry - (MagicMock(), MagicMock(), True), # No vector stores to run - ], -) -async def test_vector_store_access_check_early_returns(prisma_client, vector_store_registry, expected_result): +@pytest.mark.parametrize("prisma_client", [None, MagicMock()], ids=["no-prisma-client", "no-vector-stores-to-run"]) +async def test_vector_store_access_check_early_returns(prisma_client): """Test vector_store_access_check returns True for early exit conditions""" request_body = {"messages": [{"role": "user", "content": "test"}]} - if vector_store_registry: - vector_store_registry.get_vector_store_ids_to_run.return_value = None - - with ( - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), - patch("litellm.vector_store_registry", vector_store_registry), - ): + with patch("litellm.proxy.proxy_server.prisma_client", prisma_client): result = await vector_store_access_check( request_body=request_body, team_object=None, valid_token=None, ) - assert result == expected_result + assert result is True @pytest.mark.asyncio async def test_vector_store_access_check_skips_db_lookup_when_no_vector_stores_requested(): - """Registry returns [] (not None) for plain requests; no object permission DB lookup should happen.""" + """A plain request names no vector store, so no object permission DB lookup should happen.""" valid_token = UserAPIKeyAuth(token="test-token", object_permission_id="perm-123") team_object = MagicMock() team_object.object_permission_id = "team-permission" @@ -1590,16 +1578,8 @@ async def test_vector_store_access_check_skips_db_lookup_when_no_vector_stores_r find_unique = AsyncMock() mock_prisma_client.db.litellm_objectpermissiontable.find_unique = find_unique - mock_vector_store_registry = MagicMock() - mock_vector_store_registry.get_vector_store_ids_to_run.return_value = [] - - with ( - patch( # test-quality-ok: production auth reads these module globals; no dependency injection seam exists - "litellm.proxy.proxy_server.prisma_client", mock_prisma_client - ), - patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists - "litellm.vector_store_registry", mock_vector_store_registry - ), + with patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client ): result = await vector_store_access_check( request_body={"messages": [{"role": "user", "content": "test"}]}, @@ -1668,102 +1648,6 @@ def test_can_object_call_vector_stores_scenarios(object_permissions, vector_stor assert result is True -@pytest.mark.asyncio -async def test_vector_store_access_check_with_permissions(): - """Test vector_store_access_check with actual permission checking""" - request_body = {"tools": [{"type": "function", "function": {"name": "test"}}]} - - # Test with valid token that has access - valid_token = UserAPIKeyAuth( - token="test-token", - object_permission_id="perm-123", - models=["gpt-4"], - max_budget=100.0, - ) - - mock_prisma_client = MagicMock() - mock_permissions = LiteLLM_ObjectPermissionTable( - object_permission_id="perm-123", vector_stores=["store-1", "store-2"] - ) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=mock_permissions) - - mock_vector_store_registry = MagicMock() - mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-1"] - - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), - patch("litellm.vector_store_registry", mock_vector_store_registry), - ): - result = await vector_store_access_check( - request_body=request_body, - team_object=None, - valid_token=valid_token, - ) - - assert result is True - - # Test with denied access - mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["store-3"] - - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), - patch("litellm.vector_store_registry", mock_vector_store_registry), - ): - with pytest.raises(ProxyException) as exc_info: - await vector_store_access_check( - request_body=request_body, - team_object=None, - valid_token=valid_token, - ) - - assert exc_info.value.type == ProxyErrorTypes.key_vector_store_access_denied - - -@pytest.mark.asyncio -async def test_vector_store_access_check_with_team_permissions(): - """Ensure teams restricted to specific vector stores cannot access others.""" - request_body = {} - valid_token = UserAPIKeyAuth(token="team-test-token", object_permission_id=None) - - team_object = LiteLLM_TeamTable(team_id="team-1", object_permission_id="team-permission") - - mock_prisma_client = MagicMock() - team_permissions = LiteLLM_ObjectPermissionTable( - object_permission_id="team-permission", vector_stores=["team-store-allowed"] - ) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=team_permissions) - - mock_vector_store_registry = MagicMock() - mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["team-store-allowed"] - - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), - patch("litellm.vector_store_registry", mock_vector_store_registry), - ): - result = await vector_store_access_check( - request_body=request_body, - team_object=team_object, - valid_token=valid_token, - ) - - assert result is True - - mock_vector_store_registry.get_vector_store_ids_to_run.return_value = ["team-store-denied"] - - with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), - patch("litellm.vector_store_registry", mock_vector_store_registry), - ): - with pytest.raises(ProxyException) as exc_info: - await vector_store_access_check( - request_body=request_body, - team_object=team_object, - valid_token=valid_token, - ) - - assert exc_info.value.type == ProxyErrorTypes.team_vector_store_access_denied - - @pytest.mark.asyncio @pytest.mark.parametrize( "requested_vector_store_id,expected_error_type", @@ -1823,6 +1707,188 @@ async def test_vector_store_access_check_enforces_team_allowlist_for_rag_query( assert exc_info.value.type == expected_error_type +def _file_search_tool_request(vector_store_id: str) -> dict[str, object]: + return { + "model": "gpt-5.4-mini", + "input": "what is in this store?", + "tools": [{"type": "file_search", "vector_store_ids": [vector_store_id]}], + } + + +def _file_search_tool_string_id_request(vector_store_id: str) -> dict[str, object]: + return { + "model": "gpt-5.4-mini", + "input": "what is in this store?", + "tools": [{"type": "file_search", "vector_store_ids": vector_store_id}], + } + + +def _vector_store_ids_request(vector_store_id: str) -> dict[str, object]: + return { + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "what is in this store?"}], + "vector_store_ids": [vector_store_id], + } + + +def _drop_params_request(vector_store_ids: object) -> dict[str, object]: + return { + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "what is in this store?"}], + "drop_params": True, + "vector_store_ids": vector_store_ids, + } + + +def _allowlist(object_permission_id: str, vector_stores: list[str] | None) -> LiteLLM_ObjectPermissionTable | None: + if vector_stores is None: + return None + return LiteLLM_ObjectPermissionTable(object_permission_id=object_permission_id, vector_stores=vector_stores) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("key_vector_stores", "team_vector_stores", "request_body", "denied_by"), + [ + pytest.param( + None, + ["vs_allowed"], + _file_search_tool_request("vs_other_team"), + ProxyErrorTypes.team_vector_store_access_denied, + id="team-file-search-tool-other-store", + ), + pytest.param(None, ["vs_allowed"], _file_search_tool_request("vs_allowed"), None, id="team-file-search-tool"), + pytest.param( + ["vs_allowed"], + None, + _vector_store_ids_request("vs_other_team"), + ProxyErrorTypes.key_vector_store_access_denied, + id="key-vector-store-ids-other-store", + ), + pytest.param(["vs_allowed"], None, _vector_store_ids_request("vs_allowed"), None, id="key-vector-store-ids"), + pytest.param( + None, + ["vs_allowed"], + _file_search_tool_string_id_request("vs_other_team"), + ProxyErrorTypes.team_vector_store_access_denied, + id="team-file-search-tool-string-id-other-store", + ), + pytest.param( + None, + None, + _file_search_tool_string_id_request("vs_other_team"), + None, + id="unrestricted-file-search-tool-string-id", + ), + pytest.param( + None, + ["vs_allowed"], + _drop_params_request(["vs_other_team", {"id": "vs_allowed"}]), + ProxyErrorTypes.team_vector_store_access_denied, + id="team-drop-params-list-with-an-object-other-store", + ), + pytest.param( + None, + ["vs_allowed"], + { + **_file_search_tool_request("vs_other_team"), + "tools": [{"type": "file_search", "vector_store_ids": ["vs_other_team", 1]}], + }, + ProxyErrorTypes.team_vector_store_access_denied, + id="team-file-search-tool-with-an-int-other-store", + ), + pytest.param( + None, + ["vs_allowed"], + { + **_vector_store_ids_request("vs_other_team"), + "additional_drop_params": ["vector_store_ids"], + "vector_store_ids": {"vs_other_team": True}, + }, + None, + id="team-additional-drop-params-object-names-no-store", + ), + pytest.param( + None, + ["vs_allowed"], + _drop_params_request({"vs_other_team": True}), + None, + id="team-drop-params-object-names-no-store", + ), + ], +) +@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"]) +async def test_vector_store_access_check_enforces_allowlists_on_the_stores_the_request_names( + key_vector_stores: list[str] | None, + team_vector_stores: list[str] | None, + request_body: dict[str, object], + denied_by: ProxyErrorTypes | None, + vector_store_registry: VectorStoreRegistry | None, +) -> None: + valid_token: Final = UserAPIKeyAuth( + token="team-a-key", + team_id="team-a", + object_permission_id="key-permission", + object_permission=_allowlist("key-permission", key_vector_stores), + ) + team_object: Final = LiteLLM_TeamTable( + team_id="team-a", + object_permission_id="team-permission", + object_permission=_allowlist("team-permission", team_vector_stores), + ) + with ( + patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists + "litellm.vector_store_registry", vector_store_registry + ), + ): + await _assert_rag_query_outcome( + vector_store_access_check(request_body=request_body, team_object=team_object, valid_token=valid_token), + denied_by, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"]) +async def test_litellm_settings_drop_params_skips_a_malformed_vector_store_ids( + vector_store_registry: VectorStoreRegistry | None, +) -> None: + with patch("litellm.drop_params", True): # test-quality-ok: litellm_settings.drop_params is this module global + assert await _common_checks_for_rag_query( + {}, + _virtual_key(object_permission_id="key-permission"), + SimpleNamespace(vector_stores=["KBSTOREA"]), + request_body={ + "model": "gpt-4o-mini", + "input": "what is in this KB?", + "vector_store_ids": {"KBSTOREB": True}, + }, + vector_store_registry=vector_store_registry, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment_params", + [{"drop_params": True}, {"additional_drop_params": ["vector_store_ids"]}], + ids=["drop-params", "additional-drop-params"], +) +@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"]) +async def test_a_deployment_that_drops_vector_store_ids_skips_a_malformed_one( + deployment_params: Mapping[str, object], vector_store_registry: VectorStoreRegistry | None +) -> None: + assert await _common_checks_for_rag_query( + {}, + _virtual_key(object_permission_id="key-permission"), + SimpleNamespace(vector_stores=["KBSTOREA"]), + request_body={"model": "gpt-4o-mini", "input": "what is in this KB?", "vector_store_ids": ["KBSTOREA", 1]}, + vector_store_registry=vector_store_registry, + llm_router=_gpt_4o_mini_router(deployment_params), + ) + + _KEY_DENIED: Final = ProxyErrorTypes.key_vector_store_access_denied _TEAM_DENIED: Final = ProxyErrorTypes.team_vector_store_access_denied _USER_DENIED: Final = ProxyErrorTypes.user_vector_store_access_denied @@ -1866,6 +1932,7 @@ async def _common_checks_for_rag_query( request_body: Mapping[str, object] | None = None, user_api_key_cache: UserApiKeyCache | None = None, vector_store_registry: VectorStoreRegistry | None = None, + llm_router: "Router | None" = None, ) -> bool: permissions: Final = { "key-permission": _permission_row("key-permission", key_permission), @@ -1906,13 +1973,24 @@ async def _common_checks_for_rag_query( global_proxy_spend=None, general_settings=dict(general_settings), route="/v1/rag/query", - llm_router=None, + llm_router=llm_router, proxy_logging_obj=MagicMock(), valid_token=valid_token, request=MagicMock(spec=Request), ) +def _gpt_4o_mini_router(litellm_params: Mapping[str, object]) -> "Router": + return litellm.Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", **litellm_params}, + } + ] + ) + + async def _assert_rag_query_outcome(request: Awaitable[bool], denied_by: ProxyErrorTypes | None) -> None: if denied_by is None: assert await request is True @@ -2343,6 +2421,175 @@ async def test_deny_by_default_rejects_malformed_vector_store_ids_as_bad_request ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_body", + [ + {"vector_store_ids": {"KBSTOREB": True}}, + {"vector_store_ids": ["KBSTOREA", {"id": "KBSTOREB"}]}, + {"vector_store_ids": 7}, + ], + ids=["object-top-level-ids", "object-top-level-id", "int-top-level-ids"], +) +@pytest.mark.parametrize( + "general_settings", [{"vector_store_deny_by_default": True}, {}], ids=["deny-by-default", "legacy"] +) +@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"]) +async def test_malformed_vector_store_ids_are_rejected_as_bad_request( + request_body: Mapping[str, object], + general_settings: Mapping[str, object], + vector_store_registry: VectorStoreRegistry | None, +): + with pytest.raises(ProxyException) as exc_info: + await _common_checks_for_rag_query( + general_settings, + _virtual_key(object_permission_id="key-permission"), + SimpleNamespace(vector_stores=["KBSTOREA"]), + request_body={"model": "gpt-4o-mini", "input": "what is in this KB?", **request_body}, + vector_store_registry=vector_store_registry, + llm_router=_gpt_4o_mini_router({}), + ) + assert (exc_info.value.type, exc_info.value.param, exc_info.value.code) == ( + "invalid_request_error", + "vector_store_ids", + "400", + ) + + +_MALFORMED_TOP_LEVEL_IDS: Final = pytest.mark.parametrize( + "vector_store_ids", + [{"KBSTOREB": True}, ["KBSTOREA", {"id": "KBSTOREB"}], 7], + ids=["object-top-level-ids", "object-top-level-id", "int-top-level-ids"], +) + + +async def _team_key_query_with_top_level_ids( + key_vector_stores: list[str] | None, + team_vector_stores: list[str] | None, + vector_store_ids: object, + vector_store_registry: VectorStoreRegistry | None, +) -> bool: + return await _common_checks_for_rag_query( + {}, + _virtual_key(object_permission_id=None if key_vector_stores is None else "key-permission", team_id="team-1"), + None if key_vector_stores is None else SimpleNamespace(vector_stores=key_vector_stores), + LiteLLM_TeamTable( + team_id="team-1", object_permission_id=None if team_vector_stores is None else "team-permission" + ), + None if team_vector_stores is None else SimpleNamespace(vector_stores=team_vector_stores), + request_body={"model": "gpt-4o-mini", "input": "what is in this KB?", "vector_store_ids": vector_store_ids}, + vector_store_registry=vector_store_registry, + llm_router=_gpt_4o_mini_router({}), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "vector_store_ids", + ["KBSTOREB", {"KBSTOREB": True}, ["KBSTOREA", {"id": "KBSTOREB"}]], + ids=["string-top-level-ids", "object-top-level-ids", "object-top-level-id"], +) +@pytest.mark.parametrize( + ("key_vector_stores", "team_vector_stores"), + [(None, None), ([], [])], + ids=["no-permission-rows", "empty-allowlists"], +) +@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"]) +async def test_legacy_mode_passes_malformed_vector_store_ids_through_without_an_allowlist( + vector_store_ids: object, + key_vector_stores: list[str] | None, + team_vector_stores: list[str] | None, + vector_store_registry: VectorStoreRegistry | None, +): + assert ( + await _team_key_query_with_top_level_ids( + key_vector_stores, team_vector_stores, vector_store_ids, vector_store_registry + ) + is True + ) + + +@pytest.mark.asyncio +@_MALFORMED_TOP_LEVEL_IDS +@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"]) +async def test_legacy_mode_rejects_malformed_vector_store_ids_under_a_team_allowlist( + vector_store_ids: object, vector_store_registry: VectorStoreRegistry | None +): + with pytest.raises(ProxyException) as exc_info: + await _team_key_query_with_top_level_ids(None, ["KBSTOREA"], vector_store_ids, vector_store_registry) + assert (exc_info.value.type, exc_info.value.param, exc_info.value.code) == ( + "invalid_request_error", + "vector_store_ids", + "400", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("vector_store_ids", "denied_by"), [("KBSTOREA", None), ("KBSTOREB", _TEAM_DENIED)]) +@pytest.mark.parametrize("vector_store_registry", [VectorStoreRegistry(), None], ids=["registry", "no-registry"]) +async def test_legacy_mode_checks_a_top_level_vector_store_ids_string_as_one_store( + vector_store_ids: str, denied_by: ProxyErrorTypes | None, vector_store_registry: VectorStoreRegistry | None +): + await _assert_rag_query_outcome( + _team_key_query_with_top_level_ids(None, ["KBSTOREA"], vector_store_ids, vector_store_registry), denied_by + ) + + +def _managed_vector_store_id(unified_uuid: str, provider_vector_store_id: str) -> str: + return encode_unified_id( + generate_unified_id_string( + resource_type="vector_store", + unified_uuid=unified_uuid, + target_model_names=["gpt-4o-mini"], + provider_resource_id=provider_vector_store_id, + model_id="deployment-1", + ) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_vector_stores", "requested_id", "denied_by"), + [ + (["KBSTOREA"], _managed_vector_store_id("uuid-a", "KBSTOREA"), None), + ([_managed_vector_store_id("uuid-a", "KBSTOREA")], "KBSTOREA", None), + ([_managed_vector_store_id("uuid-a", "KBSTOREA")], _managed_vector_store_id("uuid-a2", "KBSTOREA"), None), + (["KBSTOREA"], _managed_vector_store_id("uuid-b", "KBSTOREB"), _TEAM_DENIED), + (["uuid-a"], _managed_vector_store_id("uuid-a", "KBSTOREB"), _TEAM_DENIED), + ([_managed_vector_store_id("uuid-a", "KBSTOREA")], "KBSTOREB", _TEAM_DENIED), + (["KBSTOREA"], "litellm_proxy:vector_store;unified_id,uuid-a;resource_id,KBSTOREA", _TEAM_DENIED), + ], + ids=[ + "provider-id-allows-its-managed-id", + "managed-id-allows-its-provider-id", + "managed-ids-for-one-provider-store", + "managed-id-for-another-store", + "managed-uuid-cannot-vouch-for-another-store", + "managed-id-for-another-provider-store", + "undecoded-managed-id-is-its-own-store", + ], +) +@pytest.mark.parametrize("deny_by_default", [True, False], ids=["strict", "default"]) +async def test_allowlist_compares_a_managed_vector_store_id_by_the_provider_store_it_names( + team_vector_stores: list[str], requested_id: str, denied_by: ProxyErrorTypes | None, deny_by_default: bool +): + await _assert_rag_query_outcome( + _common_checks_for_rag_query( + {"vector_store_deny_by_default": deny_by_default}, + _virtual_key(object_permission_id="key-permission", team_id="team-1"), + SimpleNamespace(vector_stores=[requested_id]), + LiteLLM_TeamTable(team_id="team-1", object_permission_id="team-permission"), + SimpleNamespace(vector_stores=team_vector_stores), + request_body={ + "model": "gpt-4o-mini", + "input": "what is in this KB?", + "tools": [{"type": "file_search", "vector_store_ids": [requested_id]}], + }, + ), + denied_by, + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("request_body", "denied_by"), diff --git a/tests/unit/vector_stores/test_vector_store_registry.py b/tests/unit/vector_stores/test_vector_store_registry.py index 1a158771f29..a6e28f258cb 100644 --- a/tests/unit/vector_stores/test_vector_store_registry.py +++ b/tests/unit/vector_stores/test_vector_store_registry.py @@ -519,6 +519,40 @@ def test_vector_store_tools_pop_exactly_the_store_ids_the_registry_reports_to_ru assert ran == checked +@pytest.mark.asyncio +async def test_vector_stores_to_run_ignore_a_vector_store_ids_object() -> None: + now: Final = datetime.now(timezone.utc) + registry: Final = VectorStoreRegistry( + [LiteLLM_ManagedVectorStore(vector_store_id="vs_registered", created_at=now, updated_at=now)] + ) + + ran: Final = ( + registry.pop_vector_stores_to_run(non_default_params={"vector_store_ids": {"vs_registered": True}}), + await registry.pop_vector_stores_to_run_with_db_fallback( + non_default_params={"vector_store_ids": {"vs_registered": True}} + ), + ) + + assert ran == ([], []) + + +@pytest.mark.asyncio +async def test_vector_stores_to_run_read_a_tuple_of_vector_store_ids() -> None: + now: Final = datetime.now(timezone.utc) + registry: Final = VectorStoreRegistry( + [LiteLLM_ManagedVectorStore(vector_store_id="vs_registered", created_at=now, updated_at=now)] + ) + + ran: Final = ( + registry.pop_vector_stores_to_run(non_default_params={"vector_store_ids": ("vs_registered",)}), + await registry.pop_vector_stores_to_run_with_db_fallback( + non_default_params={"vector_store_ids": ("vs_registered",)} + ), + ) + + assert [[store["vector_store_id"] for store in stores] for stores in ran] == [["vs_registered"], ["vs_registered"]] + + class TestRegistriesDoNotShareDefaultState: def test_adding_to_one_registry_leaves_the_next_one_empty(self): VectorStoreRegistry().add_vector_store_to_registry(