llama-index

1
from typing import TYPE_CHECKING, Optional
2

3
if TYPE_CHECKING:
4
    import botocore
5

6

7
def get_aws_service_client(
8
    service_name: Optional[str] = None,
9
    region_name: Optional[str] = None,
10
    aws_access_key_id: Optional[str] = None,
11
    aws_secret_access_key: Optional[str] = None,
12
    aws_session_token: Optional[str] = None,
13
    profile_name: Optional[str] = None,
14
    max_retries: Optional[int] = 3,
15
    timeout: Optional[float] = 60.0,
16
) -> "botocore.client.BaseClient":
17
    try:
18
        import boto3
19
        import botocore
20
    except ImportError:
21
        raise ImportError(
22
            "Please run `pip install boto3 botocore` to use AWS services."
23
        )
24

25
    config = botocore.config.Config(
26
        retries={"max_attempts": max_retries, "mode": "standard"},
27
        connect_timeout=timeout,
28
    )
29

30
    try:
31
        if not profile_name and aws_access_key_id:
32
            session = boto3.Session(
33
                aws_access_key_id=aws_access_key_id,
34
                aws_secret_access_key=aws_secret_access_key,
35
                aws_session_token=aws_session_token,
36
                region_name=region_name,
37
            )
38
            client = session.client(service_name, config=config)
39
        else:
40
            session = boto3.Session(profile_name=profile_name)
41
            if region_name:
42
                client = session.client(
43
                    service_name, region_name=region_name, config=config
44
                )
45
            else:
46
                client = session.client(service_name, config=config)
47
    except Exception as e:
48
        raise ValueError("Please verify the provided credentials.") from (e)
49

50
    return client
51