From d121029a56a94bc4b260d04e13be24a131478dce Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 14 Sep 2023 14:15:32 -0700
Subject: [PATCH] add bedrock tests

---
 litellm/tests/test_completion.py | 44 ++++++++++++++++++++++----------
 1 file changed, 30 insertions(+), 14 deletions(-)

diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index 929db53c0..03517e77d 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -634,20 +634,36 @@ def test_completion_sagemaker():
         pytest.fail(f"Error occurred: {e}")
 
 
-# def test_completion_bedrock():
-#     try:
-#         response = completion(
-#             model="bedrock/amazon.titan-tg1-large", 
-#             messages=messages,
-#             temperature=0.2,
-#             max_tokens=80,
-#             logger_fn=logger_fn
-#         )
-#         # Add any assertions here to check the response
-#         print(response)
-#     except Exception as e:
-#         pytest.fail(f"Error occurred: {e}")
-# test_completion_bedrock()
+def test_completion_bedrock_titan():
+    try:
+        response = completion(
+            model="bedrock/amazon.titan-tg1-large", 
+            messages=messages,
+            temperature=0.2,
+            max_tokens=20,
+            logger_fn=logger_fn
+        )
+        # Add any assertions here to check the response
+        print(response)
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+test_completion_bedrock_titan()
+
+
+def test_completion_bedrock_ai21():
+    try:
+        response = completion(
+            model="bedrock/ai21.j2-mid", 
+            messages=messages,
+            temperature=0.2,
+            max_tokens=20,
+            logger_fn=logger_fn
+        )
+        # Add any assertions here to check the response
+        print(response)
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+test_completion_bedrock_ai21()
 
 
 # test_completion_sagemaker()