refac

Author: tjbck

backend/open_webui/config.py

@@ -915,6 +915,7 @@ def feishu_oauth_register(oauth: OAuth):
 ####################################
 
 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)

backend/open_webui/routers/files.py

@@ -48,7 +48,7 @@
 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_handler(db_session):
                 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)

backend/open_webui/storage/provider.py

@@ -140,7 +140,7 @@ def sanitize_tag_value(s: str) -> str:
 
     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 @@ def upload_file(self, file: BinaryIO, filename: str, tags: Dict[str, str]) -> Tu
                     Tagging=tagging,
                 )
             return (
-                open(file_path, 'rb').read(),
+                contents,
                 f's3://{self.bucket_name}/{s3_key}',
             )
         except ClientError as e: