Merge pull request #39 from choinek/main

pkarw · web-flow · commit 4929da88e8f7 · 2024-11-18T12:05:48.000+01:00
[feat] #15 Add S3 storage strategy
diff --git a/.env.localhost.example b/.env.localhost.example
@@ -15,4 +15,10 @@ CELERY_RESULT_BACKEND=redis://localhost:6379/0
 OLLAMA_HOST=http://localhost:11434
 APP_ENV=development  # Default to development mode
 
-STORAGE_PROFILE_PATH=../storage_profiles
+STORAGE_PROFILE_PATH=../storage_profiles
+
+# AWS S3
+#AWS_ACCESS_KEY_ID=your-access-key-id
+#AWS_SECRET_ACCESS_KEY=your-secret-access-key
+#AWS_REGION=your-region
+#AWS_S3_BUCKET_NAME=your-bucket-name
diff --git a/README.md b/README.md
@@ -464,7 +464,7 @@ curl -X POST "http://localhost:8000/llm/generate" -H "Content-Type: application/
 
 The tool can automatically save the results using different storage strategies and storage profiles. Storage profiles are set in the `/storage_profiles` by a yaml configuration files.
 
-Example:
+### Local File System
 
 ```yaml
 strategy: local_filesystem
@@ -474,7 +474,7 @@ settings:
   create_subfolders: true
 ```
 
