Browse Source

feat: support tencent cos storage (#5297)

quicksand 10 months ago
parent
commit
147a39b984

+ 7 - 0
api/.env.example

@@ -65,6 +65,13 @@ ALIYUN_OSS_REGION=your-region
 GOOGLE_STORAGE_BUCKET_NAME=yout-bucket-name
 GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON=your-google-service-account-json-base64-string
 
+# Tencent COS Storage configuration
+TENCENT_COS_BUCKET_NAME=your-bucket-name
+TENCENT_COS_SECRET_KEY=your-secret-key
+TENCENT_COS_SECRET_ID=your-secret-id
+TENCENT_COS_REGION=your-region
+TENCENT_COS_SCHEME=your-scheme
+
 # CORS configuration
 WEB_API_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,*
 CONSOLE_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,*

+ 7 - 0
api/config.py

@@ -253,6 +253,13 @@ class Config:
         self.GOOGLE_STORAGE_BUCKET_NAME = get_env('GOOGLE_STORAGE_BUCKET_NAME')
         self.GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64 = get_env('GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64')
 
+        # Tencent Cos Storage settings
+        self.TENCENT_COS_BUCKET_NAME = get_env('TENCENT_COS_BUCKET_NAME')
+        self.TENCENT_COS_REGION = get_env('TENCENT_COS_REGION')
+        self.TENCENT_COS_SECRET_ID = get_env('TENCENT_COS_SECRET_ID')
+        self.TENCENT_COS_SECRET_KEY = get_env('TENCENT_COS_SECRET_KEY')
+        self.TENCENT_COS_SCHEME = get_env('TENCENT_COS_SCHEME')
+
         # ------------------------
         # Vector Store Configurations.
         # Currently, only support: qdrant, milvus, zilliz, weaviate, relyt, pgvector

+ 5 - 0
api/extensions/ext_storage.py

@@ -8,6 +8,7 @@ from extensions.storage.azure_storage import AzureStorage
 from extensions.storage.google_storage import GoogleStorage
 from extensions.storage.local_storage import LocalStorage
 from extensions.storage.s3_storage import S3Storage
+from extensions.storage.tencent_storage import TencentStorage
 
 
 class Storage:
@@ -32,6 +33,10 @@ class Storage:
             self.storage_runner = GoogleStorage(
                 app=app
             )
+        elif storage_type == 'tencent-cos':
+            self.storage_runner = TencentStorage(
+                app=app
+            )
         else:
             self.storage_runner = LocalStorage(app=app)
 

+ 48 - 0
api/extensions/storage/tencent_storage.py

@@ -0,0 +1,48 @@
+from collections.abc import Generator
+
+from flask import Flask
+from qcloud_cos import CosConfig, CosS3Client
+
+from extensions.storage.base_storage import BaseStorage
+
+
+class TencentStorage(BaseStorage):
+    """Implementation for tencent cos storage.
+    """
+
+    def __init__(self, app: Flask):
+        super().__init__(app)
+        app_config = self.app.config
+        self.bucket_name = app_config.get('TENCENT_COS_BUCKET_NAME')
+        config = CosConfig(
+            Region=app_config.get('TENCENT_COS_REGION'),
+            SecretId=app_config.get('TENCENT_COS_SECRET_ID'),
+            SecretKey=app_config.get('TENCENT_COS_SECRET_KEY'),
+            Scheme=app_config.get('TENCENT_COS_SCHEME'),
+        )
+        self.client = CosS3Client(config)
+
+    def save(self, filename, data):
+        self.client.put_object(Bucket=self.bucket_name, Body=data, Key=filename)
+
+    def load_once(self, filename: str) -> bytes:
+        data = self.client.get_object(Bucket=self.bucket_name, Key=filename)['Body'].get_raw_stream().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['Body'].get_stream(chunk_size=4096):
+                yield chunk
+
+        return generate()
+
+    def download(self, filename, target_filepath):
+        response = self.client.get_object(Bucket=self.bucket_name, Key=filename)
+        response['Body'].get_stream_to_file(target_filepath)
+
+    def exists(self, filename):
+        return self.client.object_exists(Bucket=self.bucket_name, Key=filename)
+
+    def delete(self, filename):
+        self.client.delete_object(Bucket=self.bucket_name, Key=filename)

+ 3 - 3
api/poetry.lock

@@ -1441,12 +1441,12 @@ test-no-images = ["pytest", "pytest-cov", "pytest-xdist", "wurlitzer"]
 
 [[package]]
 name = "cos-python-sdk-v5"
-version = "1.9.29"
+version = "1.9.30"
 description = "cos-python-sdk-v5"
 optional = false
 python-versions = "*"
 files = [
-    {file = "cos-python-sdk-v5-1.9.29.tar.gz", hash = "sha256:1bb07022368d178e7a50a3cc42e0d6cbf4b0bef2af12a3bb8436904339cdec8e"},
+    {file = "cos-python-sdk-v5-1.9.30.tar.gz", hash = "sha256:a23fd090211bf90883066d90cd74317860aa67c6d3aa80fe5e44b18c7e9b2a81"},
 ]
 
 [package.dependencies]
@@ -8921,4 +8921,4 @@ testing = ["coverage (>=5.0.3)", "zope.event", "zope.testing"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "05dd46b65d1cfeb9d295934777d70de8dbbbd20b62e6d3c976550f3134ff7a1e"
+content-hash = "7f9fada276f5050d1418e973f830c690807fc7f37e05c5f2e31319522daf0323"

+ 1 - 1
api/pyproject.toml

@@ -186,7 +186,7 @@ tencentcloud-sdk-python-hunyuan = "~3.0.1158"
 tcvectordb = "1.3.2"
 chromadb = "~0.5.0"
 tenacity = "~8.3.0"
-
+cos-python-sdk-v5 = "1.9.30"
 
 [tool.poetry.group.dev]
 optional = true

+ 1 - 0
api/requirements.txt

@@ -89,3 +89,4 @@ vanna[postgres,mysql,clickhouse,duckdb]==0.5.5
 tencentcloud-sdk-python-hunyuan~=3.0.1158
 chromadb~=0.5.0
 tenacity~=8.3.0
+cos-python-sdk-v5==1.9.30

+ 12 - 0
docker/docker-compose.yaml

@@ -96,6 +96,12 @@ services:
       GOOGLE_STORAGE_BUCKET_NAME: 'yout-bucket-name'
       # if you want to use Application Default Credentials, you can leave GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64 empty.
       GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64: 'your-google-service-account-json-base64-string'
+      # The Tencent COS storage configurations, only available when STORAGE_TYPE is `tencent-cos`.
+      TENCENT_COS_BUCKET_NAME: 'your-bucket-name'
+      TENCENT_COS_SECRET_KEY: 'your-secret-key'
+      TENCENT_COS_SECRET_ID: 'your-secret-id'
+      TENCENT_COS_REGION: 'your-region'
+      TENCENT_COS_SCHEME: 'your-scheme'
       # The type of vector store to use. Supported values are `weaviate`, `qdrant`, `milvus`, `relyt`.
       VECTOR_STORE: weaviate
       # The Weaviate endpoint URL. Only available when VECTOR_STORE is `weaviate`.
@@ -252,6 +258,12 @@ services:
       GOOGLE_STORAGE_BUCKET_NAME: 'yout-bucket-name'
       # if you want to use Application Default Credentials, you can leave GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64 empty.
       GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64: 'your-google-service-account-json-base64-string'
+      # The Tencent COS storage configurations, only available when STORAGE_TYPE is `tencent-cos`.
+      TENCENT_COS_BUCKET_NAME: 'your-bucket-name'
+      TENCENT_COS_SECRET_KEY: 'your-secret-key'
+      TENCENT_COS_SECRET_ID: 'your-secret-id'
+      TENCENT_COS_REGION: 'your-region'
+      TENCENT_COS_SCHEME: 'your-scheme'
       # The type of vector store to use. Supported values are `weaviate`, `qdrant`, `milvus`, `relyt`, `pgvector`.
       VECTOR_STORE: weaviate
       # The Weaviate endpoint URL. Only available when VECTOR_STORE is `weaviate`.