diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index 50a98d01625..97367e59023 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -5,7 +5,6 @@ from fastapi.responses import ORJSONResponse import litellm from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth.auth_checks import _can_object_call_model, can_key_call_model 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.openai_endpoint_utils import ( @@ -14,6 +13,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_query, ) from litellm.proxy.openai_files_endpoints.common_utils import ( + authorize_model_for_key, get_credentials_for_model, handle_model_based_routing, prepare_data_with_credentials, @@ -212,26 +212,7 @@ async def _authorize_model_routing_hint( ) -> None: if user_api_key_dict is None: return - - key_models: Final = getattr(user_api_key_dict, "models", None) - if not (isinstance(key_models, list) and "all-team-models" in key_models): - await can_key_call_model( - model=model, - llm_model_list=None, - valid_token=user_api_key_dict, - llm_router=llm_router, - ) - - team_models: Final = getattr(user_api_key_dict, "team_models", None) - if isinstance(team_models, list) and len(team_models) > 0: - _can_object_call_model( - model=model, - llm_router=llm_router, - models=team_models, - team_model_aliases=user_api_key_dict.team_model_aliases, - team_id=user_api_key_dict.team_id, - object_type="team", - ) + await authorize_model_for_key(model_id=model, llm_router=llm_router, user_api_key_dict=user_api_key_dict) async def _update_request_data_with_model_routing_hint( diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 52672b596ea..20484e787bd 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -609,6 +609,44 @@ async def test_vector_store_file_list_authorizes_model_query_param_before_creden llm_router.get_deployment_credentials_with_provider.assert_not_called() +@pytest.mark.asyncio +async def test_vector_store_file_list_model_query_param_enforces_project_model_grant(): + from litellm.proxy._types import LiteLLM_ProjectTableCachedObj, LiteLLM_TeamTableCachedObj + from litellm.proxy.auth.auth_checks import ProxyException + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, project_cache_key + + request = MagicMock(spec=Request) + request.query_params = {"model": "team-openai"} + request.headers = {} + + llm_router = MagicMock() + llm_router.model_group_alias = {} + cache = UserApiKeyCache() + await cache.async_set_cache( + key="team_id:team-123", + value=LiteLLM_TeamTableCachedObj(team_id="team-123", models=["team-openai"]), + ) + await cache.async_set_cache( + key=project_cache_key("proj-1"), + value=LiteLLM_ProjectTableCachedObj(project_id="proj-1", models=["other-deployment"]), + ) + user_api_key_dict = UserAPIKeyAuth(team_id="team-123", team_models=["team-openai"], project_id="proj-1") + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: proxy_server global, no seam + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: proxy_server global, no seam + ): + with pytest.raises(ProxyException): + await _update_request_data_with_model_routing_hint( + data={"vector_store_id": "vs_123"}, + request=request, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) + + llm_router.get_deployment_credentials_with_provider.assert_not_called() + + @pytest.mark.asyncio async def test_update_request_data_with_litellm_managed_vector_store_registry(): """