From b79f55eec0aa21823f0cb83173d511c932aea563 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 2 Aug 2025 19:33:17 -0700 Subject: [PATCH] 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 --- .../management_endpoints/common_utils.py | 6 +- .../key_management_endpoints.py | 12 ++- .../management_endpoints/team_endpoints.py | 2 + .../src/components/create_key_button.tsx | 73 +++++++++++++++++-- .../src/components/key_edit_view.tsx | 63 ++++++++++++++-- .../src/components/key_info_view.tsx | 14 ++++ .../src/components/networking.tsx | 37 ++++++++++ 7 files changed, 189 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index c1db60c6be..fcab831e3e 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/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: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8209d9f1f9..ad483a5ac1 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/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 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 32f48f1932..58084874f4 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/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`. diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx index 689308d890..e9e6bda598 100644 --- a/ui/litellm-dashboard/src/components/create_key_button.tsx +++ b/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 = ({ const [keyOwner, setKeyOwner] = useState("you"); const [predefinedTags, setPredefinedTags] = useState(getPredefinedTags(data)); const [guardrailsList, setGuardrailsList] = useState([]); + const [promptsList, setPromptsList] = useState([]); const [loggingSettings, setLoggingSettings] = useState([]); const [selectedCreateKeyTeam, setSelectedCreateKeyTeam] = useState(team); const [isCreateUserModalVisible, setIsCreateUserModalVisible] = useState(false); @@ -237,7 +239,17 @@ const CreateKey: React.FC = ({ } }; + 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 = ({ } 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"} > - ({ value: name, label: name }))} + /> + + + + Prompts{' '} + + e.stopPropagation()} // Prevent accordion from collapsing when clicking link + > + + + + + } + name="prompts" + className="mt-4" + help={premiumUser ? "Select existing prompts or enter new ones" : "Premium feature - Upgrade to set prompts by key"} + > + + + + 0 + ? `Current: ${keyData.metadata.prompts.join(', ')}` + : "Select or enter prompts" + } + options={promptsList.map(name => ({ value: name, label: name }))} + /> + diff --git a/ui/litellm-dashboard/src/components/key_info_view.tsx b/ui/litellm-dashboard/src/components/key_info_view.tsx index 438e496357..f97afb4b5c 100644 --- a/ui/litellm-dashboard/src/components/key_info_view.tsx +++ b/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} /> ) : (
@@ -547,6 +548,19 @@ export default function KeyInfoView({ : "Unlimited"}
+ +
+ Prompts + + {Array.isArray(currentKeyData.metadata?.prompts) && currentKeyData.metadata.prompts.length > 0 + ? currentKeyData.metadata.prompts.map((prompt, index) => ( + + {prompt} + + )) + : "No prompts specified"} + +
Models diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 1443b7f59d..bc6db4c6ad 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/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 => { + 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