tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 14:22:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Feat] Add endpoints for adding, deleting, editing guardrails in DB (#10833)

Browse Source 
* feat: add DB Add, Edit, Delete for Guardrails

* feat: endpoints for guardrail management

* add guardrail info helpers

* add presidio logo

* add createGuardrailCall
This commit is contained in:
Ishaan Jaff
2025-05-14 14:19:51 -07:00
committed by GitHub
parent 007524a972
commit ee1557afcd
9 changed files with 533 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+10
View File
@@ -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
}
litellm/proxy/guardrails/guardrail_endpoints.py
+272 -2
View File
@@ -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 <your_api_key>" \\
-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 <your_api_key>"
```
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 <your_api_key>" \\
-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 <your_api_key>"
```
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))
litellm/proxy/guardrails/guardrail_registry.py
+149 -1
View File
@@ -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)}")
litellm/proxy/schema.prisma
+10
View File
@@ -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
}
schema.prisma
+10
View File
@@ -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
}
ui/litellm-dashboard/out/assets/logos/presidio.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 61 KiB

ui/litellm-dashboard/public/assets/logos/presidio.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 61 KiB

ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx
+53
View File
@@ -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<string, string> = {
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<string, string> = {
[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 };
};
ui/litellm-dashboard/src/components/networking.tsx
+29
View File
@@ -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,