UI - Add giving keys prompt access (#13233)

* fix(create_key_button.tsx): add prompts on UI

* feat(key_management_endpoints.py): support adding prompt to key via `/key/update`

* fix(key_info_view.tsx): show existing prompts on key in key_info_view.tsx

* fix(key_edit_view.tsx): UX - disable premium feature for non-premium users

prevent accidental clicking

* fix(create_key_button.tsx): disable premium features behind flag, prevent errors

* fix(key_management_endpoints.py): fix key update logic \

* fix: fix check

* docs: document new params

litellm/proxy/management_endpoints/common_utils.py

@@ -1,7 +1,7 @@
-from typing import Any, Union, Optional
+from typing import Any, Optional, Union
 
 from litellm.proxy._types import (
-    GenerateKeyRequest,
+    KeyRequestBase,
     LiteLLM_ManagementEndpoint_MetadataFields_Premium,
     LiteLLM_TeamTable,
     LitellmUserRoles,
@@ -30,7 +30,7 @@ def _is_user_team_admin(
 
 
 def _set_object_metadata_field(
-    object_data: Union[LiteLLM_TeamTable, GenerateKeyRequest],
+    object_data: Union[LiteLLM_TeamTable, KeyRequestBase],
     field_name: str,
     value: Any,
 ) -> None:

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -586,6 +586,7 @@ async def generate_key_fn(
     - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached.
     - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
     - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
+    - prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
     - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
     - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission.
     - key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (uses default allowed routes). Defaults to "default".
@@ -808,6 +809,11 @@ def prepare_metadata_fields(
                     casted_metadata[k] = v.isoformat()
                 else:
                     casted_metadata[k] = v
+            if k in LiteLLM_ManagementEndpoint_MetadataFields_Premium:
+                from litellm.proxy.utils import _premium_user_check
+
+                _premium_user_check()
+                casted_metadata[k] = v
 
     except Exception as e:
         verbose_proxy_logger.exception(
@@ -854,7 +860,10 @@ async def prepare_key_update_data(
     data_json.pop("new_key", None)
     non_default_values = {}
     for k, v in data_json.items():
-        if k in LiteLLM_ManagementEndpoint_MetadataFields:
+        if (
+            k in LiteLLM_ManagementEndpoint_MetadataFields
+            or k in LiteLLM_ManagementEndpoint_MetadataFields_Premium
+        ):
             continue
         non_default_values[k] = v
 
@@ -957,6 +966,7 @@ async def update_key_fn(
     - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
     - models: Optional[list] - Model_name's a user is allowed to call
     - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
+    - prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
     - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
     - spend: Optional[float] - Amount spent by key
     - max_budget: Optional[float] - Max budget for key

litellm/proxy/management_endpoints/team_endpoints.py

@@ -261,6 +261,7 @@ async def new_team(  # noqa: PLR0915
     - blocked: bool - Flag indicating if the team is blocked or not - will stop all calls from keys with this team_id.
     - members: Optional[List] - Control team members via `/team/member/add` and `/team/member/delete`.
     - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
+    - prompts: Optional[List[str]] - List of prompts that the team is allowed to use.
     - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
     - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
     - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
@@ -694,6 +695,7 @@ async def update_team(
     - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
     - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
     - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.
+    - prompts: Optional[List[str]] - List of prompts that the team is allowed to use.
     - blocked: bool - Flag indicating if the team is blocked or not - will stop all calls from keys with this team_id.
     - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
     - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.

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

@@ -34,6 +34,7 @@ import {
   userFilterUICall,
   keyCreateServiceAccountCall,
   fetchMCPAccessGroups,
+  getPromptsList,
 } from "./networking";
 import VectorStoreSelector from "./vector_store_management/VectorStoreSelector";
 import { Team } from "./key_team_helpers/key_list";
@@ -169,6 +170,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({
   const [keyOwner, setKeyOwner] = useState("you");
   const [predefinedTags, setPredefinedTags] = useState(getPredefinedTags(data));
   const [guardrailsList, setGuardrailsList] = useState<string[]>([]);
+  const [promptsList, setPromptsList] = useState<string[]>([]);
   const [loggingSettings, setLoggingSettings] = useState<any[]>([]);
   const [selectedCreateKeyTeam, setSelectedCreateKeyTeam] = useState<Team | null>(team);
   const [isCreateUserModalVisible, setIsCreateUserModalVisible] = useState(false);
@@ -237,7 +239,17 @@ const CreateKey: React.FC<CreateKeyProps> = ({
       }
     };
 
+    const fetchPrompts = async () => {
+      try {
+        const response = await getPromptsList(accessToken);
+        setPromptsList(response.prompts.map(prompt => prompt.prompt_id));
+      } catch (error) {
+        console.error("Failed to fetch prompts:", error);
+      }
+    };
+
     fetchGuardrails();
+    fetchPrompts();
   }, [accessToken]);
 
   // Fetch possible user roles when component mounts
@@ -843,14 +855,61 @@ const CreateKey: React.FC<CreateKeyProps> = ({
                     }
                     name="guardrails" 
                     className="mt-4"
-                    help="Select existing guardrails or enter new ones"
+                    help={premiumUser ? "Select existing guardrails or enter new ones" : "Premium feature - Upgrade to set guardrails by key"}
                   >
-                    <Select
-                      mode="tags"
-                      style={{ width: '100%' }}
-                      placeholder="Select or enter guardrails"
-                      options={guardrailsList.map(name => ({ value: name, label: name }))}
-                    />
+                    <Tooltip 
+                      title={!premiumUser ? "Setting guardrails by key is a premium feature" : ""}
+                      placement="top"
+                    >
+                      <Select
+                        mode="tags"
+                        style={{ width: '100%' }}
+                        disabled={!premiumUser}
+                        placeholder={
+                          !premiumUser
+                            ? "Premium feature - Upgrade to set guardrails by key"
+                            : "Select or enter guardrails"
+                        }
+                        options={guardrailsList.map(name => ({ value: name, label: name }))}
+                      />
+                    </Tooltip>
+                  </Form.Item>
+                  <Form.Item 
+                    label={
+                      <span>
+                        Prompts{' '}
+                        <Tooltip title="Allow this key to use specific prompt templates">
+                          <a 
+                            href="https://docs.litellm.ai/docs/proxy/prompt_management" 
+                            target="_blank" 
+                            rel="noopener noreferrer"
+                            onClick={(e) => e.stopPropagation()} // Prevent accordion from collapsing when clicking link
+                          >
+                            <InfoCircleOutlined style={{ marginLeft: '4px' }} />
+                          </a>
+                        </Tooltip>
+                      </span>
+                    }
+                    name="prompts" 
+                    className="mt-4"
+                    help={premiumUser ? "Select existing prompts or enter new ones" : "Premium feature - Upgrade to set prompts by key"}
+                  >
+                    <Tooltip 
+                      title={!premiumUser ? "Setting prompts by key is a premium feature" : ""}
+                      placement="top"
+                    >
+                      <Select
+                        mode="tags"
+                        style={{ width: '100%' }}
+                        disabled={!premiumUser}
+                        placeholder={
+                          !premiumUser
+                            ? "Premium feature - Upgrade to set prompts by key"
+                            : "Select or enter prompts"
+                        }
+                        options={promptsList.map(name => ({ value: name, label: name }))}
+                      />
+                    </Tooltip>
                   </Form.Item>
                   <Form.Item 
                         label={

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

@@ -1,9 +1,9 @@
 import React, { useState, useEffect } from "react";
-import { Form, Input, Select, Button as AntdButton } from "antd";
+import { Form, Input, Select, Button as AntdButton, Tooltip } from "antd";
 import { Button as TremorButton, TextInput } from "@tremor/react";
 import { KeyResponse } from "./key_team_helpers/key_list";
 import { fetchTeamModels } from "../components/create_key_button";
-import { modelAvailableCall } from "./networking";
+import { modelAvailableCall, getPromptsList } from "./networking";
 import NumericalInput from "./shared/numerical_input";
 import VectorStoreSelector from "./vector_store_management/VectorStoreSelector";
 import MCPServerSelector from "./mcp_server_management/MCPServerSelector";
@@ -54,6 +54,7 @@ export function KeyEditView({
 }: KeyEditViewProps) {
   const [form] = Form.useForm();
   const [userModels, setUserModels] = useState<string[]>([]);
+  const [promptsList, setPromptsList] = useState<string[]>([]);
   const team = teams?.find(team => team.team_id === keyData.team_id);
   const [availableModels, setAvailableModels] = useState<string[]>([]);
   const [mcpAccessGroups, setMcpAccessGroups] = useState<string[]>([]);
@@ -102,6 +103,17 @@ export function KeyEditView({
       }
     };
 
+    const fetchPrompts = async () => {
+      if (!accessToken) return;
+      try {
+        const response = await getPromptsList(accessToken);
+        setPromptsList(response.prompts.map(prompt => prompt.prompt_id));
+      } catch (error) {
+        console.error("Failed to fetch prompts:", error);
+      }
+    };
+
+    fetchPrompts();
     fetchModels();
   }, [userID, userRole, accessToken, team, keyData.team_id]);
 
@@ -127,6 +139,7 @@ export function KeyEditView({
     budget_duration: getBudgetDuration(keyData.budget_duration),
     metadata: formatMetadataForDisplay(keyData.metadata),
     guardrails: keyData.metadata?.guardrails || [],
+    prompts: keyData.metadata?.prompts || [],
     vector_stores: keyData.object_permission?.vector_stores || [],
     mcp_servers_and_groups: {
       servers: keyData.object_permission?.mcp_servers || [],
@@ -138,6 +151,8 @@ export function KeyEditView({
       : []
   };
 
+  console.log("premiumUser:", premiumUser);
+
   return (
     <Form
       form={form}
@@ -200,12 +215,46 @@ export function KeyEditView({
         <Input.TextArea rows={4}  placeholder='{"gpt-4": 100, "claude-v1": 200}'/>
       </Form.Item>
 
+      
       <Form.Item label="Guardrails" name="guardrails">
-        <Select
-          mode="tags"
-          style={{ width: "100%" }}
-          placeholder="Select or enter guardrails"
-        />
+        <Tooltip 
+          title={!premiumUser ? "Setting guardrails by key is a premium feature" : ""}
+          placement="top"
+        >
+          <Select
+            mode="tags"
+            style={{ width: "100%" }}
+            disabled={!premiumUser}
+            placeholder={
+              !premiumUser
+                ? "Premium feature - Upgrade to set guardrails by key"
+                : Array.isArray(keyData.metadata?.guardrails) && keyData.metadata.guardrails.length > 0
+                  ? `Current: ${keyData.metadata.guardrails.join(', ')}`
+                  : "Select or enter guardrails"
+            }
+          />
+        </Tooltip>
+      </Form.Item>
+
+      <Form.Item label="Prompts" name="prompts">
+        <Tooltip 
+          title={!premiumUser ? "Setting prompts by key is a premium feature" : ""}
+          placement="top"
+        >
+          <Select
+            mode="tags"
+            style={{ width: "100%" }}
+            disabled={!premiumUser}
+            placeholder={
+              !premiumUser
+                ? "Premium feature - Upgrade to set prompts by key"
+                : Array.isArray(keyData.metadata?.prompts) && keyData.metadata.prompts.length > 0
+                  ? `Current: ${keyData.metadata.prompts.join(', ')}`
+                  : "Select or enter prompts"
+            }
+            options={promptsList.map(name => ({ value: name, label: name }))}
+          />
+        </Tooltip>
       </Form.Item>
 
       <Form.Item label="Vector Stores" name="vector_stores">

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

@@ -484,6 +484,7 @@ export default function KeyInfoView({
                   accessToken={accessToken}
                   userID={userID}
                   userRole={userRole}
+                  premiumUser={premiumUser}
                 />
               ) : (
                 <div className="space-y-4">
@@ -547,6 +548,19 @@ export default function KeyInfoView({
                         : "Unlimited"}
                     </Text>
                   </div>
+                  
+                  <div>
+                    <Text className="font-medium">Prompts</Text>
+                    <Text>
+                      {Array.isArray(currentKeyData.metadata?.prompts) && currentKeyData.metadata.prompts.length > 0
+                        ? currentKeyData.metadata.prompts.map((prompt, index) => (
+                          <span key={index} className="px-2 mr-2 py-1 bg-blue-100 rounded text-xs">
+                            {prompt}
+                          </span>
+                        ))
+                        : "No prompts specified"}
+                    </Text>
+                  </div>
 
                   <div>
                     <Text className="font-medium">Models</Text>

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

@@ -75,6 +75,16 @@ export interface Model {
   model_info: Object | null;
 }
 
+interface PromptSpec {
+  prompt_id: string;
+  litellm_params: Object;
+  prompt_info: Object;
+}
+
+interface ListPromptsResponse {
+  prompts: PromptSpec[];
+}
+
 export interface Organization {
   organization_id: string | null;
   organization_alias: string;
@@ -4880,6 +4890,33 @@ export const getGuardrailsList = async (accessToken: String) => {
   }
 };
 
+export const getPromptsList = async (accessToken: String) : Promise<ListPromptsResponse> => {
+  try {
+    const url = proxyBaseUrl
+      ? `${proxyBaseUrl}/prompt/list`
+      : `/prompt/list`;
+    const response = await fetch(url, {
+      method: "GET",
+      headers: {
+        [globalLitellmHeaderName]: `Bearer ${accessToken}`,
+        "Content-Type": "application/json",
+      },
+    });
+
+    if (!response.ok) {
+      const errorData = await response.text();
+      handleError(errorData);
+      throw new Error("Network response was not ok");
+    }
+
+    const data = await response.json();
+    return data;
+  } catch (error) {
+    console.error("Failed to get prompts list:", error);
+    throw error;
+  }
+};
+
 export const createGuardrailCall = async (
   accessToken: string,
   guardrailData: any