(test) catch bedrock RateLimitErrors

litellm/tests/test_completion.py

@@ -11,6 +11,7 @@ sys.path.insert(
 import pytest
 import litellm
 from litellm import embedding, completion, text_completion, completion_cost
+from litellm import RateLimitError
 
 user_message = "Write a short poem about the sky"
 messages = [{"content": user_message, "role": "user"}]
@@ -780,6 +781,8 @@ def test_completion_bedrock_titan():
         )
         # Add any assertions here to check the response
         print(response)
+    except RateLimitError:
+        pass
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 # test_completion_bedrock_titan()
@@ -796,6 +799,8 @@ def test_completion_bedrock_claude():
         )
         # Add any assertions here to check the response
         print(response)
+    except RateLimitError:
+        pass
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 # test_completion_bedrock_claude()
@@ -831,7 +836,8 @@ def test_completion_bedrock_claude_completion_auth():
         os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
         os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
         os.environ["AWS_REGION_NAME"] = aws_region_name
-    
+    except RateLimitError:
+        pass
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 # test_completion_bedrock_claude_completion_auth()
@@ -872,7 +878,8 @@ def test_completion_bedrock_claude_external_client_auth():
         os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
         os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
         os.environ["AWS_REGION_NAME"] = aws_region_name
-
+    except RateLimitError:
+        pass
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 # test_completion_bedrock_claude_external_client_auth()
@@ -890,6 +897,8 @@ def test_completion_bedrock_claude_stream():
         print(response)
         for chunk in response:
             print(chunk)
+    except RateLimitError:
+        pass
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 # test_completion_bedrock_claude_stream()
@@ -906,6 +915,8 @@ def test_completion_bedrock_ai21():
         )
         # Add any assertions here to check the response 
         print(response)
+    except RateLimitError:
+        pass
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 

litellm/tests/test_provider_specific_config.py

@@ -10,6 +10,7 @@ sys.path.insert(
 )  # Adds the parent directory to the system path
 import litellm
 from litellm import completion
+from litellm import RateLimitError
 
 #  Huggingface - Expensive to deploy models and keep them running. Maybe we can try doing this via baseten?? 
 # def hf_test_completion_tgi():
@@ -440,6 +441,8 @@ def bedrock_test_completion():
         print(f"response_2_text: {response_2_text}")
 
         assert len(response_2_text) < len(response_1_text)
+    except RateLimitError:
+        pass
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 

litellm/tests/test_streaming.py

@@ -9,7 +9,7 @@ sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
 import litellm
-from litellm import completion, acompletion, AuthenticationError, InvalidRequestError
+from litellm import completion, acompletion, AuthenticationError, InvalidRequestError, RateLimitError
 
 litellm.logging = False
 litellm.set_verbose = False
@@ -656,37 +656,38 @@ def test_completion_replicate_stream_bad_key():
 
 # test_completion_replicate_stream_bad_key()
 
-# def test_completion_bedrock_ai21_stream():
-# bedrock is currently failing tests
-#     try:
-#         litellm.set_verbose=True
-#         response = completion(
-#             model="bedrock/amazon.titan-text-express-v1", 
-#             messages=[{"role": "user", "content": "Be as verbose as possible and give as many details as possible, how does a court case get to the Supreme Court?"}],
-#             temperature=1,
-#             max_tokens=4096,
-#             stream=True,
-#         )
-#         print(response)
-#         complete_response = ""
-#         has_finish_reason = False
-#         # Add any assertions here to check the response
-#         for idx, chunk in enumerate(response):
-#             # print
-#             chunk, finished = streaming_format_tests(idx, chunk)
-#             has_finish_reason = finished
-#             complete_response += chunk
-#             if finished:
-#                 break
-#         if has_finish_reason is False:
-#             raise Exception("finish reason not set for last chunk")
-#         if complete_response.strip() == "": 
-#             raise Exception("Empty response received")
-#         print(f"completion_response: {complete_response}")
-#     except Exception as e:
-#         pytest.fail(f"Error occurred: {e}")
+def test_completion_bedrock_ai21_stream():
+    try:
+        litellm.set_verbose=True
+        response = completion(
+            model="bedrock/anthropic.claude-instant-v1", 
+            messages=[{"role": "user", "content": "Be as verbose as possible and give as many details as possible, how does a court case get to the Supreme Court?"}],
+            temperature=1,
+            max_tokens=4096,
+            stream=True,
+        )
+        print(response)
+        complete_response = ""
+        has_finish_reason = False
+        # Add any assertions here to check the response
+        for idx, chunk in enumerate(response):
+            # print
+            chunk, finished = streaming_format_tests(idx, chunk)
+            has_finish_reason = finished
+            complete_response += chunk
+            if finished:
+                break
+        if has_finish_reason is False:
+            raise Exception("finish reason not set for last chunk")
+        if complete_response.strip() == "": 
+            raise Exception("Empty response received")
+        print(f"completion_response: {complete_response}")
+    except RateLimitError:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
 
-# test_completion_bedrock_ai21_stream() 
+test_completion_bedrock_ai21_stream() 
 
 
 def test_completion_sagemaker_stream():