From 1f0cc725316a8410f0e20e10297442ce84ec2022 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 22 Aug 2024 17:24:42 -0700
Subject: [PATCH] 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 24d826f4f..f8393be9d 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 496ae1710..8ca4f37fd 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 d0ed9a699..6b831876f 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 7e9ff613a..34f14186e 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)