scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge pull request #5339 from BerriAI/litellm_add_bedrock_guardrails

Browse source 
[Feat-Proxy] add bedrock guardrails support
This commit is contained in:
Ishaan Jaff 2024-08-22 17:28:57 -07:00 • committed by GitHub
commit 2d57fab79f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 501 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
.circleci/config.yml
View file

@ -320,6 +320,9 @@ jobs:
-e APORIA_API_BASE_2=$APORIA_API_BASE_2 \
-e APORIA_API_KEY_2=$APORIA_API_KEY_2 \
-e APORIA_API_BASE_1=$APORIA_API_BASE_1 \
-e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
-e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \
-e AWS_REGION_NAME=$AWS_REGION_NAME \
-e APORIA_API_KEY_1=$APORIA_API_KEY_1 \
--name my-app \
-v $(pwd)/litellm/proxy/example_config_yaml/otel_test_config.yaml:/app/config.yaml \

135
docs/my-website/docs/proxy/guardrails/bedrock.md Normal file
View file

@ -0,0 +1,135 @@
import Image from '@theme/IdealImage';
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Bedrock
## Quick Start
### 1. Define Guardrails on your LiteLLM config.yaml
Define your guardrails under the `guardrails` section
```yaml
model_list:
- model_name: gpt-3.5-turbo
litellm_params:
model: openai/gpt-3.5-turbo
api_key: os.environ/OPENAI_API_KEY
guardrails:
- guardrail_name: "bedrock-pre-guard"
litellm_params:
guardrail: bedrock # supported values: "aporia", "bedrock", "lakera"
mode: "during_call"
guardrailIdentifier: ff6ujrregl1q # your guardrail ID on bedrock
guardrailVersion: "DRAFT" # your guardrail version on bedrock
```
#### Supported values for `mode`
- `pre_call` Run **before** LLM call, on **input**
- `post_call` Run **after** LLM call, on **input & output**
- `during_call` Run **during** LLM call, on **input** Same as `pre_call` but runs in parallel as LLM call. Response not returned until guardrail check completes
### 2. Start LiteLLM Gateway
```shell
litellm --config config.yaml --detailed_debug
```
### 3. Test request
**[Langchain, OpenAI SDK Usage Examples](../proxy/user_keys##request-format)**
<Tabs>
<TabItem label="Unsuccessful call" value = "not-allowed">
Expect this to fail since since `ishaan@berri.ai` in the request is PII
```shell
curl -i http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-npnwjPQciVRok5yNZgKmFQ" \
-d '{
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "hi my email is ishaan@berri.ai"}
],
"guardrails": ["bedrock-guard"]
}'
```
Expected response on failure
```shell
{
"error": {
"message": {
"error": "Violated guardrail policy",
"bedrock_guardrail_response": {
"action": "GUARDRAIL_INTERVENED",
"assessments": [
{
"topicPolicy": {
"topics": [
{
"action": "BLOCKED",
"name": "Coffee",
"type": "DENY"
}
]
}
}
],
"blockedResponse": "Sorry, the model cannot answer this question. coffee guardrail applied ",
"output": [
{
"text": "Sorry, the model cannot answer this question. coffee guardrail applied "
}
],
"outputs": [
{
"text": "Sorry, the model cannot answer this question. coffee guardrail applied "
}
],
"usage": {
"contentPolicyUnits": 0,
"contextualGroundingPolicyUnits": 0,
"sensitiveInformationPolicyFreeUnits": 0,
"sensitiveInformationPolicyUnits": 0,
"topicPolicyUnits": 1,
"wordPolicyUnits": 0
}
}
},
"type": "None",
"param": "None",
"code": "400"
}
}
```
</TabItem>
<TabItem label="Successful Call " value = "allowed">
```shell
curl -i http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-npnwjPQciVRok5yNZgKmFQ" \
-d '{
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "hi what is the weather"}
],
"guardrails": ["bedrock-guard"]
}'
```
</TabItem>
</Tabs>

2
docs/my-website/sidebars.js
View file

@ -54,7 +54,7 @@ const sidebars = {
{
type: "category",
label: "🛡️ [Beta] Guardrails",
items: ["proxy/guardrails/quick_start", "proxy/guardrails/aporia_api", "proxy/guardrails/lakera_ai"],
items: ["proxy/guardrails/quick_start", "proxy/guardrails/aporia_api", "proxy/guardrails/lakera_ai", "proxy/guardrails/bedrock"],
},
{
type: "category",

8
litellm/proxy/example_config_yaml/otel_test_config.yaml
View file

@ -21,4 +21,10 @@ guardrails:
guardrail: aporia # supported values: "aporia", "bedrock", "lakera"
mode: "post_call"
api_key: os.environ/APORIA_API_KEY_2
api_base: os.environ/APORIA_API_BASE_2
api_base: os.environ/APORIA_API_BASE_2
- guardrail_name: "bedrock-pre-guard"
litellm_params:
guardrail: bedrock # supported values: "aporia", "bedrock", "lakera"
mode: "pre_call"
guardrailIdentifier: ff6ujrregl1q
guardrailVersion: "DRAFT"

289
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py Normal file
View file

@ -0,0 +1,289 @@
# +-------------------------------------------------------------+
#
# Use Bedrock Guardrails for your LLM calls
#
# +-------------------------------------------------------------+
# Thank you users! We ❤️ you! - Krrish & Ishaan
import os
import sys
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import asyncio
import json
import sys
import traceback
import uuid
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional, Union
import aiohttp
import httpx
from fastapi import HTTPException
import litellm
from litellm import get_secret
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.logging_utils import (
convert_litellm_response_object_to_str,
)
from litellm.llms.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
_get_async_httpx_client,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata
from litellm.types.guardrails import (
BedrockContentItem,
BedrockRequest,
BedrockTextContent,
GuardrailEventHooks,
)
GUARDRAIL_NAME = "bedrock"
class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
def __init__(
self,
guardrailIdentifier: Optional[str] = None,
guardrailVersion: Optional[str] = None,
**kwargs,
):
self.async_handler = _get_async_httpx_client()
self.guardrailIdentifier = guardrailIdentifier
self.guardrailVersion = guardrailVersion
# store kwargs as optional_params
self.optional_params = kwargs
super().__init__(**kwargs)
def convert_to_bedrock_format(
self,
messages: Optional[List[Dict[str, str]]] = None,
response: Optional[Union[Any, litellm.ModelResponse]] = None,
) -> BedrockRequest:
bedrock_request: BedrockRequest = BedrockRequest(source="INPUT")
bedrock_request_content: List[BedrockContentItem] = []
if messages:
for message in messages:
content = message.get("content")
if isinstance(content, str):
bedrock_content_item = BedrockContentItem(
text=BedrockTextContent(text=content)
)
bedrock_request_content.append(bedrock_content_item)
bedrock_request["content"] = bedrock_request_content
if response:
bedrock_request["source"] = "OUTPUT"
if isinstance(response, litellm.ModelResponse):
for choice in response.choices:
if isinstance(choice, litellm.Choices):
if choice.message.content and isinstance(
choice.message.content, str
):
bedrock_content_item = BedrockContentItem(
text=BedrockTextContent(text=choice.message.content)
)
bedrock_request_content.append(bedrock_content_item)
bedrock_request["content"] = bedrock_request_content
return bedrock_request
#### CALL HOOKS - proxy only ####
def _load_credentials(
self,
):
try:
from botocore.credentials import Credentials
except ImportError as e:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
## CREDENTIALS ##
# pop aws_secret_access_key, aws_access_key_id, aws_session_token, aws_region_name from kwargs, since completion calls fail with them
aws_secret_access_key = self.optional_params.pop("aws_secret_access_key", None)
aws_access_key_id = self.optional_params.pop("aws_access_key_id", None)
aws_session_token = self.optional_params.pop("aws_session_token", None)
aws_region_name = self.optional_params.pop("aws_region_name", None)
aws_role_name = self.optional_params.pop("aws_role_name", None)
aws_session_name = self.optional_params.pop("aws_session_name", None)
aws_profile_name = self.optional_params.pop("aws_profile_name", None)
aws_bedrock_runtime_endpoint = self.optional_params.pop(
"aws_bedrock_runtime_endpoint", None
) # https://bedrock-runtime.{region_name}.amazonaws.com
aws_web_identity_token = self.optional_params.pop(
"aws_web_identity_token", None
)
aws_sts_endpoint = self.optional_params.pop("aws_sts_endpoint", None)
### SET REGION NAME ###
if aws_region_name is None:
# check env #
litellm_aws_region_name = get_secret("AWS_REGION_NAME", None)
if litellm_aws_region_name is not None and isinstance(
litellm_aws_region_name, str
):
aws_region_name = litellm_aws_region_name
standard_aws_region_name = get_secret("AWS_REGION", None)
if standard_aws_region_name is not None and isinstance(
standard_aws_region_name, str
):
aws_region_name = standard_aws_region_name
if aws_region_name is None:
aws_region_name = "us-west-2"
credentials: Credentials = self.get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_region_name=aws_region_name,
aws_session_name=aws_session_name,
aws_profile_name=aws_profile_name,
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
)
return credentials, aws_region_name
def _prepare_request(
self,
credentials,
data: BedrockRequest,
optional_params: dict,
aws_region_name: str,
extra_headers: Optional[dict] = None,
):
try:
import boto3
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import Credentials
except ImportError as e:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name)
api_base = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply"
encoded_data = json.dumps(data).encode("utf-8")
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
request = AWSRequest(
method="POST", url=api_base, data=encoded_data, headers=headers
)
sigv4.add_auth(request)
prepped_request = request.prepare()
return prepped_request
async def make_bedrock_api_request(
self, kwargs: dict, response: Optional[Union[Any, litellm.ModelResponse]] = None
):
credentials, aws_region_name = self._load_credentials()
request_data: BedrockRequest = self.convert_to_bedrock_format(
messages=kwargs.get("messages"), response=response
)
prepared_request = self._prepare_request(
credentials=credentials,
data=request_data,
optional_params=self.optional_params,
aws_region_name=aws_region_name,
)
verbose_proxy_logger.debug(
"Bedrock AI request body: %s, url %s, headers: %s",
request_data,
prepared_request.url,
prepared_request.headers,
)
_json_data = json.dumps(request_data) # type: ignore
response = await self.async_handler.post(
url=prepared_request.url,
json=request_data, # type: ignore
headers=prepared_request.headers,
)
verbose_proxy_logger.debug("Bedrock AI response: %s", response.text)
if response.status_code == 200:
# check if the response was flagged
_json_response = response.json()
if _json_response.get("action") == "GUARDRAIL_INTERVENED":
raise HTTPException(
status_code=400,
detail={
"error": "Violated guardrail policy",
"bedrock_guardrail_response": _json_response,
},
)
else:
verbose_proxy_logger.error(
"Bedrock AI: error in response. Status code: %s, response: %s",
response.status_code,
response.text,
)
async def async_moderation_hook( ### 👈 KEY CHANGE ###
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal["completion", "embeddings", "image_generation"],
):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
from litellm.types.guardrails import GuardrailEventHooks
event_type: GuardrailEventHooks = GuardrailEventHooks.during_call
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
return
new_messages: Optional[List[dict]] = data.get("messages")
if new_messages is not None:
await self.make_bedrock_api_request(kwargs=data)
add_guardrail_to_applied_guardrails_header(
request_data=data, guardrail_name=self.guardrail_name
)
else:
verbose_proxy_logger.warning(
"Bedrock AI: not running guardrail. No messages in data"
)
pass
async def async_post_call_success_hook(
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
response,
):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
from litellm.types.guardrails import GuardrailEventHooks
if (
self.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.post_call
)
is not True
):
return
new_messages: Optional[List[dict]] = data.get("messages")
if new_messages is not None:
await self.make_bedrock_api_request(kwargs=data, response=response)
add_guardrail_to_applied_guardrails_header(
request_data=data, guardrail_name=self.guardrail_name
)
else:
verbose_proxy_logger.warning(
"Bedrock AI: not running guardrail. No messages in data"
)

