diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e3b7df4e9c..8cfa5587b6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -52,9 +52,14 @@ def common_checks( and len(team_object.models) > 0 and _model not in team_object.models ): - raise Exception( - f"Team={team_object.team_id} not allowed to call model={_model}. Allowed team models = {team_object.models}" - ) + # this means the team has access to all models on the proxy + if "all-proxy-models" in team_object.models: + # this means the team has access to all models on the proxy + pass + else: + raise Exception( + f"Team={team_object.team_id} not allowed to call model={_model}. Allowed team models = {team_object.models}" + ) # 3. If team is in budget if ( team_object.max_budget is not None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0cfc661ece..e8c0a1a289 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -637,6 +637,12 @@ async def user_api_key_auth( len(valid_token.models) == 0 ): # assume an empty model list means all models are allowed to be called pass + elif ( + isinstance(valid_token.models, list) + and "all-proxy-models" in valid_token.models + ): + # Admin UI - Special alias to allow `all_models` + pass else: try: data = await request.json() diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 8430aa7fd2..aa1fbc2fca 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -49,6 +49,7 @@ from litellm.proxy.proxy_server import ( spend_key_fn, view_spend_logs, user_info, + info_key_fn, ) from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token from litellm._logging import verbose_proxy_logger @@ -245,6 +246,44 @@ def test_call_with_valid_model(prisma_client): pytest.fail(f"An exception occurred - {str(e)}") +def test_call_with_valid_model_using_all_models(prisma_client): + # Make a call to a key with model = `all-proxy-models` this is an Alias from LiteLLM Admin UI + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + try: + + async def test(): + await litellm.proxy.proxy_server.prisma_client.connect() + request = GenerateKeyRequest(models=["all-proxy-models"]) + key = await generate_key_fn(data=request) + print(key) + + generated_key = key.key + bearer_token = "Bearer " + generated_key + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + async def return_body(): + return b'{"model": "mistral"}' + + request.body = return_body + + # use generated key to auth in + result = await user_api_key_auth(request=request, api_key=bearer_token) + print("result from user auth with new key", result) + + # call /key/info for key - models == "all-proxy-models" + key_info = await info_key_fn(key=generated_key) + print("key_info", key_info) + models = key_info["info"]["models"] + assert models == ["all-proxy-models"] + + asyncio.run(test()) + except Exception as e: + pytest.fail(f"An exception occurred - {str(e)}") + + def test_call_with_user_over_budget(prisma_client): # 5. Make a call with a key over budget, expect to fail setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx index 17476232c1..99cbd7b6e6 100644 --- a/ui/litellm-dashboard/src/components/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/create_key_button.tsx @@ -91,19 +91,6 @@ const CreateKey: React.FC = ({ } }; - - const handleModelSelection = (selectedModels: string[]) => { - if (selectedModels.includes("all_models")) { - // Select all models except "All Models" - const allModelsExceptAll = team ? team.models : userModels; - form.setFieldsValue({ - models: allModelsExceptAll - }); - } - }; - - - const handleCopy = () => { message.success('API Key copied to clipboard'); }; @@ -154,11 +141,7 @@ const CreateKey: React.FC = ({ mode="multiple" placeholder="Select models" style={{ width: "100%" }} - onChange={(selectedModels) => handleModelSelection(selectedModels)} > - {team && team.models ? ( team.models.map((model: string) => (