Merge pull request #2801 from BerriAI/litellm_support_all_models_as_a_ui_alias

[UI] use all_models alias

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

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()

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)

ui/litellm-dashboard/src/components/create_key_button.tsx

@@ -91,19 +91,6 @@ const CreateKey: React.FC<CreateKeyProps> = ({
     }
   };
 
-
-  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<CreateKeyProps> = ({
                   mode="multiple"
                   placeholder="Select models"
                   style={{ width: "100%" }}
-                  onChange={(selectedModels) => handleModelSelection(selectedModels)}
                 >
-                  <Option key="all_models" value="all_models">
-                    All Models
-                  </Option>
                   {team && team.models ? (
                     team.models.map((model: string) => (
                       <Option key={model} value={model}>

ui/litellm-dashboard/src/components/teams.tsx

@@ -115,11 +115,17 @@ const Team: React.FC<TeamProps> = ({
                   <Input />
                 </Form.Item>
                 <Form.Item label="Models" name="models">
+                  <Select2.Option key="all-proxy-models" value="all-proxy-models">
+                      All Models on Proxy
+                    </Select2.Option>
                   <Select2
                     mode="multiple"
                     placeholder="Select models"
                     style={{ width: "100%" }}
                   >
+                    <Select2.Option key="all-proxy-models" value="all-proxy-models">
+                      All Models on Proxy
+                    </Select2.Option>
                     {userModels && userModels.map((model) => (
                       <Select2.Option key={model} value={model}>
                         {model}
@@ -216,15 +222,7 @@ const handleEditSubmit = async (formValues: Record<string, any>) => {
     setIsDeleteModalOpen(true);
   };
 
-  const handleModelSelection = (selectedModels: string[]) => {
-    if (selectedModels.includes("all_models")) {
-      // Select all models except "All Models"
-      const allModelsExceptAll = userModels.filter(model => model !== "all");
-      form.setFieldsValue({
-        models: allModelsExceptAll
-      });
-    }
-  };
+
   
 
   const confirmDelete = async () => {
@@ -365,8 +363,8 @@ const handleEditSubmit = async (formValues: Record<string, any>) => {
                             {Array.isArray(team.models) ? (
                               <div style={{ display: "flex", flexDirection: "column" }}>
                                 {team.models.length === 0 ? (
-                                  <Badge size={"xs"} className="mb-1" color="purple">
-                                    <Text>All Models</Text>
+                                  <Badge size={"xs"} className="mb-1" color="blue">
+                                    <Text>all-proxy-models</Text>
                                   </Badge>
                                 ) : (
                                   team.models.map((model: string, index: number) => (
@@ -484,10 +482,9 @@ const handleEditSubmit = async (formValues: Record<string, any>) => {
                     mode="multiple"
                     placeholder="Select models"
                     style={{ width: "100%" }}
-                    onChange={(selectedModels) => handleModelSelection(selectedModels)}
                   >
-                    <Select2.Option key="all_models" value="all_models">
-                      All Models
+                    <Select2.Option key="all-proxy-models" value="all-proxy-models">
+                      All Models on Proxy
                     </Select2.Option>
                     {userModels.map((model) => (
                       <Select2.Option key={model} value={model}>

ui/litellm-dashboard/src/components/view_key_table.tsx

@@ -100,15 +100,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
       
     }
 
-    const handleModelSelection = (selectedModels: string[]) => {
-      if (selectedModels.includes("all_models")) {
-        // Select all models except "All Models"
-        const allModelsExceptAll = selectedTeam ? selectedTeam.models : userModels;
-        form.setFieldsValue({
-          models: allModelsExceptAll
-        });
-      }
-    };
+
   
     const handleOk = () => {
       form
@@ -155,11 +147,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
                   mode="multiple"
                   placeholder="Select models"
                   style={{ width: "100%" }}
-                  onChange={(selectedModels) => handleModelSelection(selectedModels)}
-                >
-                  <Option key="all_models" value="all_models">
-                    All Models
-                  </Option>
+                >                
                   {selectedTeam && selectedTeam.models ? (
                     selectedTeam.models.map((model: string) => (
                       <Option key={model} value={model}>
@@ -429,8 +417,8 @@ const handleEditSubmit = async (formValues: Record<string, any>) => {
                             ))
                           ) : (
                             // If selected team is None or selected team's models are empty, show all models
-                            <Badge size={"xs"} className="mb-1" color="purple">
-                              <Text>All Models</Text>
+                            <Badge size={"xs"} className="mb-1" color="blue">
+                              <Text>all-proxy-models</Text>
                             </Badge>
                           )}
                         </>