feat: support tencent cos storage (#5297)

This commit is contained in:
quicksand 2024-06-17 19:18:52 +08:00 committed by GitHub
parent 7a758a35fe
commit 147a39b984
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 84 additions and 4 deletions

View File

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

View File

@ -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

View File

@ -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)

View File

@ -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)

6
api/poetry.lock generated
View File

@ -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"

View File

@ -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

View File

@ -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

View File

@ -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`.