Merge pull request #1919 from BerriAI/litellm_bedrock_set_timeouts

[FEAT] Bedrock set timeouts on litellm.completion

litellm/llms/bedrock.py

@@ -526,6 +526,7 @@ def completion(
     optional_params=None,
     litellm_params=None,
     logger_fn=None,
+    timeout=None,
 ):
     exception_mapping_worked = False
     try:
@@ -553,6 +554,7 @@ def completion(
                 aws_role_name=aws_role_name,
                 aws_session_name=aws_session_name,
                 aws_profile_name=aws_profile_name,
+                timeout=timeout,
             )
 
         model = model

litellm/main.py

@@ -1590,6 +1590,7 @@ def completion(
                 logger_fn=logger_fn,
                 encoding=encoding,
                 logging_obj=logging,
+                timeout=timeout,
             )
 
             if "stream" in optional_params and optional_params["stream"] == True:

litellm/tests/test_timeout.py

@@ -37,6 +37,28 @@ def test_timeout():
 # test_timeout()
 
 
+def test_bedrock_timeout():
+    # this Will Raise a timeout
+    litellm.set_verbose = True
+    try:
+        response = litellm.completion(
+            model="bedrock/anthropic.claude-instant-v1",
+            timeout=0.01,
+            messages=[{"role": "user", "content": "hello, write a 20 pg essay"}],
+        )
+        pytest.fail("Did not raise error `openai.APITimeoutError`")
+    except openai.APITimeoutError as e:
+        print(
+            "Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e
+        )
+        print(type(e))
+        pass
+    except Exception as e:
+        pytest.fail(
+            f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}"
+        )
+
+
 def test_hanging_request_azure():
     litellm.set_verbose = True
     import asyncio

litellm/utils.py

@@ -6643,6 +6643,13 @@ def exception_type(
                         llm_provider="bedrock",
                         response=original_exception.response,
                     )
+                if "Connect timeout on endpoint URL" in error_str:
+                    exception_mapping_worked = True
+                    raise Timeout(
+                        message=f"BedrockException: Timeout Error - {error_str}",
+                        model=model,
+                        llm_provider="bedrock",
+                    )
                 if hasattr(original_exception, "status_code"):
                     if original_exception.status_code == 500:
                         exception_mapping_worked = True