From 14a6ce367d2d23057d8131dbecba127b3c345b69 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 15:40:58 -0700 Subject: [PATCH 1/5] add types for BedrockMessage --- litellm/types/guardrails.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 66c2a535ad6..13992beec58 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -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): + text: str + + +class BedrockContentItem(TypedDict): + text: BedrockTextContent + + +class BedrockMessage(TypedDict): + source: Literal["INPUT", "OUTPUT"] + content: List[BedrockContentItem] From 7d55047ab9f99926d6147cc2b6c448c25e4c684d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 16:09:55 -0700 Subject: [PATCH 2/5] add bedrock guardrails support --- .../guardrail_hooks/bedrock_guardrails.py | 273 ++++++++++++++++++ litellm/proxy/guardrails/init_guardrails.py | 18 +- litellm/proxy/proxy_config.yaml | 12 +- litellm/types/guardrails.py | 6 +- 4 files changed, 296 insertions(+), 13 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py new file mode 100644 index 00000000000..6c7ea4d906b --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -0,0 +1,273 @@ +# +-------------------------------------------------------------+ +# +# 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, + ) -> BedrockRequest: + bedrock_request: BedrockRequest = BedrockRequest(source="INPUT") + if messages: + bedrock_request_content: List[BedrockContentItem] = [] + 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 + + 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): + + credentials, aws_region_name = self._load_credentials() + request_data: BedrockRequest = self.convert_to_bedrock_format( + messages=kwargs.get("messages") + ) + 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 + + # """ + # Use this for the post call moderation with Guardrails + # """ + # event_type: GuardrailEventHooks = GuardrailEventHooks.post_call + # if self.should_run_guardrail(data=data, event_type=event_type) is not True: + # return + + # response_str: Optional[str] = convert_litellm_response_object_to_str(response) + # if response_str is not None: + # await self.make_bedrock_api_request( + # response_string=response_str, new_messages=data.get("messages", []) + # ) + + # add_guardrail_to_applied_guardrails_header( + # request_data=data, guardrail_name=self.guardrail_name + # ) + + # pass diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index ad99daf9556..f0e2a9e2eca 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -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, diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 65c7f70525c..d8e88cec704 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -6,13 +6,9 @@ model_list: api_base: https://exampleopenaiendpoint-production.up.railway.app/ guardrails: - - guardrail_name: "lakera-pre-guard" + - guardrail_name: "bedrock-pre-guard" litellm_params: - guardrail: lakera # supported values: "aporia", "bedrock", "lakera" + guardrail: bedrock # 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 - \ No newline at end of file + guardrailIdentifier: ff6ujrregl1q + guardrailVersion: "DRAFT" \ No newline at end of file diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 13992beec58..10f4be7e1eb 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -100,14 +100,14 @@ class GuardrailEventHooks(str, Enum): during_call = "during_call" -class BedrockTextContent(TypedDict): +class BedrockTextContent(TypedDict, total=False): text: str -class BedrockContentItem(TypedDict): +class BedrockContentItem(TypedDict, total=False): text: BedrockTextContent -class BedrockMessage(TypedDict): +class BedrockRequest(TypedDict, total=False): source: Literal["INPUT", "OUTPUT"] content: List[BedrockContentItem] From 499b6b33688dd43a9b1cf55c5ffa8370b5568dc1 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 16:25:22 -0700 Subject: [PATCH 3/5] doc bedrock guardrails --- .../docs/proxy/guardrails/bedrock.md | 135 ++++++++++++++++++ docs/my-website/sidebars.js | 2 +- 2 files changed, 136 insertions(+), 1 deletion(-) create mode 100644 docs/my-website/docs/proxy/guardrails/bedrock.md diff --git a/docs/my-website/docs/proxy/guardrails/bedrock.md b/docs/my-website/docs/proxy/guardrails/bedrock.md new file mode 100644 index 00000000000..ac8aa1c1b59 --- /dev/null +++ b/docs/my-website/docs/proxy/guardrails/bedrock.md @@ -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)** + + + + +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" + } +} + +``` + + + + + +```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"] + }' +``` + + + + + + diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index ab94ed5b421..b907a113045 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -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", From 9e3d573bcb6c14dc517ae9f930b952cd6b472698 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 16:34:43 -0700 Subject: [PATCH 4/5] add async_post_call_success_hook --- .../guardrail_hooks/bedrock_guardrails.py | 78 +++++++++++-------- litellm/proxy/proxy_config.yaml | 7 +- 2 files changed, 50 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 6c7ea4d906b..d11f58a3eab 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -67,10 +67,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): 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: - bedrock_request_content: List[BedrockContentItem] = [] for message in messages: content = message.get("content") if isinstance(content, str): @@ -80,7 +82,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): 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 #### @@ -172,11 +186,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return prepped_request - async def make_bedrock_api_request(self, kwargs: dict): + 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") + messages=kwargs.get("messages"), response=response ) prepared_request = self._prepare_request( credentials=credentials, @@ -242,32 +258,32 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) 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 + 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 - # """ - # Use this for the post call moderation with Guardrails - # """ - # event_type: GuardrailEventHooks = GuardrailEventHooks.post_call - # if self.should_run_guardrail(data=data, event_type=event_type) is not True: - # return + if ( + self.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.post_call + ) + is not True + ): + return - # response_str: Optional[str] = convert_litellm_response_object_to_str(response) - # if response_str is not None: - # await self.make_bedrock_api_request( - # response_string=response_str, new_messages=data.get("messages", []) - # ) - - # add_guardrail_to_applied_guardrails_header( - # request_data=data, guardrail_name=self.guardrail_name - # ) - - # pass + 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" + ) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index d8e88cec704..d0ed9a69904 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,14 +1,13 @@ model_list: - model_name: gpt-4 litellm_params: - model: openai/fake - api_key: fake-key - api_base: https://exampleopenaiendpoint-production.up.railway.app/ + model: openai/gpt-4 + 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" + mode: "post_call" guardrailIdentifier: ff6ujrregl1q guardrailVersion: "DRAFT" \ No newline at end of file From 1f0cc725316a8410f0e20e10297442ce84ec2022 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 17:24:42 -0700 Subject: [PATCH 5/5] test bedrock guardrails --- .circleci/config.yml | 3 +++ .../example_config_yaml/otel_test_config.yaml | 8 ++++++- litellm/proxy/proxy_config.yaml | 2 +- tests/otel_tests/test_guardrails.py | 23 +++++++++++++++++++ 4 files changed, 34 insertions(+), 2 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 24d826f4f61..f8393be9dff 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -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 \ diff --git a/litellm/proxy/example_config_yaml/otel_test_config.yaml b/litellm/proxy/example_config_yaml/otel_test_config.yaml index 496ae1710dc..8ca4f37fd6a 100644 --- a/litellm/proxy/example_config_yaml/otel_test_config.yaml +++ b/litellm/proxy/example_config_yaml/otel_test_config.yaml @@ -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 \ No newline at end of file + 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" \ No newline at end of file diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index d0ed9a69904..6b831876f04 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,5 +1,5 @@ model_list: - - model_name: gpt-4 + - model_name: fake-openai-endpoint litellm_params: model: openai/gpt-4 api_key: os.environ/OPENAI_API_KEY diff --git a/tests/otel_tests/test_guardrails.py b/tests/otel_tests/test_guardrails.py index 7e9ff613a4d..34f14186e13 100644 --- a/tests/otel_tests/test_guardrails.py +++ b/tests/otel_tests/test_guardrails.py @@ -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)