feat: Add google storage support (#3887)

Co-authored-by: miendinh <miendinh@users.noreply.github.com>
This commit is contained in:
miendinh 2024-04-27 17:26:52 +07:00 committed by GitHub
parent 661b30784e
commit 9ad489d133
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 43 additions and 1 deletions

View File

@ -54,6 +54,9 @@ ALIYUN_OSS_BUCKET_NAME=your-bucket-name
ALIYUN_OSS_ACCESS_KEY=your-access-key ALIYUN_OSS_ACCESS_KEY=your-access-key
ALIYUN_OSS_SECRET_KEY=your-secret-key ALIYUN_OSS_SECRET_KEY=your-secret-key
ALIYUN_OSS_ENDPOINT=your-endpoint ALIYUN_OSS_ENDPOINT=your-endpoint
# Google Storage configuration
GOOGLE_STORAGE_BUCKET_NAME=yout-bucket-name
GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON=your-google-service-account-json-base64-string
# CORS configuration # CORS configuration
WEB_API_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,* WEB_API_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,*

View File

@ -213,6 +213,8 @@ class Config:
self.ALIYUN_OSS_ACCESS_KEY=get_env('ALIYUN_OSS_ACCESS_KEY') self.ALIYUN_OSS_ACCESS_KEY=get_env('ALIYUN_OSS_ACCESS_KEY')
self.ALIYUN_OSS_SECRET_KEY=get_env('ALIYUN_OSS_SECRET_KEY') self.ALIYUN_OSS_SECRET_KEY=get_env('ALIYUN_OSS_SECRET_KEY')
self.ALIYUN_OSS_ENDPOINT=get_env('ALIYUN_OSS_ENDPOINT') self.ALIYUN_OSS_ENDPOINT=get_env('ALIYUN_OSS_ENDPOINT')
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')
# ------------------------ # ------------------------
# Vector Store Configurations. # Vector Store Configurations.

View File

@ -1,3 +1,4 @@
import base64
import os import os
import shutil import shutil
from collections.abc import Generator from collections.abc import Generator
@ -11,6 +12,7 @@ from azure.storage.blob import AccountSasPermissions, BlobServiceClient, Resourc
from botocore.client import Config from botocore.client import Config
from botocore.exceptions import ClientError from botocore.exceptions import ClientError
from flask import Flask from flask import Flask
from google.cloud import storage as GoogleStorage
class Storage: class Storage:
@ -51,6 +53,10 @@ class Storage:
self.bucket_name, self.bucket_name,
connect_timeout=30 connect_timeout=30
) )
elif self.storage_type == 'google-storage':
self.bucket_name = app.config.get('GOOGLE_STORAGE_BUCKET_NAME')
service_account_json = base64.b64decode(app.config.get('GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64')).decode('utf-8')
self.client = GoogleStorage.Client().from_service_account_json(service_account_json)
else: else:
self.folder = app.config.get('STORAGE_LOCAL_PATH') self.folder = app.config.get('STORAGE_LOCAL_PATH')
if not os.path.isabs(self.folder): if not os.path.isabs(self.folder):
@ -64,6 +70,10 @@ class Storage:
blob_container.upload_blob(filename, data) blob_container.upload_blob(filename, data)
elif self.storage_type == 'aliyun-oss': elif self.storage_type == 'aliyun-oss':
self.client.put_object(filename, data) self.client.put_object(filename, data)
elif self.storage_type == 'google-storage':
bucket = self.client.get_bucket(self.bucket_name)
blob = bucket.blob(filename)
blob.upload_from_file(data)
else: else:
if not self.folder or self.folder.endswith('/'): if not self.folder or self.folder.endswith('/'):
filename = self.folder + filename filename = self.folder + filename
@ -99,6 +109,10 @@ class Storage:
elif self.storage_type == 'aliyun-oss': elif self.storage_type == 'aliyun-oss':
with closing(self.client.get_object(filename)) as obj: with closing(self.client.get_object(filename)) as obj:
data = obj.read() data = obj.read()
elif self.storage_type == 'google-storage':
bucket = self.client.get_bucket(self.bucket_name)
blob = bucket.get_blob(filename)
data = blob.download_as_bytes()
else: else:
if not self.folder or self.folder.endswith('/'): if not self.folder or self.folder.endswith('/'):
filename = self.folder + filename filename = self.folder + filename
@ -135,6 +149,12 @@ class Storage:
with closing(self.client.get_object(filename)) as obj: with closing(self.client.get_object(filename)) as obj:
while chunk := obj.read(4096): while chunk := obj.read(4096):
yield chunk yield chunk
elif self.storage_type == 'google-storage':
bucket = self.client.get_bucket(self.bucket_name)
blob = bucket.get_blob(filename)
with closing(blob.open(mode='rb')) as blob_stream:
while chunk := blob_stream.read(4096):
yield chunk
else: else:
if not self.folder or self.folder.endswith('/'): if not self.folder or self.folder.endswith('/'):
filename = self.folder + filename filename = self.folder + filename
@ -161,6 +181,12 @@ class Storage:
blob_data.readinto(my_blob) blob_data.readinto(my_blob)
elif self.storage_type == 'aliyun-oss': elif self.storage_type == 'aliyun-oss':
self.client.get_object_to_file(filename, target_filepath) self.client.get_object_to_file(filename, target_filepath)
elif self.storage_type == 'google-storage':
bucket = self.client.get_bucket(self.bucket_name)
blob = bucket.get_blob(filename)
with open(target_filepath, "wb") as my_blob:
blob_data = blob.download_blob()
blob_data.readinto(my_blob)
else: else:
if not self.folder or self.folder.endswith('/'): if not self.folder or self.folder.endswith('/'):
filename = self.folder + filename filename = self.folder + filename
@ -185,6 +211,10 @@ class Storage:
return blob.exists() return blob.exists()
elif self.storage_type == 'aliyun-oss': elif self.storage_type == 'aliyun-oss':
return self.client.object_exists(filename) return self.client.object_exists(filename)
elif self.storage_type == 'google-storage':
bucket = self.client.get_bucket(self.bucket_name)
blob = bucket.blob(filename)
return blob.exists()
else: else:
if not self.folder or self.folder.endswith('/'): if not self.folder or self.folder.endswith('/'):
filename = self.folder + filename filename = self.folder + filename
@ -201,6 +231,9 @@ class Storage:
blob_container.delete_blob(filename) blob_container.delete_blob(filename)
elif self.storage_type == 'aliyun-oss': elif self.storage_type == 'aliyun-oss':
self.client.delete_object(filename) self.client.delete_object(filename)
elif self.storage_type == 'google-storage':
bucket = self.client.get_bucket(self.bucket_name)
bucket.delete_blob(filename)
else: else:
if not self.folder or self.folder.endswith('/'): if not self.folder or self.folder.endswith('/'):
filename = self.folder + filename filename = self.folder + filename

