Added capability to customize read and connection timeout for bedrock models

autogen/oai/bedrock.py

@@ -56,6 +56,8 @@ def __init__(self, **kwargs: Any):
         self._aws_session_token = kwargs.get("aws_session_token", None)
         self._aws_region = kwargs.get("aws_region", None)
         self._aws_profile_name = kwargs.get("aws_profile_name", None)
+        self._aws_read_timeout = kwargs.get("aws_read_timeout", 60)
+        self._aws_connect_timeout = kwargs.get("aws_connect_timeout", 60)
 
         if not self._aws_access_key:
             self._aws_access_key = os.getenv("AWS_ACCESS_KEY_ID")
@@ -77,6 +79,8 @@ def __init__(self, **kwargs: Any):
             region_name=self._aws_region,
             signature_version="v4",
             retries={"max_attempts": self._retries, "mode": "standard"},
+            read_timeout=self._aws_read_timeout,
+            connect_timeout=self._aws_connect_timeout,
         )
 
         session = boto3.Session(

autogen/oai/client.py

@@ -485,7 +485,7 @@ def _configure_azure_openai(self, config: Dict[str, Any], openai_config: Dict[st
     def _configure_openai_config_for_bedrock(self, config: Dict[str, Any], openai_config: Dict[str, Any]) -> None:
         """Update openai_config with AWS credentials from config."""
         required_keys = ["aws_access_key", "aws_secret_key", "aws_region"]
-        optional_keys = ["aws_session_token", "aws_profile_name"]
+        optional_keys = ["aws_session_token", "aws_profile_name", "aws_read_timeout", "aws_connect_timeout"]
         for key in required_keys:
             if key in config:
                 openai_config[key] = config[key]

test/oai/test_bedrock.py

@@ -58,6 +58,8 @@ def test_parsing_params(bedrock_client):
         # "aws_secret_access_key": "test_secret_access_key",
         # "aws_session_token": "test_session_token",
         # "aws_profile_name": "test_profile_name",
+        # "aws_read_timeout": 120,
+        # "aws_connect_timeout: 30,
         "model": "anthropic.claude-3-sonnet-20240229-v1:0",
         "temperature": 0.8,
         "topP": 0.6,

website/docs/topics/non-openai-models/cloud-bedrock.ipynb

@@ -66,6 +66,8 @@
     "- aws_secret_key (or environment variable: AWS_SECRET_KEY)\n",
     "- aws_session_token (or environment variable: AWS_SESSION_TOKEN)\n",
     "- aws_profile_name\n",
+    "- aws_read_timeout (int or float having default value of 60 seconds)\n",
+    "- aws_connect_timeout (int or float having default value of 60 seconds)\n",
     "\n",
     "Beyond the authentication credentials, the only mandatory parameters are `api_type` and `model`.\n",
     "\n",
@@ -99,6 +101,8 @@
     "        \"aws_secret_key\": \"\",\n",
     "        \"aws_session_token\": \"\",\n",
     "        \"aws_profile_name\": \"\",\n",
+    "        \"aws_read_timeout\": 120,\n",
+    "        \"aws_connect_timeout\": 30,\n",
     "    },\n",
     "    {\n",
     "        \"api_type\": \"bedrock\",\n",