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 = ({ const [modelsToPick, setModelsToPick] = useState([]); const [keyOwner, setKeyOwner] = useState("you"); const [predefinedTags, setPredefinedTags] = useState(getPredefinedTags(data)); - + const [guardrailsList, setGuardrailsList] = useState([]); const handleOk = () => { setIsModalVisible(false); @@ -121,6 +124,22 @@ const CreateKey: React.FC = ({ 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) => { try { const newKeyAlias = formValues?.key_alias ?? ""; @@ -392,6 +411,33 @@ const CreateKey: React.FC = ({ > + + Guardrails{' '} + + e.stopPropagation()} // Prevent accordion from collapsing when clicking link + > + + + + + } + name="guardrails" + className="mt-8" + help="Select existing guardrails or enter new ones" + > + ({ value: name, label: name }))} + /> + = ({ setSelectedToken(null); }; - const handleEditSubmit = async (formValues: Record) => { +const handleEditSubmit = async (formValues: Record) => { /** * Call API to update team with teamId and values * @@ -747,12 +803,23 @@ const ViewKeyTable: React.FC = ({ // 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 = ({ 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) => {