-for Google drive:
+### Google Drive
 
 ```yaml
 strategy: google_drive
@@ -489,8 +489,56 @@ Where the `service_account_file` is a `json` file with authorization credentials
 
 Note: Service Account is different account that the one you're using for Google workspace (files will not be visible in the UI)
 
+### Amazon S3 - Cloud Object Storage
+
+```yaml
+strategy: aws_s3
+settings:
+  bucket_name: ${AWS_S3_BUCKET_NAME}
+  region: ${AWS_REGION}
+  access_key: ${AWS_ACCESS_KEY_ID}
+  secret_access_key: ${AWS_SECRET_ACCESS_KEY}
+```
+
+#### Requirements for AWS S3 Access Key
+
+1. **Access Key Ownership**  
+   The access key must belong to an IAM user or role with permissions for S3 operations.
+
+2. **IAM Policy Example**  
+   The IAM policy attached to the user or role must allow the necessary actions. Below is an example of a policy granting access to an S3 bucket:
+   ```json
+   {
+       "Version": "2012-10-17",
+       "Statement": [
+           {
+               "Effect": "Allow",
+               "Action": [
+                   "s3:PutObject",
+                   "s3:GetObject",
+                   "s3:ListBucket",
+                   "s3:DeleteObject"
+               ],
+               "Resource": [
+                   "arn:aws:s3:::your-bucket-name",
+                   "arn:aws:s3:::your-bucket-name/*"
+               ]
+           }
+       ]
+   }
+   ```
+
+Next, populate the appropriate `.env` file (e.g., .env, .env.localhost) with the required AWS credentials:
+
+```bash
+AWS_ACCESS_KEY_ID=your-access-key-id
+AWS_SECRET_ACCESS_KEY=your-secret-access-key
+AWS_REGION=your-region
+AWS_S3_BUCKET_NAME=your-bucket-name
+```
+
 ## License
-This project is licensed under the GNU General Public License. See the [LICENSE](LICENSE.md) file for details.
+This project is licensed under the GNU General Public License. See the [LICENSE](LICENSE) file for details.
 
 **Important note on [marker](https://github.com/VikParuchuri/marker) license***:
 
diff --git a/app/requirements.txt b/app/requirements.txt
@@ -17,3 +17,4 @@ google-auth-oauthlib
 transformers==4.45.2
 surya-ocr==0.4.14
 marker-pdf==0.2.6
+boto3
diff --git a/app/storage_manager.py b/app/storage_manager.py
@@ -2,19 +2,22 @@
 import yaml
 from storage_strategies.local_filesystem import LocalFilesystemStorageStrategy
 from storage_strategies.google_drive import GoogleDriveStorageStrategy
+from storage_strategies.aws_s3 import AWSS3StorageStrategy
 from pathlib import Path
 
 class StorageManager:
     def __init__(self, profile_name):
         profile_path = os.path.join(os.getenv('STORAGE_PROFILE_PATH', '/storage_profiles'), f'{profile_name}.yaml')
         with open(profile_path, 'r') as file:
             self.profile = yaml.safe_load(file)
-        
+
         strategy = self.profile['strategy']
         if strategy == 'local_filesystem':
             self.strategy = LocalFilesystemStorageStrategy(self.profile)
         elif strategy == 'google_drive':
             self.strategy = GoogleDriveStorageStrategy(self.profile)
+        elif strategy == 'aws_s3':
+            self.strategy = AWSS3StorageStrategy(self.profile)
         else:
             raise ValueError(f"Unknown storage strategy '{strategy}'")
 
@@ -28,4 +31,4 @@ def list(self):
         return self.strategy.list()
 
     def delete(self, file_name):
-        self.strategy.delete(file_name)
+        self.strategy.delete(file_name)
diff --git a/app/storage_strategies/aws_s3.py b/app/storage_strategies/aws_s3.py
@@ -0,0 +1,81 @@
+import boto3
+from botocore.exceptions import EndpointConnectionError, ClientError
+from storage_strategies.storage_strategy import StorageStrategy
+
+class AWSS3StorageStrategy(StorageStrategy):
+    def __init__(self, context):
+        super().__init__(context)
+
+        self.bucket_name = self.resolve_placeholder(context['settings'].get('bucket_name'))
+        self.region = self.resolve_placeholder(context['settings'].get('region'))
+        self.access_key = self.resolve_placeholder(context['settings'].get('access_key'))
+        self.secret_access_key = self.resolve_placeholder(context['settings'].get('secret_access_key'))
+
+        try:
+            self.s3_client = boto3.client(
+                's3',
+                aws_access_key_id=self.access_key,
+                aws_secret_access_key=self.secret_access_key,
+                region_name=self.region
+            )
+            self.s3_client.head_bucket(Bucket=self.bucket_name)
+        except EndpointConnectionError as e:
+            raise RuntimeError(
+                f"{str(e)}\n"
+                "Check your AWS_REGION and AWS_S3_BUCKET_NAME environment variables."
+            ) from e
+        except ClientError as e:
+            error_code = e.response.get('Error', {}).get('Code', 'Unknown')
+            if error_code in ('400', '403'):
+                raise RuntimeError(
+                    f"{str(e)}\n"
+                    "Error: Please check your AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY."
+                ) from e
+            raise
+
+    def save(self, file_name, dest_file_name, content):
+        formatted_file_name = self.format_file_name(file_name, dest_file_name)
+
+        try:
+            self.s3_client.put_object(
+                Bucket=self.bucket_name,
+                Key=formatted_file_name,
+                Body=content.encode('utf-8')
+            )
+        except ClientError as e:
+            raise RuntimeError(
+                f"{str(e)}\n"
+                f"Error saving file '{file_name}' as '{formatted_file_name}' to bucket '{self.bucket_name}'."
+            ) from e
+
+    def load(self, file_name):
+        try:
+            response = self.s3_client.get_object(Bucket=self.bucket_name, Key=file_name)
+            return response['Body'].read().decode('utf-8')
+        except ClientError as e:
+            error_code = e.response['Error']['Code']
+            if error_code == 'NoSuchKey':
+                return None
+            raise RuntimeError(
+                f"{str(e)}\n"
+                f"Error loading file '{file_name}' from bucket '{self.bucket_name}'."
+            ) from e
+
+    def list(self):
+        try:
+            response = self.s3_client.list_objects_v2(Bucket=self.bucket_name)
+            return [item['Key'] for item in response.get('Contents', [])]
+        except ClientError as e:
+            raise RuntimeError(
+                f"{str(e)}\n"
+                f"Error listing objects in bucket '{self.bucket_name}'."
+            ) from e
+
+    def delete(self, file_name):
+        try:
+            self.s3_client.delete_object(Bucket=self.bucket_name, Key=file_name)
+        except ClientError as e:
+            raise RuntimeError(
+                f"{str(e)}\n"
+                f"Error deleting file '{file_name}' from bucket '{self.bucket_name}'."
+            ) from e
diff --git a/app/storage_strategies/storage_strategy.py b/app/storage_strategies/storage_strategy.py
@@ -2,6 +2,7 @@
 import yaml
 from datetime import datetime
 from pathlib import Path
+from string import Template
 
 class StorageStrategy:
     def __init__(self, context):
@@ -18,7 +19,7 @@ def list(self):
 
     def delete(self, file_name):
         raise NotImplementedError("Subclasses must implement this method")
-    
+
     def format_file_name(self, file_name, format_string):
         return format_string.format(file_fullname=file_name,  # file_name with path
                                      file_name=Path(file_name).stem,  # file_name without path
@@ -28,4 +29,15 @@ def format_file_name(self, file_name, format_string):
                                      dd=datetime.now().strftime('%d'),
                                      HH=datetime.now().strftime('%H'),
                                      MM=datetime.now().strftime('%M'),
-                                     SS=datetime.now().strftime('%S'))
+                                     SS=datetime.now().strftime('%S'))
+
+    def resolve_placeholder(self, value, default=None):
+        if not value:
+            return default
+        try:
+            return Template(value).substitute(os.environ)
+        except KeyError as e:
+            if default:
+                return default
+            else:
+                raise ValueError(f"Environment variable '{e.args[0]}' is missing, and no default value is provided.")
diff --git a/storage_profiles/s3.yaml b/storage_profiles/s3.yaml
@@ -0,0 +1,6 @@
+strategy: aws_s3
+settings:
+  bucket_name: ${AWS_S3_BUCKET_NAME}
+  region: ${AWS_REGION}
+  access_key: ${AWS_ACCESS_KEY_ID}
+  secret_access_key: ${AWS_SECRET_ACCESS_KEY}
