From 518777fde3f0651188c365e644d013ce6e2d2010 Mon Sep 17 00:00:00 2001 From: iceyao Date: Mon, 9 Sep 2024 16:24:24 +0800 Subject: [PATCH] add volcengine tos storage --- api/.env.example | 9 +++- api/configs/middleware/__init__.py | 2 + .../storage/volcengine_tos_storage_config.py | 34 +++++++++++++ api/extensions/ext_storage.py | 3 ++ api/extensions/storage/volcengine_storage.py | 48 +++++++++++++++++++ 5 files changed, 95 insertions(+), 1 deletion(-) create mode 100644 api/configs/middleware/storage/volcengine_tos_storage_config.py create mode 100644 api/extensions/storage/volcengine_storage.py diff --git a/api/.env.example b/api/.env.example index 2320a4042d7eaa..a3db406aeabdce 100644 --- a/api/.env.example +++ b/api/.env.example @@ -39,7 +39,7 @@ DB_DATABASE=dify # Storage configuration # use for store upload files, private keys... -# storage type: local, s3, azure-blob, google-storage, tencent-cos, huawei-obs +# storage type: local, s3, azure-blob, google-storage, tencent-cos, huawei-obs, volcengine-tos STORAGE_TYPE=local STORAGE_LOCAL_PATH=storage S3_USE_AWS_MANAGED_IAM=false @@ -86,6 +86,13 @@ OCI_ACCESS_KEY=your-access-key OCI_SECRET_KEY=your-secret-key OCI_REGION=your-region +# Volcengine tos Storage configuration +VOLCENGINE_TOS_ENDPOINT=your-endpoint +VOLCENGINE_TOS_BUCKET_NAME=your-bucket-name +VOLCENGINE_TOS_ACCESS_KEY=your-access-key +VOLCENGINE_TOS_SECRET_KEY=your-secret-key +VOLCENGINE_TOS_REGION=your-region + # CORS configuration WEB_API_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,* CONSOLE_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,* diff --git a/api/configs/middleware/__init__.py b/api/configs/middleware/__init__.py index 18c5ba8df36ea9..e017c2c5b828ef 100644 --- a/api/configs/middleware/__init__.py +++ b/api/configs/middleware/__init__.py @@ -12,6 +12,7 @@ from configs.middleware.storage.huawei_obs_storage_config import HuaweiCloudOBSStorageConfig from configs.middleware.storage.oci_storage_config import OCIStorageConfig from configs.middleware.storage.tencent_cos_storage_config import TencentCloudCOSStorageConfig +from configs.middleware.storage.volcengine_tos_storage_config import VolcengineTOSStorageConfig from configs.middleware.vdb.analyticdb_config import AnalyticdbConfig from configs.middleware.vdb.chroma_config import ChromaConfig from configs.middleware.vdb.elasticsearch_config import ElasticsearchConfig @@ -201,6 +202,7 @@ class MiddlewareConfig( GoogleCloudStorageConfig, TencentCloudCOSStorageConfig, HuaweiCloudOBSStorageConfig, + VolcengineTOSStorageConfig, S3StorageConfig, OCIStorageConfig, # configs of vdb and vdb providers diff --git a/api/configs/middleware/storage/volcengine_tos_storage_config.py b/api/configs/middleware/storage/volcengine_tos_storage_config.py new file mode 100644 index 00000000000000..a0e09a3cc76711 --- /dev/null +++ b/api/configs/middleware/storage/volcengine_tos_storage_config.py @@ -0,0 +1,34 @@ +from typing import Optional + +from pydantic import BaseModel, Field + + +class VolcengineTOSStorageConfig(BaseModel): + """ + Volcengine tos storage configs + """ + + VOLCENGINE_TOS_BUCKET_NAME: Optional[str] = Field( + description="Volcengine TOS Bucket Name", + default=None, + ) + + VOLCENGINE_TOS_ACCESS_KEY: Optional[str] = Field( + description="Volcengine TOS Access Key", + default=None, + ) + + VOLCENGINE_TOS_SECRET_KEY: Optional[str] = Field( + description="Volcengine TOS Secret Key", + default=None, + ) + + VOLCENGINE_TOS_ENDPOINT: Optional[str] = Field( + description="Volcengine TOS Endpoint URL", + default=None, + ) + + VOLCENGINE_TOS_REGION: Optional[str] = Field( + description="Volcengine TOS Region", + default=None, + ) diff --git a/api/extensions/ext_storage.py b/api/extensions/ext_storage.py index 6f572e3e81edb1..5ce18b72925774 100644 --- a/api/extensions/ext_storage.py +++ b/api/extensions/ext_storage.py @@ -11,6 +11,7 @@ from extensions.storage.oci_storage import OCIStorage from extensions.storage.s3_storage import S3Storage from extensions.storage.tencent_storage import TencentStorage +from extensions.storage.volcengine_storage import VolcengineStorage class Storage: @@ -33,6 +34,8 @@ def init_app(self, app: Flask): self.storage_runner = OCIStorage(app=app) elif storage_type == "huawei-obs": self.storage_runner = HuaweiStorage(app=app) + elif storage_type == "volcengine-tos": + self.storage_runner = VolcengineStorage(app=app) else: self.storage_runner = LocalStorage(app=app) diff --git a/api/extensions/storage/volcengine_storage.py b/api/extensions/storage/volcengine_storage.py new file mode 100644 index 00000000000000..f74ad2ee6dc036 --- /dev/null +++ b/api/extensions/storage/volcengine_storage.py @@ -0,0 +1,48 @@ +from collections.abc import Generator + +import tos +from flask import Flask + +from extensions.storage.base_storage import BaseStorage + + +class VolcengineStorage(BaseStorage): + """Implementation for Volcengine TOS storage.""" + + def __init__(self, app: Flask): + super().__init__(app) + app_config = self.app.config + self.bucket_name = app_config.get("VOLCENGINE_TOS_BUCKET_NAME") + self.client = tos.TosClientV2( + ak=app_config.get("VOLCENGINE_TOS_ACCESS_KEY"), + sk=app_config.get("VOLCENGINE_TOS_SECRET_KEY"), + endpoint=app_config.get("VOLCENGINE_TOS_ENDPOINT"), + region=app_config.get("VOLCENGINE_TOS_REGION"), + ) + + def save(self, filename, data): + self.client.put_object(bucket=self.bucket_name, key=filename, content=data) + + def load_once(self, filename: str) -> bytes: + data = self.client.get_object(bucket=self.bucket_name, key=filename).read() + return data + + def load_stream(self, filename: str) -> Generator: + def generate(filename: str = filename) -> Generator: + response = self.client.get_object(bucket=self.bucket_name, key=filename) + while chunk := response.read(4096): + yield chunk + + return generate() + + def download(self, filename, target_filepath): + self.client.get_object_to_file(bucket=self.bucket_name, key=filename, file_path=target_filepath) + + def exists(self, filename): + res = self.client.head_object(bucket=self.bucket_name, key=filename) + if res.status_code != 200: + return False + return True + + def delete(self, filename): + self.client.delete_object(bucket=self.bucket_name, key=filename)