View File

@ -44,6 +44,7 @@ google-auth-httplib2==0.2.0
google-generativeai==0.5.0 google-generativeai==0.5.0
google-search-results==2.4.2 google-search-results==2.4.2
googleapis-common-protos==1.63.0 googleapis-common-protos==1.63.0
google-cloud-storage==2.16.0
replicate~=0.22.0 replicate~=0.22.0
websocket-client~=1.7.0 websocket-client~=1.7.0
dashscope[tokenizer]~=1.17.0 dashscope[tokenizer]~=1.17.0

View File

@ -70,7 +70,7 @@ services:
# If you want to enable cross-origin support, # If you want to enable cross-origin support,
# you must use the HTTPS protocol and set the configuration to `SameSite=None, Secure=true, HttpOnly=true`. # you must use the HTTPS protocol and set the configuration to `SameSite=None, Secure=true, HttpOnly=true`.
# #
# The type of storage to use for storing user files. Supported values are `local` and `s3` and `azure-blob`, Default: `local` # The type of storage to use for storing user files. Supported values are `local` and `s3` and `azure-blob` and `google-storage`, Default: `local`
STORAGE_TYPE: local STORAGE_TYPE: local
# The path to the local storage directory, the directory relative the root path of API service codes or absolute path. Default: `storage` or `/home/john/storage`. # The path to the local storage directory, the directory relative the root path of API service codes or absolute path. Default: `storage` or `/home/john/storage`.
# only available when STORAGE_TYPE is `local`. # only available when STORAGE_TYPE is `local`.
@ -86,6 +86,9 @@ services:
AZURE_BLOB_ACCOUNT_KEY: 'difyai' AZURE_BLOB_ACCOUNT_KEY: 'difyai'
AZURE_BLOB_CONTAINER_NAME: 'difyai-container' AZURE_BLOB_CONTAINER_NAME: 'difyai-container'
AZURE_BLOB_ACCOUNT_URL: 'https://<your_account_name>.blob.core.windows.net' AZURE_BLOB_ACCOUNT_URL: 'https://<your_account_name>.blob.core.windows.net'
# The Google storage configurations, only available when STORAGE_TYPE is `google-storage`.
GOOGLE_STORAGE_BUCKET_NAME: 'yout-bucket-name'
GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64: 'your-google-service-account-json-base64-string'
# The type of vector store to use. Supported values are `weaviate`, `qdrant`, `milvus`, `relyt`. # The type of vector store to use. Supported values are `weaviate`, `qdrant`, `milvus`, `relyt`.
VECTOR_STORE: weaviate VECTOR_STORE: weaviate
# The Weaviate endpoint URL. Only available when VECTOR_STORE is `weaviate`. # The Weaviate endpoint URL. Only available when VECTOR_STORE is `weaviate`.