OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-14 18:33:15 +00:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

chore: format backend

Browse Source
This commit is contained in:
Timothy Jaeryang Baek 2025-01-21 19:33:33 -08:00
parent 4c86f39a84
commit 8d3c73aed5
3 changed files with 25 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
backend/open_webui/config.py
View File

@ -663,7 +663,9 @@ S3_BUCKET_NAME = os.environ.get("S3_BUCKET_NAME", None)
S3_ENDPOINT_URL = os.environ.get("S3_ENDPOINT_URL", None) S3_ENDPOINT_URL = os.environ.get("S3_ENDPOINT_URL", None)
GCS_BUCKET_NAME = os.environ.get("GCS_BUCKET_NAME", None) GCS_BUCKET_NAME = os.environ.get("GCS_BUCKET_NAME", None)
GOOGLE_APPLICATION_CREDENTIALS_JSON = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON", None) GOOGLE_APPLICATION_CREDENTIALS_JSON = os.environ.get(
"GOOGLE_APPLICATION_CREDENTIALS_JSON", None
)
#################################### ####################################
# File Upload DIR # File Upload DIR

11
backend/open_webui/storage/provider.py
View File

@ -142,12 +142,15 @@ class S3StorageProvider(StorageProvider):
# Always delete from local storage # Always delete from local storage
LocalStorageProvider.delete_all_files() LocalStorageProvider.delete_all_files()
class GCSStorageProvider(StorageProvider): class GCSStorageProvider(StorageProvider):
def __init__(self): def __init__(self):
self.bucket_name = GCS_BUCKET_NAME self.bucket_name = GCS_BUCKET_NAME
if GOOGLE_APPLICATION_CREDENTIALS_JSON: if GOOGLE_APPLICATION_CREDENTIALS_JSON:
self.gcs_client = storage.Client.from_service_account_info(info=json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)) self.gcs_client = storage.Client.from_service_account_info(
info=json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
)
else: else:
# if no credentials json is provided, credentials will be picked up from the environment # if no credentials json is provided, credentials will be picked up from the environment
# if running on local environment, credentials would be user credentials # if running on local environment, credentials would be user credentials
@ -165,7 +168,7 @@ class GCSStorageProvider(StorageProvider):
except GoogleCloudError as e: except GoogleCloudError as e:
raise RuntimeError(f"Error uploading file to GCS: {e}") raise RuntimeError(f"Error uploading file to GCS: {e}")
def get_file(self, file_path:str) -> str: def get_file(self, file_path: str) -> str:
"""Handles downloading of the file from GCS storage.""" """Handles downloading of the file from GCS storage."""
try: try:
filename = file_path.removeprefix("gs://").split("/")[1] filename = file_path.removeprefix("gs://").split("/")[1]
@ -177,7 +180,7 @@ class GCSStorageProvider(StorageProvider):
except NotFound as e: except NotFound as e:
raise RuntimeError(f"Error downloading file from GCS: {e}") raise RuntimeError(f"Error downloading file from GCS: {e}")
def delete_file(self, file_path:str) -> None: def delete_file(self, file_path: str) -> None:
"""Handles deletion of the file from GCS storage.""" """Handles deletion of the file from GCS storage."""
try: try:
filename = file_path.removeprefix("gs://").split("/")[1] filename = file_path.removeprefix("gs://").split("/")[1]
@ -203,6 +206,7 @@ class GCSStorageProvider(StorageProvider):
# Always delete from local storage # Always delete from local storage
LocalStorageProvider.delete_all_files() LocalStorageProvider.delete_all_files()
def get_storage_provider(storage_provider: str): def get_storage_provider(storage_provider: str):
if storage_provider == "local": if storage_provider == "local":
Storage = LocalStorageProvider() Storage = LocalStorageProvider()
@ -214,4 +218,5 @@ def get_storage_provider(storage_provider: str):
raise RuntimeError(f"Unsupported storage provider: {storage_provider}") raise RuntimeError(f"Unsupported storage provider: {storage_provider}")
return Storage return Storage
Storage = get_storage_provider(STORAGE_PROVIDER) Storage = get_storage_provider(STORAGE_PROVIDER)

7
backend/open_webui/test/apps/webui/storage/test_provider.py
View File

@ -104,7 +104,6 @@ class TestS3StorageProvider:
self.file_bytesio_empty = io.BytesIO() self.file_bytesio_empty = io.BytesIO()
super().__init__() super().__init__()
def test_upload_file(self, monkeypatch, tmp_path): def test_upload_file(self, monkeypatch, tmp_path):
upload_dir = mock_upload_dir(monkeypatch, tmp_path) upload_dir = mock_upload_dir(monkeypatch, tmp_path)
# S3 checks # S3 checks
@ -182,6 +181,7 @@ class TestS3StorageProvider:
assert not (upload_dir / self.filename).exists() assert not (upload_dir / self.filename).exists()
assert not (upload_dir / self.filename_extra).exists() assert not (upload_dir / self.filename_extra).exists()
class TestGCSStorageProvider: class TestGCSStorageProvider:
Storage = provider.GCSStorageProvider() Storage = provider.GCSStorageProvider()
Storage.bucket_name = "my-bucket" Storage.bucket_name = "my-bucket"
@ -261,7 +261,10 @@ class TestGCSStorageProvider:
object = self.Storage.bucket.get_blob(self.filename_extra) object = self.Storage.bucket.get_blob(self.filename_extra)
assert (upload_dir / self.filename_extra).exists() assert (upload_dir / self.filename_extra).exists()
assert (upload_dir / self.filename_extra).read_bytes() == self.file_content assert (upload_dir / self.filename_extra).read_bytes() == self.file_content
assert self.Storage.bucket.get_blob(self.filename_extra).name == self.filename_extra assert (
self.Storage.bucket.get_blob(self.filename_extra).name
== self.filename_extra
)
assert self.file_content == object.download_as_bytes() assert self.file_content == object.download_as_bytes()
self.Storage.delete_all_files() self.Storage.delete_all_files()