(feat) /guardrails/list show guardrail info params (#7442)

* add GuardrailInfoResponse

* add list_guardrails

* test_get_guardrails_list_response
This commit is contained in:
Ishaan Jaff 2024-12-27 14:35:00 -08:00 committed by GitHub
parent 8ef5b4e94c
commit 3e7794d880
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 128 additions and 8 deletions

View file

@ -8,24 +8,71 @@ from fastapi import APIRouter, Depends, HTTPException, status
from litellm.proxy._types import CommonProxyErrors from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.guardrails import GuardrailInfoResponse, ListGuardrailsResponse
#### GUARDRAILS ENDPOINTS #### #### GUARDRAILS ENDPOINTS ####
router = APIRouter() router = APIRouter()
def _get_guardrail_names_from_config(guardrails_config: List[Dict]) -> List[str]: def _get_guardrails_list_response(
return [guardrail["guardrail_name"] for guardrail in guardrails_config] guardrails_config: List[Dict],
) -> ListGuardrailsResponse:
"""
Helper function to get the guardrails list response
"""
guardrail_configs: List[GuardrailInfoResponse] = []
for guardrail in guardrails_config:
guardrail_configs.append(
GuardrailInfoResponse(
guardrail_name=guardrail.get("guardrail_name"),
guardrail_info=guardrail.get("guardrail_info"),
)
)
return ListGuardrailsResponse(guardrails=guardrail_configs)
@router.get( @router.get(
"/guardrails/list", "/guardrails/list",
tags=["Guardrails"], tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)], dependencies=[Depends(user_api_key_auth)],
response_model=ListGuardrailsResponse,
) )
async def list_guardrails(): async def list_guardrails():
""" """
Enterprise Feature
List the guardrails that are available on the proxy server List the guardrails that are available on the proxy server
👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start)
Example Request:
```bash
curl -X GET "http://localhost:4000/guardrails/list" -H "Authorization: Bearer <your_api_key>"
```
Example Response:
```json
{
"guardrails": [
{
"guardrail_name": "bedrock-pre-guard",
"guardrail_info": {
"params": [
{
"name": "toxicity_score",
"type": "float",
"description": "Score between 0-1 indicating content toxicity level"
},
{
"name": "pii_detection",
"type": "boolean"
}
]
}
}
]
}
```
""" """
from litellm.proxy.proxy_server import premium_user, proxy_config from litellm.proxy.proxy_server import premium_user, proxy_config
@ -47,4 +94,4 @@ async def list_guardrails():
detail={"error": "No guardrails found in config"}, detail={"error": "No guardrails found in config"},
) )
return _get_guardrail_names_from_config(config["guardrails"]) return _get_guardrails_list_response(_guardrails_config)

View file

@ -11,9 +11,22 @@ model_list:
litellm_params: litellm_params:
model: bedrock/* model: bedrock/*
litellm_settings:
callbacks: ["otel"]
# callback_settings: guardrails:
# otel: - guardrail_name: "bedrock-pre-guard"
# message_logging: False litellm_params:
guardrail: bedrock # supported values: "aporia", "bedrock", "lakera"
mode: "post_call"
guardrailIdentifier: ff6ujrregl1q
guardrailVersion: "DRAFT"
guardrail_info:
params:
- name: "toxicity_score"
type: "float"
description: "Score between 0-1 indicating content toxicity level"
- name: "pii_detection"
type: "boolean"
litellm_settings:
callbacks: ["datadog"]

View file

@ -136,3 +136,12 @@ class BedrockRequest(TypedDict, total=False):
class DynamicGuardrailParams(TypedDict): class DynamicGuardrailParams(TypedDict):
extra_body: Dict[str, Any] extra_body: Dict[str, Any]
class GuardrailInfoResponse(BaseModel):
guardrail_name: Optional[str]
guardrail_info: Optional[Dict] # This will contain all other fields
class ListGuardrailsResponse(BaseModel):
guardrails: List[GuardrailInfoResponse]

View file

@ -1,3 +1,7 @@
"""
Test custom guardrail + unit tests for guardrails
"""
import io import io
import os import os
import sys import sys
@ -29,6 +33,8 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata
from litellm.types.guardrails import GuardrailEventHooks from litellm.types.guardrails import GuardrailEventHooks
from litellm.proxy.guardrails.guardrail_endpoints import _get_guardrails_list_response
from litellm.types.guardrails import GuardrailInfoResponse, ListGuardrailsResponse
def test_get_guardrail_from_metadata(): def test_get_guardrail_from_metadata():
@ -143,3 +149,48 @@ def test_get_guardrail_dynamic_request_body_params():
} }
} }
assert guardrail.get_guardrail_dynamic_request_body_params(data) == {} assert guardrail.get_guardrail_dynamic_request_body_params(data) == {}
def test_get_guardrails_list_response():
# Test case 1: Valid guardrails config
sample_config = [
{
"guardrail_name": "test-guard",
"guardrail_info": {
"params": [
{
"name": "toxicity_score",
"type": "float",
"description": "Score between 0-1",
}
]
},
}
]
response = _get_guardrails_list_response(sample_config)
assert isinstance(response, ListGuardrailsResponse)
assert len(response.guardrails) == 1
assert response.guardrails[0].guardrail_name == "test-guard"
assert response.guardrails[0].guardrail_info == {
"params": [
{
"name": "toxicity_score",
"type": "float",
"description": "Score between 0-1",
}
]
}
# Test case 2: Empty guardrails config
empty_response = _get_guardrails_list_response([])
assert isinstance(empty_response, ListGuardrailsResponse)
assert len(empty_response.guardrails) == 0
# Test case 3: Missing optional fields
minimal_config = [{"guardrail_name": "minimal-guard"}]
minimal_response = _get_guardrails_list_response(minimal_config)
assert isinstance(minimal_response, ListGuardrailsResponse)
assert len(minimal_response.guardrails) == 1
assert minimal_response.guardrails[0].guardrail_name == "minimal-guard"
assert minimal_response.guardrails[0].guardrail_info is None