diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index c2a8bd11d0..1d6f3b5211 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -470,4 +470,14 @@ model LiteLLM_ManagedVectorStoresTable { created_at DateTime @default(now()) updated_at DateTime @updatedAt litellm_credential_name String? +} + +// Guardrails table for storing guardrail configurations +model LiteLLM_GuardrailsTable { + guardrail_id String @id @default(uuid()) + guardrail_name String @unique + litellm_params Json + guardrail_info Json? + created_at DateTime @default(now()) + updated_at DateTime @updatedAt } \ No newline at end of file diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 7407d6fb12..75bbd5f710 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -4,14 +4,21 @@ CRUD ENDPOINTS FOR GUARDRAILS from typing import Dict, List, Optional, cast -from fastapi import APIRouter, Depends +from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.guardrails import GuardrailInfoResponse, ListGuardrailsResponse +from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry +from litellm.types.guardrails import ( + Guardrail, + GuardrailInfoResponse, + ListGuardrailsResponse, +) #### GUARDRAILS ENDPOINTS #### router = APIRouter() +GUARDRAIL_REGISTRY = GuardrailRegistry() def _get_guardrails_list_response( @@ -83,3 +90,266 @@ async def list_guardrails(): return _get_guardrails_list_response([]) return _get_guardrails_list_response(_guardrails_config) + + +class CreateGuardrailRequest(BaseModel): + guardrail: Guardrail + + +@router.post( + "/guardrails", + tags=["Guardrails"], + dependencies=[Depends(user_api_key_auth)], +) +async def create_guardrail(request: CreateGuardrailRequest): + """ + Create a new guardrail + + 👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start) + + Example Request: + ```bash + curl -X POST "http://localhost:4000/guardrails" \\ + -H "Authorization: Bearer " \\ + -H "Content-Type: application/json" \\ + -d '{ + "guardrail": { + "guardrail_name": "my-bedrock-guard", + "litellm_params": { + "guardrail": "bedrock", + "mode": "pre_call", + "guardrailIdentifier": "ff6ujrregl1q", + "guardrailVersion": "DRAFT", + "default_on": true + }, + "guardrail_info": { + "description": "Bedrock content moderation guardrail" + } + } + }' + ``` + + Example Response: + ```json + { + "guardrail_id": "123e4567-e89b-12d3-a456-426614174000", + "guardrail_name": "my-bedrock-guard", + "litellm_params": { + "guardrail": "bedrock", + "mode": "pre_call", + "guardrailIdentifier": "ff6ujrregl1q", + "guardrailVersion": "DRAFT", + "default_on": true + }, + "guardrail_info": { + "description": "Bedrock content moderation guardrail" + }, + "created_at": "2023-11-09T12:34:56.789Z", + "updated_at": "2023-11-09T12:34:56.789Z" + } + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Prisma client not initialized") + + try: + result = await GUARDRAIL_REGISTRY.add_guardrail_to_db( + guardrail=request.guardrail, prisma_client=prisma_client + ) + return result + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/guardrails/{guardrail_id}", + tags=["Guardrails"], + dependencies=[Depends(user_api_key_auth)], +) +async def get_guardrail(guardrail_id: str): + """ + Get a guardrail by ID + + 👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start) + + Example Request: + ```bash + curl -X GET "http://localhost:4000/guardrails/123e4567-e89b-12d3-a456-426614174000" \\ + -H "Authorization: Bearer " + ``` + + Example Response: + ```json + { + "guardrail_id": "123e4567-e89b-12d3-a456-426614174000", + "guardrail_name": "my-bedrock-guard", + "litellm_params": { + "guardrail": "bedrock", + "mode": "pre_call", + "guardrailIdentifier": "ff6ujrregl1q", + "guardrailVersion": "DRAFT", + "default_on": true + }, + "guardrail_info": { + "description": "Bedrock content moderation guardrail" + }, + "created_at": "2023-11-09T12:34:56.789Z", + "updated_at": "2023-11-09T12:34:56.789Z" + } + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Prisma client not initialized") + + try: + result = await GUARDRAIL_REGISTRY.get_guardrail_by_id_from_db( + guardrail_id=guardrail_id, prisma_client=prisma_client + ) + if result is None: + raise HTTPException( + status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" + ) + return result + except HTTPException as e: + raise e + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +class UpdateGuardrailRequest(BaseModel): + guardrail: Guardrail + + +@router.put( + "/guardrails/{guardrail_id}", + tags=["Guardrails"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest): + """ + Update an existing guardrail + + 👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start) + + Example Request: + ```bash + curl -X PUT "http://localhost:4000/guardrails/123e4567-e89b-12d3-a456-426614174000" \\ + -H "Authorization: Bearer " \\ + -H "Content-Type: application/json" \\ + -d '{ + "guardrail": { + "guardrail_name": "updated-bedrock-guard", + "litellm_params": { + "guardrail": "bedrock", + "mode": "pre_call", + "guardrailIdentifier": "ff6ujrregl1q", + "guardrailVersion": "1.0", + "default_on": true + }, + "guardrail_info": { + "description": "Updated Bedrock content moderation guardrail" + } + } + }' + ``` + + Example Response: + ```json + { + "guardrail_id": "123e4567-e89b-12d3-a456-426614174000", + "guardrail_name": "updated-bedrock-guard", + "litellm_params": { + "guardrail": "bedrock", + "mode": "pre_call", + "guardrailIdentifier": "ff6ujrregl1q", + "guardrailVersion": "1.0", + "default_on": true + }, + "guardrail_info": { + "description": "Updated Bedrock content moderation guardrail" + }, + "created_at": "2023-11-09T12:34:56.789Z", + "updated_at": "2023-11-09T13:45:12.345Z" + } + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Prisma client not initialized") + + try: + # Check if guardrail exists + existing_guardrail = await GUARDRAIL_REGISTRY.get_guardrail_by_id_from_db( + guardrail_id=guardrail_id, prisma_client=prisma_client + ) + + if existing_guardrail is None: + raise HTTPException( + status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" + ) + + result = await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=request.guardrail, + prisma_client=prisma_client, + ) + return result + except HTTPException as e: + raise e + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +@router.delete( + "/guardrails/{guardrail_id}", + tags=["Guardrails"], + dependencies=[Depends(user_api_key_auth)], +) +async def delete_guardrail(guardrail_id: str): + """ + Delete a guardrail + + 👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start) + + Example Request: + ```bash + curl -X DELETE "http://localhost:4000/guardrails/123e4567-e89b-12d3-a456-426614174000" \\ + -H "Authorization: Bearer " + ``` + + Example Response: + ```json + { + "message": "Guardrail 123e4567-e89b-12d3-a456-426614174000 deleted successfully" + } + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Prisma client not initialized") + + try: + # Check if guardrail exists + existing_guardrail = await GUARDRAIL_REGISTRY.get_guardrail_by_id_from_db( + guardrail_id=guardrail_id, prisma_client=prisma_client + ) + + if existing_guardrail is None: + raise HTTPException( + status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" + ) + + result = await GUARDRAIL_REGISTRY.delete_guardrail_from_db( + guardrail_id=guardrail_id, prisma_client=prisma_client + ) + return result + except HTTPException as e: + raise e + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 6b43fd474b..0557e908d4 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -1,6 +1,10 @@ # litellm/proxy/guardrails/guardrail_registry.py -from litellm.types.guardrails import SupportedGuardrailIntegrations +from datetime import datetime, timezone +from typing import List, Optional + +from litellm.proxy.utils import PrismaClient +from litellm.types.guardrails import Guardrail, SupportedGuardrailIntegrations from .guardrail_initializers import ( initialize_aim, @@ -21,3 +25,147 @@ guardrail_registry = { SupportedGuardrailIntegrations.HIDE_SECRETS.value: initialize_hide_secrets, SupportedGuardrailIntegrations.GURDRAILS_AI.value: initialize_guardrails_ai, } + + +class GuardrailRegistry: + """ + Registry for guardrails + + Handles adding, removing, and getting guardrails in DB + in memory + """ + + def __init__(self): + pass + + ########################################################### + ########### DB management helpers for guardrails ########### + ############################################################ + async def add_guardrail_to_db( + self, guardrail: Guardrail, prisma_client: PrismaClient + ): + """ + Add a guardrail to the database + """ + try: + guardrail_name = guardrail.get("guardrail_name") + litellm_params = guardrail.get("litellm_params", {}) + guardrail_info = guardrail.get("guardrail_info", {}) + + # Create guardrail in DB + created_guardrail = await prisma_client.db.litellm_guardrailstable.create( + data={ + "guardrail_name": guardrail_name, + "litellm_params": litellm_params, + "guardrail_info": guardrail_info, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } + ) + + # Add guardrail_id to the returned guardrail object + guardrail_dict = dict(guardrail) + guardrail_dict["guardrail_id"] = created_guardrail.guardrail_id + + return guardrail_dict + except Exception as e: + raise Exception(f"Error adding guardrail to DB: {str(e)}") + + async def delete_guardrail_from_db( + self, guardrail_id: str, prisma_client: PrismaClient + ): + """ + Delete a guardrail from the database + """ + try: + # Delete from DB + await prisma_client.db.litellm_guardrailstable.delete( + where={"guardrail_id": guardrail_id} + ) + + return {"message": f"Guardrail {guardrail_id} deleted successfully"} + except Exception as e: + raise Exception(f"Error deleting guardrail from DB: {str(e)}") + + async def update_guardrail_in_db( + self, guardrail_id: str, guardrail: Guardrail, prisma_client: PrismaClient + ): + """ + Update a guardrail in the database + """ + try: + guardrail_name = guardrail.get("guardrail_name") + litellm_params = guardrail.get("litellm_params", {}) + guardrail_info = guardrail.get("guardrail_info", {}) + + # Update in DB + updated_guardrail = await prisma_client.db.litellm_guardrailstable.update( + where={"guardrail_id": guardrail_id}, + data={ + "guardrail_name": guardrail_name, + "litellm_params": litellm_params, + "guardrail_info": guardrail_info, + "updated_at": datetime.now(timezone.utc), + }, + ) + + # Convert to dict and return + return dict(updated_guardrail) + except Exception as e: + raise Exception(f"Error updating guardrail in DB: {str(e)}") + + async def get_all_guardrails_from_db( + self, prisma_client: PrismaClient + ) -> List[Guardrail]: + """ + Get all guardrails from the database + """ + try: + guardrails_from_db = ( + await prisma_client.db.litellm_guardrailstable.find_many( + order={"created_at": "desc"}, + ) + ) + + guardrails: List[Guardrail] = [] + for guardrail in guardrails_from_db: + guardrails.append(Guardrail(**(dict(guardrail)))) + + return guardrails + except Exception as e: + raise Exception(f"Error getting guardrails from DB: {str(e)}") + + async def get_guardrail_by_id_from_db( + self, guardrail_id: str, prisma_client: PrismaClient + ) -> Optional[Guardrail]: + """ + Get a guardrail by its ID from the database + """ + try: + guardrail = await prisma_client.db.litellm_guardrailstable.find_unique( + where={"guardrail_id": guardrail_id} + ) + + if not guardrail: + return None + + return Guardrail(**(dict(guardrail))) + except Exception as e: + raise Exception(f"Error getting guardrail from DB: {str(e)}") + + async def get_guardrail_by_name_from_db( + self, guardrail_name: str, prisma_client: PrismaClient + ) -> Optional[Guardrail]: + """ + Get a guardrail by its name from the database + """ + try: + guardrail = await prisma_client.db.litellm_guardrailstable.find_unique( + where={"guardrail_name": guardrail_name} + ) + + if not guardrail: + return None + + return Guardrail(**(dict(guardrail))) + except Exception as e: + raise Exception(f"Error getting guardrail from DB: {str(e)}") diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index c2a8bd11d0..1d6f3b5211 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -470,4 +470,14 @@ model LiteLLM_ManagedVectorStoresTable { created_at DateTime @default(now()) updated_at DateTime @updatedAt litellm_credential_name String? +} + +// Guardrails table for storing guardrail configurations +model LiteLLM_GuardrailsTable { + guardrail_id String @id @default(uuid()) + guardrail_name String @unique + litellm_params Json + guardrail_info Json? + created_at DateTime @default(now()) + updated_at DateTime @updatedAt } \ No newline at end of file diff --git a/schema.prisma b/schema.prisma index c2a8bd11d0..1d6f3b5211 100644 --- a/schema.prisma +++ b/schema.prisma @@ -470,4 +470,14 @@ model LiteLLM_ManagedVectorStoresTable { created_at DateTime @default(now()) updated_at DateTime @updatedAt litellm_credential_name String? +} + +// Guardrails table for storing guardrail configurations +model LiteLLM_GuardrailsTable { + guardrail_id String @id @default(uuid()) + guardrail_name String @unique + litellm_params Json + guardrail_info Json? + created_at DateTime @default(now()) + updated_at DateTime @updatedAt } \ No newline at end of file diff --git a/ui/litellm-dashboard/out/assets/logos/presidio.png b/ui/litellm-dashboard/out/assets/logos/presidio.png new file mode 100644 index 0000000000..cce9139017 Binary files /dev/null and b/ui/litellm-dashboard/out/assets/logos/presidio.png differ diff --git a/ui/litellm-dashboard/public/assets/logos/presidio.png b/ui/litellm-dashboard/public/assets/logos/presidio.png new file mode 100644 index 0000000000..cce9139017 Binary files /dev/null and b/ui/litellm-dashboard/public/assets/logos/presidio.png differ diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx new file mode 100644 index 0000000000..2c8f6108d6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx @@ -0,0 +1,53 @@ +export enum GuardrailProviders { + PresidioPII = "Presidio PII", + Aporia = "Aporia", + AimSecurity = "Aim Security", + Bedrock = "Amazon Bedrock", + GuardrailsAI = "Guardrails.ai", + LakeraAI = "Lakera AI", + Custom = "Custom Guardrail", + PromptInjection = "Prompt Injection Detection", +} + +export const guardrail_provider_map: Record = { + Aporia: "aporia", + AimSecurity: "aim", + Bedrock: "bedrock", + GuardrailsAI: "guardrails_ai", + LakeraAI: "lakera", + PromptInjection: "detect_prompt_injection", + PresidioPII: "presidio", +}; + +const asset_logos_folder = '../ui/assets/logos/'; + +export const guardrailLogoMap: Record = { + [GuardrailProviders.Aporia]: `${asset_logos_folder}aporia.svg`, + [GuardrailProviders.AimSecurity]: `${asset_logos_folder}aim.svg`, + [GuardrailProviders.Bedrock]: `${asset_logos_folder}bedrock.svg`, + [GuardrailProviders.GuardrailsAI]: `${asset_logos_folder}guardrails_ai.svg`, + [GuardrailProviders.LakeraAI]: `${asset_logos_folder}lakera.svg`, + [GuardrailProviders.PromptInjection]: `${asset_logos_folder}prompt_injection.svg`, + [GuardrailProviders.PresidioPII]: `${asset_logos_folder}presidio.svg` +}; + +export const getGuardrailLogoAndName = (guardrailValue: string): { logo: string, displayName: string } => { + if (!guardrailValue) { + return { logo: "", displayName: "-" }; + } + + // Find the enum key by matching guardrail_provider_map values + const enumKey = Object.keys(guardrail_provider_map).find( + key => guardrail_provider_map[key].toLowerCase() === guardrailValue.toLowerCase() + ); + + if (!enumKey) { + return { logo: "", displayName: guardrailValue }; + } + + // Get the display name from GuardrailProviders enum and logo from map + const displayName = GuardrailProviders[enumKey as keyof typeof GuardrailProviders]; + const logo = guardrailLogoMap[displayName as keyof typeof guardrailLogoMap]; + + return { logo, displayName }; +}; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 02ca8ff81e..4af134fa1a 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -4303,6 +4303,35 @@ export const getGuardrailsList = async (accessToken: String) => { } }; +export const createGuardrailCall = async (accessToken: string, guardrailData: any) => { + try { + const url = proxyBaseUrl ? `${proxyBaseUrl}/guardrails` : `/guardrails`; + + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + guardrail: guardrailData + }), + }); + + if (!response.ok) { + const errorData = await response.text(); + handleError(errorData); + throw new Error(errorData); + } + + const data = await response.json(); + console.log("Create guardrail response:", data); + return data; + } catch (error) { + console.error("Failed to create guardrail:", error); + throw error; + } +}; export const uiSpendLogDetailsCall = async ( accessToken: string,