add volcengine tos storage (#8164)

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,*

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

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,
+    )

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)
 

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)

api/pyproject.toml

@@ -192,6 +192,7 @@ azure-ai-ml = "^1.19.0"
 azure-ai-inference = "^1.0.0b3"
 volcengine-python-sdk = {extras = ["ark"], version = "^1.0.98"}
 oci = "^2.133.0"
+tos = "^2.7.1"
 [tool.poetry.group.indriect.dependencies]
 kaleido = "0.2.1"
 rank-bm25 = "~0.2.2"

docker/.env.example

@@ -329,6 +329,18 @@ HUAWEI_OBS_ACCESS_KEY=your-access-key
 # The server url of the HUAWEI OBS service.
 HUAWEI_OBS_SERVER=your-server-url
 
+# Volcengine TOS Configuration
+# The name of the Volcengine TOS bucket to use for storing files.
+VOLCENGINE_TOS_BUCKET_NAME=your-bucket-name
+# The secret key to use for authenticating with the Volcengine TOS service.
+VOLCENGINE_TOS_SECRET_KEY=your-secret-key
+# The access key to use for authenticating with the Volcengine TOS service.
+VOLCENGINE_TOS_ACCESS_KEY=your-access-key
+# The endpoint of the Volcengine TOS service.
+VOLCENGINE_TOS_ENDPOINT=your-server-url
+# The region of the Volcengine TOS service.
+VOLCENGINE_TOS_REGION=your-region
+
 # ------------------------------
 # Vector Database Configuration
 # ------------------------------

docker/docker-compose.yaml

@@ -90,6 +90,11 @@ x-shared-env: &shared-api-worker-env
   OCI_ACCESS_KEY: ${OCI_ACCESS_KEY:-}
   OCI_SECRET_KEY: ${OCI_SECRET_KEY:-}
   OCI_REGION: ${OCI_REGION:-}
+  VOLCENGINE_TOS_BUCKET_NAME: ${VOLCENGINE_TOS_BUCKET_NAME:-}
+  VOLCENGINE_TOS_SECRET_KEY: ${VOLCENGINE_TOS_SECRET_KEY:-}
+  VOLCENGINE_TOS_ACCESS_KEY: ${VOLCENGINE_TOS_ACCESS_KEY:-}
+  VOLCENGINE_TOS_ENDPOINT: ${VOLCENGINE_TOS_ENDPOINT:-}
+  VOLCENGINE_TOS_REGION: ${VOLCENGINE_TOS_REGION:-}
   VECTOR_STORE: ${VECTOR_STORE:-weaviate}
   WEAVIATE_ENDPOINT: ${WEAVIATE_ENDPOINT:-http://weaviate:8080}
   WEAVIATE_API_KEY: ${WEAVIATE_API_KEY:-WVF5YThaHlkYwhGUSmCRgsX3tD5ngdN8pkih}