18
litellm/proxy/guardrails/init_guardrails.py
View file

@ -96,8 +96,10 @@ def init_guardrails_v2(all_guardrails: dict):
litellm_params = LitellmParams(
guardrail=litellm_params_data["guardrail"],
mode=litellm_params_data["mode"],
api_key=litellm_params_data["api_key"],
api_base=litellm_params_data["api_base"],
api_key=litellm_params_data.get("api_key"),
api_base=litellm_params_data.get("api_base"),
guardrailIdentifier=litellm_params_data.get("guardrailIdentifier"),
guardrailVersion=litellm_params_data.get("guardrailVersion"),
)
if (
@ -134,6 +136,18 @@ def init_guardrails_v2(all_guardrails: dict):
event_hook=litellm_params["mode"],
)
litellm.callbacks.append(_aporia_callback) # type: ignore
if litellm_params["guardrail"] == "bedrock":
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockGuardrail,
)
_bedrock_callback = BedrockGuardrail(
guardrail_name=guardrail["guardrail_name"],
event_hook=litellm_params["mode"],
guardrailIdentifier=litellm_params["guardrailIdentifier"],
guardrailVersion=litellm_params["guardrailVersion"],
)
litellm.callbacks.append(_bedrock_callback) # type: ignore
elif litellm_params["guardrail"] == "lakera":
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import (
lakeraAI_Moderation,

20
litellm/proxy/proxy_config.yaml
View file

@ -1,20 +1,14 @@
model_list:
- model_name: gpt-4
- model_name: fake-openai-endpoint
litellm_params:
model: openai/fake
api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/
litellm_settings:
success_callback: ["prometheus"]
failure_callback: ["prometheus"]
guardrails:
- guardrail_name: "lakera-pre-guard"
- guardrail_name: "bedrock-pre-guard"
litellm_params:
guardrail: lakera # supported values: "aporia", "bedrock", "lakera"
mode: "during_call"
api_key: os.environ/LAKERA_API_KEY
api_base: os.environ/LAKERA_API_BASE
category_thresholds:
prompt_injection: 0.1
jailbreak: 0.1
guardrail: bedrock # supported values: "aporia", "bedrock", "lakera"
mode: "post_call"
guardrailIdentifier: ff6ujrregl1q
guardrailVersion: "DRAFT"

21
litellm/types/guardrails.py
View file

@ -1,5 +1,5 @@
from enum import Enum
from typing import Dict, List, Optional, TypedDict
from typing import Dict, List, Literal, Optional, TypedDict
from pydantic import BaseModel, ConfigDict
from typing_extensions import Required, TypedDict
@ -76,8 +76,14 @@ class LitellmParams(TypedDict, total=False):
mode: str
api_key: str
api_base: Optional[str]
# Lakera specific params
category_thresholds: Optional[LakeraCategoryThresholds]
# Bedrock specific params
guardrailIdentifier: Optional[str]
guardrailVersion: Optional[str]
class Guardrail(TypedDict):
guardrail_name: str
@ -92,3 +98,16 @@ class GuardrailEventHooks(str, Enum):
pre_call = "pre_call"
post_call = "post_call"
during_call = "during_call"
class BedrockTextContent(TypedDict, total=False):
text: str
class BedrockContentItem(TypedDict, total=False):
text: BedrockTextContent
class BedrockRequest(TypedDict, total=False):
source: Literal["INPUT", "OUTPUT"]
content: List[BedrockContentItem]

23
tests/otel_tests/test_guardrails.py
View file

@ -144,6 +144,7 @@ async def test_no_llm_guard_triggered():
assert "x-litellm-applied-guardrails" not in headers
@pytest.mark.asyncio
async def test_guardrails_with_api_key_controls():
"""
@ -194,3 +195,25 @@ async def test_guardrails_with_api_key_controls():
except Exception as e:
print(e)
assert "Aporia detected and blocked PII" in str(e)
@pytest.mark.asyncio
async def test_bedrock_guardrail_triggered():
"""
- Tests a request where our bedrock guardrail should be triggered
- Assert that the guardrails applied are returned in the response headers
"""
async with aiohttp.ClientSession() as session:
try:
response, headers = await chat_completion(
session,
"sk-1234",
model="fake-openai-endpoint",
messages=[{"role": "user", "content": f"Hello do you like coffee?"}],
guardrails=["bedrock-pre-guard"],
)
pytest.fail("Should have thrown an exception")
except Exception as e:
print(e)
assert "GUARDRAIL_INTERVENED" in str(e)
assert "Violated guardrail policy" in str(e)