From 91aebf36dcc6a5c693ebec7099ac45fc49fac0ca Mon Sep 17 00:00:00 2001 From: Wh1isper <9573586@qq.com> Date: Thu, 27 Mar 2025 21:17:57 +0800 Subject: [PATCH 1/5] feat: Add read timeout and connect timeout for bedrock provider --- pydantic_ai_slim/pydantic_ai/providers/bedrock.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py index 79c4e679ba..2cd53e97d4 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py @@ -1,5 +1,6 @@ from __future__ import annotations as _annotations +import os from typing import overload from pydantic_ai.exceptions import UserError @@ -8,6 +9,7 @@ try: import boto3 from botocore.client import BaseClient + from botocore.config import Config from botocore.exceptions import NoRegionError except ImportError as _import_error: raise ImportError( @@ -42,6 +44,8 @@ def __init__( aws_access_key_id: str | None = None, aws_secret_access_key: str | None = None, aws_session_token: str | None = None, + aws_read_timeout: float | None = None, + aws_connect_timeout: float | None = None, ) -> None: ... def __init__( @@ -52,6 +56,8 @@ def __init__( aws_access_key_id: str | None = None, aws_secret_access_key: str | None = None, aws_session_token: str | None = None, + aws_read_timeout: float | None = None, + aws_connect_timeout: float | None = None, ) -> None: """Initialize the Bedrock provider. @@ -61,6 +67,8 @@ def __init__( aws_access_key_id: The AWS access key ID. aws_secret_access_key: The AWS secret access key. aws_session_token: The AWS session token. + aws_read_timeout: The read timeout for bedrock client. + aws_connect_timeout: The connect timeout for bedrock client. """ if bedrock_client is not None: self._client = bedrock_client @@ -72,6 +80,12 @@ def __init__( aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token, region_name=region_name, + config=Config( + read_timeout=aws_read_timeout + or float(os.getenv('AWS_READ_TIMEOUT', 300)), # Need more time for big tool call + connect_timeout=aws_connect_timeout + or float(os.getenv('AWS_CONNECT_TIMEOUT', 60)), # Boto3 use 60 as default + ), ) except NoRegionError as exc: # pragma: no cover raise UserError('You must provide a `region_name` or a boto3 client for Bedrock Runtime.') from exc From fc65e8b3763e5fd34e629c90e3fcf4615d587fa4 Mon Sep 17 00:00:00 2001 From: Wh1isper <9573586@qq.com> Date: Thu, 27 Mar 2025 21:21:35 +0800 Subject: [PATCH 2/5] change comment --- pydantic_ai_slim/pydantic_ai/providers/bedrock.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py index 2cd53e97d4..adf2e30d72 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py @@ -82,7 +82,7 @@ def __init__( region_name=region_name, config=Config( read_timeout=aws_read_timeout - or float(os.getenv('AWS_READ_TIMEOUT', 300)), # Need more time for big tool call + or float(os.getenv('AWS_READ_TIMEOUT', 300)), # Need more time for long tool call connect_timeout=aws_connect_timeout or float(os.getenv('AWS_CONNECT_TIMEOUT', 60)), # Boto3 use 60 as default ), From 269cdb72f94328cbc6279c78337cc9219829daca Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Mon, 31 Mar 2025 11:27:59 +0200 Subject: [PATCH 3/5] Add test --- pydantic_ai_slim/pydantic_ai/providers/bedrock.py | 13 +++++-------- tests/providers/test_bedrock.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py index adf2e30d72..1a8980e057 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py @@ -67,25 +67,22 @@ def __init__( aws_access_key_id: The AWS access key ID. aws_secret_access_key: The AWS secret access key. aws_session_token: The AWS session token. - aws_read_timeout: The read timeout for bedrock client. - aws_connect_timeout: The connect timeout for bedrock client. + aws_read_timeout: The read timeout for Bedrock client. + aws_connect_timeout: The connect timeout for Bedrock client. """ if bedrock_client is not None: self._client = bedrock_client else: try: + read_timeout = aws_read_timeout or float(os.getenv('AWS_READ_TIMEOUT', 300)) + connect_timeout = aws_connect_timeout or float(os.getenv('AWS_CONNECT_TIMEOUT', 60)) self._client = boto3.client( # type: ignore[reportUnknownMemberType] 'bedrock-runtime', aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token, region_name=region_name, - config=Config( - read_timeout=aws_read_timeout - or float(os.getenv('AWS_READ_TIMEOUT', 300)), # Need more time for long tool call - connect_timeout=aws_connect_timeout - or float(os.getenv('AWS_CONNECT_TIMEOUT', 60)), # Boto3 use 60 as default - ), + config=Config(read_timeout=read_timeout, connect_timeout=connect_timeout), ) except NoRegionError as exc: # pragma: no cover raise UserError('You must provide a `region_name` or a boto3 client for Bedrock Runtime.') from exc diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index 5fce0072c5..6f07c6c964 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -1,4 +1,7 @@ +from typing import cast + import pytest +from mypy_boto3_bedrock_runtime import BedrockRuntimeClient from ..conftest import TestEnv, try_import @@ -15,3 +18,15 @@ def test_bedrock_provider(env: TestEnv): assert isinstance(provider, BedrockProvider) assert provider.name == 'bedrock' assert provider.base_url == 'https://bedrock-runtime.us-east-1.amazonaws.com' + + +def test_bedrock_provider_timeout(env: TestEnv): + env.set('AWS_READ_TIMEOUT', '1') + env.set('AWS_CONNECT_TIMEOUT', '1') + provider = BedrockProvider() + assert isinstance(provider, BedrockProvider) + assert provider.name == 'bedrock' + + config = cast(BedrockRuntimeClient, provider.client).meta.config + assert config.read_timeout == 1 # type: ignore + assert config.connect_timeout == 1 # type: ignore From 4d83527cb0e9eaadfc1e72c3aa95bd798c33ddc8 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Mon, 31 Mar 2025 11:29:27 +0200 Subject: [PATCH 4/5] Move import inside --- tests/providers/test_bedrock.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index 6f07c6c964..f8fd43e91d 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -1,11 +1,12 @@ from typing import cast import pytest -from mypy_boto3_bedrock_runtime import BedrockRuntimeClient from ..conftest import TestEnv, try_import with try_import() as imports_successful: + from mypy_boto3_bedrock_runtime import BedrockRuntimeClient + from pydantic_ai.providers.bedrock import BedrockProvider From fd071d564ecc05554fbff9a6d8dcd519f77c03f4 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Mon, 31 Mar 2025 11:32:10 +0200 Subject: [PATCH 5/5] Add default region --- tests/providers/test_bedrock.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index f8fd43e91d..9f62fe6879 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -22,6 +22,7 @@ def test_bedrock_provider(env: TestEnv): def test_bedrock_provider_timeout(env: TestEnv): + env.set('AWS_DEFAULT_REGION', 'us-east-1') env.set('AWS_READ_TIMEOUT', '1') env.set('AWS_CONNECT_TIMEOUT', '1') provider = BedrockProvider()