diff --git a/autogen/oai/bedrock.py b/autogen/oai/bedrock.py index a07073eef299..2de3fecd8b1a 100644 --- a/autogen/oai/bedrock.py +++ b/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( diff --git a/notebook/agentchat_nested_chats_chess_altmodels.ipynb b/notebook/agentchat_nested_chats_chess_altmodels.ipynb index 8980a87e8818..a50849cffc24 100644 --- a/notebook/agentchat_nested_chats_chess_altmodels.ipynb +++ b/notebook/agentchat_nested_chats_chess_altmodels.ipynb @@ -130,7 +130,7 @@ " move: Annotated[\n", " str,\n", " \"Call this tool to make a move after you have the list of legal moves and want to make a move. Takes UCI format, e.g. e2e4 or e7e5 or e7e8q.\",\n", - " ]\n", + " ],\n", ") -> Annotated[str, \"Result of the move.\"]:\n", " move = chess.Move.from_uci(move)\n", " board.push_uci(str(move))\n", diff --git a/test/oai/test_bedrock.py b/test/oai/test_bedrock.py index 42502acf691c..5fb78ca68285 100644 --- a/test/oai/test_bedrock.py +++ b/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, diff --git a/website/docs/topics/non-openai-models/cloud-bedrock.ipynb b/website/docs/topics/non-openai-models/cloud-bedrock.ipynb index 422598dd2fe1..3246dccbbfc0 100644 --- a/website/docs/topics/non-openai-models/cloud-bedrock.ipynb +++ b/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",