From 84fb5aead804a22173d7b5143cee0e6a9d35dd8e Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 23 Jan 2025 18:01:54 -0800
Subject: [PATCH] (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 dca6904937edc2fac4560b4b7211b55f8c2a6d3f.

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

This reverts commit c4780479a990f532083bdf1e9fa88750eff00098.

* 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 8f7b9ae1af71075f6129a49bfc4af094eb2398c7.

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

This reverts commit a609139dded26b5f59dc55dd68ad0baa470b6e53.

* fix edit guardrail on ui

* fix list_guardrails
---
 litellm/proxy/_types.py                       |   2 +-
 .../proxy/guardrails/guardrail_endpoints.py   |  12 +-
 litellm/proxy/proxy_config.yaml               |  16 ++-
 .../src/components/create_key_button.tsx      |  48 +++++++-
 .../src/components/networking.tsx             |  28 +++++
 .../src/components/user_dashboard.tsx         |   1 -
 .../src/components/view_key_table.tsx         | 113 ++++++++++++++----
 7 files changed, 185 insertions(+), 35 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 937baf9c382..c5fe5518786 100644
--- a/litellm/proxy/_types.py
+++ b/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
diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py
index 7ca64b89a5f..b1919b45ef8 100644
--- a/litellm/proxy/guardrails/guardrail_endpoints.py
+++ b/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
-
-    if not premium_user:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail={
-                "error": CommonProxyErrors.not_premium_user.value,
-            },
-        )
+    from litellm.proxy.proxy_server import proxy_config
 
     config = proxy_config.config
 
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index b8b69fd2821..088c8a05145 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/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"
\ No newline at end of file
diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx
index 7060128dc96..59e7e275a62 100644
--- a/ui/litellm-dashboard/src/components/create_key_button.tsx
+++ b/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
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx
index f219b563f9a..159bc035d19 100644
--- a/ui/litellm-dashboard/src/components/networking.tsx
+++ b/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;
+  }
+};
diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx
index e0b3c5959ed..e00f3602327 100644
--- a/ui/litellm-dashboard/src/components/user_dashboard.tsx
+++ b/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)}`);
diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx
index 144c30c0d20..ea339d5bd58 100644
--- a/ui/litellm-dashboard/src/components/view_key_table.tsx
+++ b/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"
@@ -731,7 +787,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
     setSelectedToken(null);
   };
 
-  const handleEditSubmit = async (formValues: Record<string, any>) => {
+const handleEditSubmit = async (formValues: Record<string, any>) => {
   /**
    * Call API to update team with teamId and values
    * 
@@ -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,21 +839,27 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
 
   console.log("handleEditSubmit:", formValues);
 
-  let newKeyValues = await keyUpdateCall(accessToken, formValues);
-  console.log("handleEditSubmit: newKeyValues", newKeyValues);
+  try {
+    let newKeyValues = await keyUpdateCall(accessToken, formValues);
+    console.log("handleEditSubmit: newKeyValues", newKeyValues);
 
-  // Update the keys with the update key
-  if (data) {
-    const updatedData = data.map((key) =>
-      key.token === currentKey ? newKeyValues : key
-    );
-    setData(updatedData);
+    // Update the keys with the update key
+    if (data) {
+      const updatedData = data.map((key) =>
+        key.token === currentKey ? newKeyValues : key
+      );
+      setData(updatedData);
+    }
+    message.success("Key updated successfully");
+
+    setEditModalVisible(false);
+    setSelectedToken(null);
+  } catch (error) {
+    console.error("Error updating key:", error);
+    message.error("Failed to update key");
   }
-  message.success("Key updated successfully");
+};
 
-  setEditModalVisible(false);
-  setSelectedToken(null);
-  };
 
 
   const handleDelete = async (token: any) => {