| 18 | log = logging.getLogger(__name__) |
| 19 | |
| 20 | class BedrockClient(ModelClient): |
| 21 | __doc__ = r"""A component wrapper for the AWS Bedrock API client. |
| 22 | |
| 23 | AWS Bedrock provides a unified API that gives access to various foundation models |
| 24 | including Amazon's own models and third-party models like Anthropic Claude. |
| 25 | |
| 26 | Example: |
| 27 | ```python |
| 28 | from api.bedrock_client import BedrockClient |
| 29 | |
| 30 | client = BedrockClient() |
| 31 | generator = adal.Generator( |
| 32 | model_client=client, |
| 33 | model_kwargs={"model": "anthropic.claude-3-sonnet-20240229-v1:0"} |
| 34 | ) |
| 35 | ``` |
| 36 | """ |
| 37 | |
| 38 | def __init__( |
| 39 | self, |
| 40 | aws_access_key_id: Optional[str] = None, |
| 41 | aws_secret_access_key: Optional[str] = None, |
| 42 | aws_session_token: Optional[str] = None, |
| 43 | aws_region: Optional[str] = None, |
| 44 | aws_role_arn: Optional[str] = None, |
| 45 | *args, |
| 46 | **kwargs |
| 47 | ) -> None: |
| 48 | """Initialize the AWS Bedrock client. |
| 49 | |
| 50 | Args: |
| 51 | aws_access_key_id: AWS access key ID. If not provided, will use environment variable AWS_ACCESS_KEY_ID. |
| 52 | aws_secret_access_key: AWS secret access key. If not provided, will use environment variable AWS_SECRET_ACCESS_KEY. |
| 53 | aws_session_token: AWS session token. If not provided, will use environment variable AWS_SESSION_TOKEN. |
| 54 | aws_region: AWS region. If not provided, will use environment variable AWS_REGION. |
| 55 | aws_role_arn: AWS IAM role ARN for role-based authentication. If not provided, will use environment variable AWS_ROLE_ARN. |
| 56 | """ |
| 57 | super().__init__(*args, **kwargs) |
| 58 | from api.config import ( |
| 59 | AWS_ACCESS_KEY_ID, |
| 60 | AWS_SECRET_ACCESS_KEY, |
| 61 | AWS_SESSION_TOKEN, |
| 62 | AWS_REGION, |
| 63 | AWS_ROLE_ARN, |
| 64 | ) |
| 65 | |
| 66 | self.aws_access_key_id = aws_access_key_id or AWS_ACCESS_KEY_ID |
| 67 | self.aws_secret_access_key = aws_secret_access_key or AWS_SECRET_ACCESS_KEY |
| 68 | self.aws_session_token = aws_session_token or AWS_SESSION_TOKEN |
| 69 | self.aws_region = aws_region or AWS_REGION or "us-east-1" |
| 70 | self.aws_role_arn = aws_role_arn or AWS_ROLE_ARN |
| 71 | |
| 72 | self.sync_client = self.init_sync_client() |
| 73 | self.async_client = None # Initialize async client only when needed |
| 74 | |
| 75 | @classmethod |
| 76 | def from_dict(cls, data: Dict[str, Any]): |
| 77 | """Create an instance from a dictionary.""" |
no outgoing calls
no test coverage detected