From b97ea585b2cdb31aea1dde8b6c062b984e0a19a5 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 27 Nov 2025 03:49:30 +0530 Subject: [PATCH] Add method for extracting vector store ids from path params (#16566) * Add method for extracting vector store ids from path params * Add vector id handling from path * Move method to utils --- litellm/proxy/auth/user_api_key_auth.py | 7 ++ .../proxy/common_utils/http_parsing_utils.py | 75 +++++++++++++++++++ .../vector_stores/vector_store_registry.py | 6 +- .../proxy/auth/test_auth_checks.py | 49 ++++++++++++ 4 files changed, 135 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ba5747a43a..5552089d46 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -8,6 +8,7 @@ Returns a UserAPIKeyAuth object if the API key is valid """ import asyncio +import re import secrets from datetime import datetime, timezone from typing import List, Optional, Tuple, cast @@ -54,6 +55,7 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, + populate_request_with_path_params, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -1202,6 +1204,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) + + @tracer.wrap() async def user_api_key_auth( request: Request, @@ -1223,6 +1227,9 @@ async def user_api_key_auth( """ request_data = await _read_request_body(request=request) + request_data = populate_request_with_path_params( + request_data=request_data, request=request + ) route: str = get_request_route(request=request) ## CHECK IF ROUTE IS ALLOWED diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 8d8d176e23..59b3ec20b4 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -324,3 +324,78 @@ def get_tags_from_request_body(request_body: dict) -> List[str]: ###################################### return [tag for tag in combined_tags if isinstance(tag, str)] + +def populate_request_with_path_params( + request_data: dict, request: Request +) -> dict: + """ + Copy FastAPI path params into the request payload so downstream checks + (e.g. vector store RBAC) see them the same way as body params. + + Since path_params may not be available during dependency injection, + we parse the URL path directly for known patterns. + + Args: + request_data: The request data dictionary to populate + request: The FastAPI Request object + + Returns: + dict: Updated request_data with path parameters added + """ + # Try to get path_params if available (sometimes populated by FastAPI) + path_params = getattr(request, "path_params", None) + if isinstance(path_params, dict) and path_params: + for key, value in path_params.items(): + if key == "vector_store_id": + request_data.setdefault("vector_store_id", value) + existing_ids = request_data.get("vector_store_ids") + if isinstance(existing_ids, list): + if value not in existing_ids: + existing_ids.append(value) + else: + request_data["vector_store_ids"] = [value] + continue + request_data.setdefault(key, value) + verbose_proxy_logger.debug( + f"populate_request_with_path_params: Found path_params, vector_store_ids={request_data.get('vector_store_ids')}" + ) + return request_data + + # Fallback: parse the URL path directly to extract vector_store_id + _add_vector_store_id_from_path(request_data=request_data, request=request) + + return request_data + + +def _add_vector_store_id_from_path(request_data: dict, request: Request) -> None: + """ + Parse the request path to find /vector_stores/{vector_store_id}/... segments. + + When found, ensure both vector_store_id and vector_store_ids are populated. + + Args: + request_data: The request data dictionary to populate + request: The FastAPI Request object + """ + path = request.url.path + vector_store_match = re.search(r"/vector_stores/([^/]+)/", path) + if vector_store_match: + vector_store_id = vector_store_match.group(1) + verbose_proxy_logger.debug( + f"populate_request_with_path_params: Extracted vector_store_id={vector_store_id} from path={path}" + ) + request_data.setdefault("vector_store_id", vector_store_id) + existing_ids = request_data.get("vector_store_ids") + if isinstance(existing_ids, list): + if vector_store_id not in existing_ids: + existing_ids.append(vector_store_id) + else: + request_data["vector_store_ids"] = [vector_store_id] + verbose_proxy_logger.debug( + f"populate_request_with_path_params: Updated request_data with vector_store_ids={request_data.get('vector_store_ids')}" + ) + else: + verbose_proxy_logger.debug( + f"populate_request_with_path_params: No vector_store_id present in path={path}" + ) + diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 9578c8e349..78c8d7cf2e 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -136,14 +136,16 @@ class VectorStoreRegistry: vector_store_ids: List[str] = [] # 1. check if vector_store_ids is provided in the non_default_params - vector_store_ids = non_default_params.get("vector_store_ids", None) or [] + vector_store_ids_param = non_default_params.get("vector_store_ids") + if isinstance(vector_store_ids_param, list): + vector_store_ids.extend(vector_store_ids_param) # 2. check if vector_store_ids is provided as a tool in the request vector_store_ids = self._get_vector_store_ids_from_tool_calls( tools=tools, vector_store_ids=vector_store_ids ) - return vector_store_ids + return list(dict.fromkeys(vector_store_ids)) def get_and_pop_recognised_vector_store_tools( self, tools: Optional[List[Dict]] = None, vector_store_ids: Optional[List[str]] = None diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 0c7accf632..3d4b68ce44 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -397,6 +397,55 @@ async def test_vector_store_access_check_with_permissions(): 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 = MagicMock() + team_object.object_permission_id = "team-permission" + + mock_prisma_client = MagicMock() + team_permissions = MagicMock() + team_permissions.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 + + def test_can_object_call_model_with_alias(): """Test that can_object_call_model works with model aliases""" from litellm import Router