(UI) - Set/edit guardrails on a virtual key (#7954)

* Revert "JWT Auth - `enforce_rbac` support + UI team view, spend calc fix (#7863)"

This reverts commit dca6904937.

* Revert "Litellm dev 01 10 2025 p2 (#7679)"

This reverts commit c4780479a9.

* ui - allow setting guardrails on a key

* working edit guardrails

* fix edit guardrails on a key

* Revert "Revert "JWT Auth - `enforce_rbac` support + UI team view, spend calc fix (#7863)""

This reverts commit 8f7b9ae1af.

* Revert "Revert "Litellm dev 01 10 2025 p2 (#7679)""

This reverts commit a609139dde.

* fix edit guardrail on ui

* fix list_guardrails

litellm/proxy/_types.py

@@ -2288,7 +2288,6 @@ class ProxyStateVariables(TypedDict):
 UI_TEAM_ID = "litellm-dashboard"
 
 
-
 class JWTAuthBuilderResult(TypedDict):
     is_proxy_admin: bool
     team_object: Optional[LiteLLM_TeamTable]
@@ -2301,6 +2300,7 @@ class JWTAuthBuilderResult(TypedDict):
     end_user_id: Optional[str]
     org_id: Optional[str]
 
+
 class ClientSideFallbackModel(TypedDict, total=False):
     """
     Dictionary passed when client configuring input

litellm/proxy/guardrails/guardrail_endpoints.py

@@ -6,7 +6,6 @@ from typing import Dict, List, Optional, cast
 
 from fastapi import APIRouter, Depends, HTTPException, status
 
-from litellm.proxy._types import CommonProxyErrors
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.types.guardrails import GuardrailInfoResponse, ListGuardrailsResponse
 
@@ -40,7 +39,6 @@ def _get_guardrails_list_response(
 )
 async def list_guardrails():
     """
-    ✨ Enterprise Feature
     List the guardrails that are available on the proxy server
 
     👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start)
@@ -74,15 +72,7 @@ async def list_guardrails():
     }
     ```
     """
-    from litellm.proxy.proxy_server import premium_user, proxy_config
+    from litellm.proxy.proxy_server import proxy_config
-
-    if not premium_user:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail={
-                "error": CommonProxyErrors.not_premium_user.value,
-            },
-        )
 
     config = proxy_config.config
 

litellm/proxy/proxy_config.yaml

@@ -14,6 +14,20 @@ model_list:
       api_base: https://exampleopenaiendpoint-production.up.railway.app/
 
 
-
 litellm_settings:
   callbacks: ["prometheus"]
+
+
+guardrails:
+  - guardrail_name: "bedrock-pre-guard"
+    litellm_params:
+      guardrail: bedrock  # supported values: "aporia", "bedrock", "lakera"
+      mode: "during_call"
+      guardrailIdentifier: ff6ujrregl1q
+      guardrailVersion: "DRAFT"
+  - guardrail_name: "bedrock-post-guard"
+    litellm_params:
+      guardrail: bedrock  # supported values: "aporia", "bedrock", "lakera"
+      mode: "post_call"
+      guardrailIdentifier: ff6ujrregl1q
+      guardrailVersion: "DRAFT"

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

@@ -27,7 +27,10 @@ import {
   keyCreateCall,
   slackBudgetAlertsHealthCheck,
   modelAvailableCall,
+  getGuardrailsList,
 } from "./networking";
+import { InfoCircleOutlined } from '@ant-design/icons';
+import { Tooltip } from 'antd';
 
 const { Option } = Select;
 
@@ -81,7 +84,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({
   const [modelsToPick, setModelsToPick] = useState([]);
   const [keyOwner, setKeyOwner] = useState("you");
   const [predefinedTags, setPredefinedTags] = useState(getPredefinedTags(data));
-
+  const [guardrailsList, setGuardrailsList] = useState<string[]>([]);
 
   const handleOk = () => {
     setIsModalVisible(false);
@@ -121,6 +124,22 @@ const CreateKey: React.FC<CreateKeyProps> = ({
     fetchUserModels();
   }, [accessToken, userID, userRole]);
 
+  useEffect(() => {
+    const fetchGuardrails = async () => {
+      try {
+        const response = await getGuardrailsList(accessToken);
+        const guardrailNames = response.guardrails.map(
+          (g: { guardrail_name: string }) => g.guardrail_name
+        );
+        setGuardrailsList(guardrailNames);
+      } catch (error) {
+        console.error("Failed to fetch guardrails:", error);
+      }
+    };
+
+    fetchGuardrails();
+  }, [accessToken]);
+
   const handleCreate = async (formValues: Record<string, any>) => {
     try {
       const newKeyAlias = formValues?.key_alias ?? "";
@@ -392,6 +411,33 @@ const CreateKey: React.FC<CreateKeyProps> = ({
                 >
                   <TextInput placeholder="" />
                 </Form.Item>
+                <Form.Item 
+                  label={
+                    <span>
+                      Guardrails{' '}
+                      <Tooltip title="Setup your first guardrail">
+                        <a 
+                          href="https://docs.litellm.ai/docs/proxy/guardrails/quick_start" 
+                          target="_blank" 
+                          rel="noopener noreferrer"
+                          onClick={(e) => e.stopPropagation()} // Prevent accordion from collapsing when clicking link
+                        >
+                          <InfoCircleOutlined style={{ marginLeft: '4px' }} />
+                        </a>
+                      </Tooltip>
+                    </span>
+                  }
+                  name="guardrails" 
+                  className="mt-8"
+                  help="Select existing guardrails or enter new ones"
+                >
+                  <Select
+                    mode="tags"
+                    style={{ width: '100%' }}
+                    placeholder="Select or enter guardrails"
+                    options={guardrailsList.map(name => ({ value: name, label: name }))}
+                  />
+                </Form.Item>
 
                 <Form.Item label="Metadata" name="metadata" className="mt-8">
                   <Input.TextArea

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

@@ -3050,3 +3050,31 @@ export const getProxyUISettings = async (
   }
 };
 
+
+
+export const getGuardrailsList = async (accessToken: String) => {
+  try {
+    let url = proxyBaseUrl ? `${proxyBaseUrl}/guardrails/list` : `/guardrails/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();
+    console.log("Guardrails list response:", data);
+    return data;
+  } catch (error) {
+    console.error("Failed to fetch guardrails list:", error);
+    throw error;
+  }
+};

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

