diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index b599ff575d..fe25dda1c1 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -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) diff --git a/backend/open_webui/routers/files.py b/backend/open_webui/routers/files.py index b5e8ea83ce..66d7539278 100644 --- a/backend/open_webui/routers/files.py +++ b/backend/open_webui/routers/files.py @@ -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) diff --git a/backend/open_webui/storage/provider.py b/backend/open_webui/storage/provider.py index 0886c851f0..f70f3e862b 100644 --- a/backend/open_webui/storage/provider.py +++ b/backend/open_webui/storage/provider.py @@ -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: