File size: 7,418 Bytes
fcaa164 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 |
# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
import os
from pathlib import Path, PurePath
from typing import Optional, Tuple
from warnings import warn
from camel.loaders import File, create_file_from_raw_bytes
from camel.storages.object_storages.base import BaseObjectStorage
class AmazonS3Storage(BaseObjectStorage):
r"""A class to connect with AWS S3 object storage to put and get objects
from one S3 bucket. The class will first try to use the credentials passed
as arguments, if not provided, it will look for the environment variables
`AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`. If none of these are
provided, it will try to use the local credentials (will be created if
logged in with AWS CLI).
Args:
bucket_name (str): The name of the S3 bucket.
create_if_not_exists (bool, optional): Whether to create the bucket if
it does not exist. Defaults to True.
access_key_id (Optional[str], optional): The AWS access key ID.
Defaults to None.
secret_access_key (Optional[str], optional): The AWS secret access key.
Defaults to None.
anonymous (bool, optional): Whether to use anonymous access. Defaults
to False.
References:
https://aws.amazon.com/pm/serv-s3/
https://aws.amazon.com/cli/
"""
def __init__(
self,
bucket_name: str,
create_if_not_exists: bool = True,
access_key_id: Optional[str] = None,
secret_access_key: Optional[str] = None,
anonymous: bool = False,
) -> None:
self._bucket_name = bucket_name
self._create_if_not_exists = create_if_not_exists
aws_key_id = access_key_id or os.getenv("AWS_ACCESS_KEY_ID")
aws_secret_key = secret_access_key or os.getenv(
"AWS_SECRET_ACCESS_KEY"
)
if not all([aws_key_id, aws_secret_key]) and not anonymous:
warn(
"AWS access key not configured. Local credentials will be "
"used."
)
# Make all the empty values None
aws_key_id = None
aws_secret_key = None
import botocore.session
from botocore import UNSIGNED
from botocore.config import Config
session = botocore.session.get_session()
if not anonymous:
self._client = session.create_client(
"s3",
aws_access_key_id=aws_key_id,
aws_secret_access_key=aws_secret_key,
)
else:
self._client = session.create_client(
"s3", config=Config(signature_version=UNSIGNED)
)
self._prepare_and_check()
def _prepare_and_check(self) -> None:
r"""Check privileges and existence of the bucket."""
from botocore.exceptions import ClientError, NoCredentialsError
try:
self._client.head_bucket(Bucket=self._bucket_name)
except ClientError as e:
error_code = e.response['Error']['Code']
if error_code == '403':
raise PermissionError(
f"Failed to access bucket {self._bucket_name}: "
f"No permission."
)
elif error_code == '404':
if self._create_if_not_exists:
self._client.create_bucket(Bucket=self._bucket_name)
warn(
f"Bucket {self._bucket_name} not found. Automatically "
f"created."
)
else:
raise FileNotFoundError(
f"Failed to access bucket {self._bucket_name}: Not "
f"found."
)
else:
raise e
except NoCredentialsError as e:
raise PermissionError("No AWS credentials found.") from e
@staticmethod
def canonicalize_path(file_path: PurePath) -> Tuple[str, str]:
r"""Canonicalize file path for Amazon S3.
Args:
file_path (PurePath): The path to be canonicalized.
Returns:
Tuple[str, str]: The canonicalized file key and file name.
"""
return file_path.as_posix(), file_path.name
def _put_file(self, file_key: str, file: File) -> None:
r"""Put a file to the Amazon S3 bucket.
Args:
file_key (str): The path to the object in the bucket.
file (File): The file to be uploaded.
"""
self._client.put_object(
Bucket=self._bucket_name, Key=file_key, Body=file.raw_bytes
)
def _get_file(self, file_key: str, filename: str) -> File:
r"""Get a file from the Amazon S3 bucket.
Args:
file_key (str): The path to the object in the bucket.
filename (str): The name of the file.
Returns:
File: The object from the S3 bucket.
"""
response = self._client.get_object(
Bucket=self._bucket_name, Key=file_key
)
raw_bytes = response["Body"].read()
return create_file_from_raw_bytes(raw_bytes, filename)
def _upload_file(
self, local_file_path: Path, remote_file_key: str
) -> None:
r"""Upload a local file to the Amazon S3 bucket.
Args:
local_file_path (Path): The path to the local file to be uploaded.
remote_file_key (str): The path to the object in the bucket.
"""
with open(local_file_path, "rb") as f:
self._client.put_object(
Bucket=self._bucket_name, Key=remote_file_key, Body=f
)
def _download_file(
self,
local_file_path: Path,
remote_file_key: str,
) -> None:
r"""Download a file from the Amazon S3 bucket to the local system.
Args:
local_file_path (Path): The path to the local file to be saved.
remote_file_key (str): The key of the object in the bucket.
"""
file = self._client.get_object(
Bucket=self._bucket_name,
Key=remote_file_key,
)
with open(local_file_path, "wb") as f:
f.write(file["Body"].read())
def _object_exists(self, file_key: str) -> bool:
r"""
Check if the object exists in the Amazon S3 bucket.
Args:
file_key: The key of the object in the bucket.
Returns:
bool: Whether the object exists in the bucket.
"""
try:
self._client.head_object(Bucket=self._bucket_name, Key=file_key)
return True
except self._client.exceptions.ClientError:
return False
|