@@ -206,7 +206,6 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
 
             setUserSpendData(response["user_info"]);
             console.log(`userSpendData: ${JSON.stringify(userSpendData)}`)
-            
             const teamsArray = [...response["teams"]];
             if (teamsArray.length > 0) {
               console.log(`response['teams']: ${JSON.stringify(teamsArray)}`);

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

@@ -1,6 +1,6 @@
 "use client";
 import React, { useEffect, useState } from "react";
-import { keyDeleteCall, modelAvailableCall } from "./networking";
+import { keyDeleteCall, modelAvailableCall, getGuardrailsList } from "./networking";
 import { add } from 'date-fns';
 import { InformationCircleIcon, StatusOnlineIcon, TrashIcon, PencilAltIcon, RefreshIcon } from "@heroicons/react/outline";
 import { keySpendLogsCall, PredictedSpendLogsCall, keyUpdateCall, modelInfoCall, regenerateKeyCall } from "./networking";
@@ -26,6 +26,7 @@ import {
   TextInput,
   Textarea,
 } from "@tremor/react";
+import { InfoCircleOutlined } from '@ant-design/icons';
 import { Select as Select3, SelectItem, MultiSelect, MultiSelectItem } from "@tremor/react";
 import {
   Button as Button2,
@@ -130,6 +131,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   const [newExpiryTime, setNewExpiryTime] = useState<string | null>(null);
 
   const [knownTeamIDs, setKnownTeamIDs] = useState(initialKnownTeamIDs);
+  const [guardrailsList, setGuardrailsList] = useState<string[]>([]);
 
   // Function to check if user is admin of a team
   const isUserTeamAdmin = (team: any) => {
@@ -311,22 +313,49 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
     const [keyTeam, setKeyTeam] = useState(selectedTeam);
     const [errorModels, setErrorModels] = useState<string[]>([]);
     const [errorBudget, setErrorBudget] = useState<boolean>(false);
+    const [guardrailsList, setGuardrailsList] = useState<string[]>([]);
+
+    useEffect(() => {
+      const fetchGuardrails = async () => {
+        try {
+          const response = await getGuardrailsList(accessToken);
+          const guardrailNames = response.guardrails.map(
+            (g: { guardrail_name: string }) => g.guardrail_name
+          );
+          setGuardrailsList(guardrailNames);
+        } catch (error) {
+          console.error("Failed to fetch guardrails:", error);
+        }
+      };
+
+      fetchGuardrails();
+    }, [accessToken]);
 
     let metadataString = '';
     try {
-      metadataString = JSON.stringify(token.metadata, null, 2);
+      // Create a copy of metadata without guardrails for display
+      const displayMetadata = { ...token.metadata };
+      delete displayMetadata.guardrails;
+      metadataString = JSON.stringify(displayMetadata, null, 2);
     } catch (error) {
       console.error("Error stringifying metadata:", error);
-      // You can choose a fallback, such as an empty string or a warning message
       metadataString = '';
     }
 
-    // Ensure token is defined and handle gracefully if not
+    // Extract existing guardrails from metadata
+    let existingGuardrails: string[] = [];
+    try {
+      existingGuardrails = token.metadata?.guardrails || [];
+    } catch (error) {
+      console.error("Error extracting guardrails:", error);
+    }
+
     const initialValues = token ? {
       ...token,
       budget_duration: token.budget_duration,
-      metadata: metadataString
+      metadata: metadataString,
-    } : { metadata: metadataString };
+      guardrails: existingGuardrails
+    } : { metadata: metadataString, guardrails: [] };
 
 
     const handleOk = () => {
@@ -508,6 +537,33 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
               >
                 <InputNumber step={1} precision={1} width={200} />
               </Form.Item>
+              <Form.Item
+                label={
+                  <span>
+                    Guardrails{' '}
+                    <Tooltip title="Setup your first guardrail">
+                      <a 
+                        href="https://docs.litellm.ai/docs/proxy/guardrails/quick_start" 
+                        target="_blank" 
+                        rel="noopener noreferrer"
+                        onClick={(e) => e.stopPropagation()}
+                      >
+                        <InfoCircleOutlined style={{ marginLeft: '4px' }} />
+                      </a>
+                    </Tooltip>
+                  </span>
+                }
+                name="guardrails" 
+                className="mt-8"
+                help="Select existing guardrails or enter new ones"
+              >
+                <Select
+                  mode="tags"
+                  style={{ width: '100%' }}
+                  placeholder="Select or enter guardrails"
+                  options={guardrailsList.map(name => ({ value: name, label: name }))}
+                />
+              </Form.Item>
               <Form.Item
                 label="Metadata (ensure this is valid JSON)"
                 name="metadata"
@@ -747,12 +803,23 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   // Convert metadata back to an object if it exists and is a string
   if (formValues.metadata && typeof formValues.metadata === 'string') {
     try {
-      formValues.metadata = JSON.parse(formValues.metadata);
+      const parsedMetadata = JSON.parse(formValues.metadata);
+      // Only add guardrails if they are set in form values
+      formValues.metadata = {
+        ...parsedMetadata,
+        ...(formValues.guardrails?.length > 0 ? { guardrails: formValues.guardrails } : {})
+      };
     } catch (error) {
       console.error("Error parsing metadata JSON:", error);
       message.error("Invalid metadata JSON for formValue " + formValues.metadata);
       return;
     }
+  } else {
+    // If metadata is not a string (or doesn't exist), only add guardrails if they are set
+    formValues.metadata = {
+      ...(formValues.metadata || {}),
+      ...(formValues.guardrails?.length > 0 ? { guardrails: formValues.guardrails } : {})
+    };
   }
 
   // Convert the budget_duration back to the API expected format
@@ -772,6 +839,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
 
   console.log("handleEditSubmit:", formValues);
 
+  try {
     let newKeyValues = await keyUpdateCall(accessToken, formValues);
     console.log("handleEditSubmit: newKeyValues", newKeyValues);
 
@@ -786,9 +854,14 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
 
     setEditModalVisible(false);
     setSelectedToken(null);
+  } catch (error) {
+    console.error("Error updating key:", error);
+    message.error("Failed to update key");
+  }
 };
 
 
+
   const handleDelete = async (token: any) => {
     console.log("handleDelete:", token);
     if (token.token == null) {