From 46555ab78b3fdfe15edcb271bc51379edfba9584 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 29 Jul 2024 14:13:08 -0700
Subject: [PATCH] test - bedrock guardrailConfig

---
 litellm/tests/test_bedrock_completion.py | 33 ++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py
index 220e3f62f01..498cfb5d894 100644
--- a/litellm/tests/test_bedrock_completion.py
+++ b/litellm/tests/test_bedrock_completion.py
@@ -81,6 +81,39 @@ def test_completion_bedrock_claude_completion_auth():
 # test_completion_bedrock_claude_completion_auth()
 
 
+def test_completion_bedrock_guardrails():
+    import os
+
+    litellm.set_verbose = True
+
+    try:
+        response = completion(
+            model="anthropic.claude-v2",
+            messages=[
+                {
+                    "content": "where do i buy coffee from? ",
+                    "role": "user",
+                }
+            ],
+            max_tokens=10,
+            guardrailConfig={
+                "guardrailIdentifier": "ff6ujrregl1q",
+                "guardrailVersion": "DRAFT",
+                "trace": "disabled",
+            },
+        )
+        # Add any assertions here to check the response
+        print(response)
+        assert (
+            "Sorry, the model cannot answer this question. coffee guardrail applied"
+            in response.choices[0].message.content
+        )
+    except RateLimitError:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
 def test_completion_bedrock_claude_2_1_completion_auth():
     print("calling bedrock claude 2.1 completion params auth")
     import os