mirror of
https://github.com/open-webui/open-webui.git
synced 2026-07-10 12:01:29 -05:00
refac
This commit is contained in:
@@ -915,6 +915,7 @@ if CUSTOM_NAME:
|
||||
####################################
|
||||
|
||||
STORAGE_PROVIDER = os.environ.get('STORAGE_PROVIDER', 'local') # defaults to local, s3
|
||||
STORAGE_LOCAL_CACHE = os.environ.get('STORAGE_LOCAL_CACHE', 'true').lower() == 'true'
|
||||
|
||||
S3_ACCESS_KEY_ID = os.environ.get('S3_ACCESS_KEY_ID', None)
|
||||
S3_SECRET_ACCESS_KEY = os.environ.get('S3_SECRET_ACCESS_KEY', None)
|
||||
|
||||
@@ -48,7 +48,7 @@ from open_webui.routers.audio import transcribe
|
||||
from open_webui.storage.provider import Storage
|
||||
|
||||
|
||||
from open_webui.config import BYPASS_ADMIN_ACCESS_CONTROL
|
||||
from open_webui.config import BYPASS_ADMIN_ACCESS_CONTROL, STORAGE_LOCAL_CACHE, STORAGE_PROVIDER, UPLOAD_DIR
|
||||
from open_webui.utils.auth import get_admin_user, get_verified_user
|
||||
from open_webui.utils.misc import strict_match_mime_type
|
||||
from pydantic import BaseModel
|
||||
@@ -88,6 +88,20 @@ def _is_text_file(file_path: str, chunk_size: int = 8192) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _cleanup_local_cache(file_path: str) -> None:
|
||||
"""Remove the local cached copy of a cloud-stored file after processing."""
|
||||
if STORAGE_LOCAL_CACHE or STORAGE_PROVIDER == 'local':
|
||||
return
|
||||
try:
|
||||
local_filename = os.path.basename(file_path)
|
||||
local_path = os.path.join(UPLOAD_DIR, local_filename)
|
||||
if os.path.isfile(local_path):
|
||||
os.remove(local_path)
|
||||
log.debug(f'Cleaned up local cache: {local_path}')
|
||||
except OSError as e:
|
||||
log.warning(f'Failed to clean up local cache for {file_path}: {e}')
|
||||
|
||||
|
||||
async def process_uploaded_file(
|
||||
request,
|
||||
file,
|
||||
@@ -150,11 +164,14 @@ async def process_uploaded_file(
|
||||
db=db_session,
|
||||
)
|
||||
|
||||
if db:
|
||||
await _process_handler(db)
|
||||
else:
|
||||
async with get_async_db_context() as db_session:
|
||||
await _process_handler(db_session)
|
||||
try:
|
||||
if db:
|
||||
await _process_handler(db)
|
||||
else:
|
||||
async with get_async_db_context() as db_session:
|
||||
await _process_handler(db_session)
|
||||
finally:
|
||||
_cleanup_local_cache(file_path)
|
||||
|
||||
|
||||
@router.post('/', response_model=FileModelResponse)
|
||||
|
||||
@@ -140,7 +140,7 @@ class S3StorageProvider(StorageProvider):
|
||||
|
||||
def upload_file(self, file: BinaryIO, filename: str, tags: Dict[str, str]) -> Tuple[bytes, str]:
|
||||
"""Handles uploading of the file to S3 storage."""
|
||||
_, file_path = LocalStorageProvider.upload_file(file, filename, tags)
|
||||
contents, file_path = LocalStorageProvider.upload_file(file, filename, tags)
|
||||
s3_key = os.path.join(self.key_prefix, filename)
|
||||
try:
|
||||
self.s3_client.upload_file(file_path, self.bucket_name, s3_key)
|
||||
@@ -153,7 +153,7 @@ class S3StorageProvider(StorageProvider):
|
||||
Tagging=tagging,
|
||||
)
|
||||
return (
|
||||
open(file_path, 'rb').read(),
|
||||
contents,
|
||||
f's3://{self.bucket_name}/{s3_key}',
|
||||
)
|
||||
except ClientError as e:
|
||||
|
||||
Reference in New Issue
Block a user