test bedrock guardrails

.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 \

litellm/proxy/example_config_yaml/otel_test_config.yaml

@@ -22,3 +22,9 @@ guardrails:
       mode: "post_call"
       api_key: os.environ/APORIA_API_KEY_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"

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

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)