refac

backend/open_webui/env.py

@@ -594,6 +594,15 @@ try:
 except (ValueError, TypeError):
     AIOHTTP_CLIENT_TIMEOUT_TOOL_SERVER_DATA = 10
 
+AIOHTTP_FILE_STREAM_CHUNK_SIZE = os.getenv('AIOHTTP_FILE_STREAM_CHUNK_SIZE', str(1024 * 1024))
+try:
+    AIOHTTP_FILE_STREAM_CHUNK_SIZE = int(AIOHTTP_FILE_STREAM_CHUNK_SIZE)
+except Exception:
+    AIOHTTP_FILE_STREAM_CHUNK_SIZE = 1024 * 1024
+
+if AIOHTTP_FILE_STREAM_CHUNK_SIZE <= 0:
+    AIOHTTP_FILE_STREAM_CHUNK_SIZE = 1024 * 1024
+
 
 # SSL verification for tool server connections specifically.
 # Accepts "True", "False", or a path to a CA bundle file.

backend/open_webui/routers/audio.py

@@ -49,6 +49,7 @@ from open_webui.env import (
     AIOHTTP_CLIENT_SESSION_SSL,
     AIOHTTP_CLIENT_TIMEOUT,
     AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST,
+    AIOHTTP_FILE_STREAM_CHUNK_SIZE,
     BYPASS_PYDUB_PREPROCESSING,
     DEVICE_TYPE,
     ENABLE_FORWARD_USER_INFO_HEADERS,
@@ -678,15 +679,19 @@ async def _transcribe_openai(request, file_path, filename, languages, file_dir,
                 for key, value in payload.items():
                     form_data.add_field(key, str(value))
 
-                with open(file_path, 'rb') as audio_file:
-                    form_data.add_field('file', audio_file, filename=filename)
+                async def audio_chunks():
+                    async with aiofiles.open(file_path, 'rb') as audio_file:
+                        while chunk := await audio_file.read(AIOHTTP_FILE_STREAM_CHUNK_SIZE):
+                            yield chunk
 
-                    r = await session.post(
-                        url=f'{api_base_url}/audio/transcriptions',
-                        headers=headers,
-                        data=form_data,
-                        ssl=AIOHTTP_CLIENT_SESSION_SSL,
-                    )
+                form_data.add_field('file', audio_chunks(), filename=filename)
+
+                r = await session.post(
+                    url=f'{api_base_url}/audio/transcriptions',
+                    headers=headers,
+                    data=form_data,
+                    ssl=AIOHTTP_CLIENT_SESSION_SSL,
+                )
             if r.status == 200:
                 break
 
@@ -823,13 +828,18 @@ async def _transcribe_azure(request, file_path, filename, file_dir, id):
         base_url or f'https://{region}.api.cognitive.microsoft.com'
     ) + '/speechtotext/transcriptions:transcribe?api-version=2024-11-15'
 
-    form_data = aiohttp.FormData()
-    form_data.add_field('definition', definition)
-    form_data.add_field('audio', open(file_path, 'rb'), filename=filename)
-
     r = None
     try:
         session = await get_session()
+        form_data = aiohttp.FormData()
+        form_data.add_field('definition', definition)
+
+        async def audio_chunks():
+            async with aiofiles.open(file_path, 'rb') as audio_file:
+                while chunk := await audio_file.read(AIOHTTP_FILE_STREAM_CHUNK_SIZE):
+                    yield chunk
+
+        form_data.add_field('audio', audio_chunks(), filename=filename)
         r = await session.post(
             url=endpoint,
             data=form_data,
@@ -1002,7 +1012,12 @@ async def _transcribe_mistral(request, file_path, filename, metadata, file_dir,
             if language:
                 form_data.add_field('language', language)
 
-            form_data.add_field('file', open(file_path, 'rb'), filename=filename, content_type=mime_type)
+            async def audio_chunks():
+                async with aiofiles.open(file_path, 'rb') as audio_file:
+                    while chunk := await audio_file.read(AIOHTTP_FILE_STREAM_CHUNK_SIZE):
+                        yield chunk
+
+            form_data.add_field('file', audio_chunks(), filename=filename, content_type=mime_type)
 
             r = await session.post(
                 url=f'{api_base_url}/audio/transcriptions',
@@ -1094,7 +1109,7 @@ async def transcribe(request: Request, file_path: str, metadata: Optional[dict]
         for chunk_path in chunk_paths:
             if chunk_path != file_path and os.path.isfile(chunk_path):
                 try:
-                    os.remove(chunk_path)
+                    await asyncio.to_thread(os.remove, chunk_path)
                 except Exception:
                     pass
 
@@ -1208,12 +1223,8 @@ async def transcription(
         if not os.path.realpath(file_path).startswith(os.path.realpath(file_dir)):
             raise ValueError('Invalid file path detected')
 
-        def _write_upload():
-            with open(file_path, 'wb') as f:
-                f.write(contents)
-
-        # Audio uploads can be large; write to disk off the event loop.
-        await asyncio.to_thread(_write_upload)
+        async with aiofiles.open(file_path, 'wb') as f:
+            await f.write(contents)
 
         try:
             metadata = None

backend/open_webui/routers/images.py

@@ -13,6 +13,7 @@ from types import SimpleNamespace
 from typing import Optional
 from urllib.parse import quote, urlparse
 
+import aiofiles
 import aiohttp
 from fastapi import APIRouter, Depends, HTTPException, Request, UploadFile
 from fastapi.responses import FileResponse
@@ -940,10 +941,10 @@ async def image_edits(
                 if isinstance(file_response, FileResponse):
                     file_path = file_response.path
 
-                    with open(file_path, 'rb') as f:
-                        file_bytes = f.read()
-                        image_data = base64.b64encode(file_bytes).decode('utf-8')
-                        mime_type, _ = mimetypes.guess_type(file_path)
+                    async with aiofiles.open(file_path, 'rb') as f:
+                        file_bytes = await f.read()
+                    image_data = base64.b64encode(file_bytes).decode('utf-8')
+                    mime_type, _ = mimetypes.guess_type(file_path)
 
                     return f'data:{mime_type};base64,{image_data}'
             return data

backend/open_webui/routers/ollama.py

@@ -11,6 +11,7 @@ from datetime import datetime
 from typing import Optional, Union
 from urllib.parse import urlparse
 
+import aiofiles
 import aiohttp
 from aiocache import cached
 from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile
@@ -25,6 +26,7 @@ from open_webui.env import (
     AIOHTTP_CLIENT_SESSION_SSL,
     AIOHTTP_CLIENT_TIMEOUT,
     AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST,
+    AIOHTTP_FILE_STREAM_CHUNK_SIZE,
     BYPASS_MODEL_ACCESS_CONTROL,
     ENABLE_FORWARD_USER_INFO_HEADERS,
     FORWARD_SESSION_INFO_HEADER_CHAT_ID,
@@ -1529,7 +1531,7 @@ async def download_file_stream(
     file_url: str,
     file_path: str,
     file_name: str,
-    chunk_size: int = 1024 * 1024,
+    chunk_size: int = AIOHTTP_FILE_STREAM_CHUNK_SIZE,
 ):
     """Stream a model file download from *file_url*, then push the blob to Ollama."""
     current_size = os.path.getsize(file_path) if os.path.exists(file_path) else 0
@@ -1544,37 +1546,38 @@ async def download_file_stream(
     ) as response:
         total_size = int(response.headers.get('content-length', 0)) + current_size
 
-        with open(file_path, 'ab+') as f:
+        async with aiofiles.open(file_path, 'ab') as f:
             async for data in response.content.iter_chunked(chunk_size):
                 current_size += len(data)
-                f.write(data)
+                await f.write(data)
 
-                done = current_size == total_size
-                progress = round((current_size / total_size) * 100, 2)
+                progress_total = total_size or current_size
+                progress = round((current_size / progress_total) * 100, 2)
                 yield f'data: {{"progress": {progress}, "completed": {current_size}, "total": {total_size}}}\n\n'
 
-            if done:
-                f.close()
-                hashed = await asyncio.to_thread(calculate_sha256, file_path, chunk_size)
+            done = True
+            hashed = await asyncio.to_thread(calculate_sha256, file_path, chunk_size)
 
-                def _read_blob():
-                    with open(file_path, 'rb') as blob_f:
-                        return blob_f.read()
+            blob_url = f'{ollama_url}/api/blobs/sha256:{hashed}'
+            blob_size = await asyncio.to_thread(os.path.getsize, file_path)
 
-                blob_data = await asyncio.to_thread(_read_blob)
+            async def blob_chunks():
+                async with aiofiles.open(file_path, 'rb') as blob_file:
+                    while chunk := await blob_file.read(chunk_size):
+                        yield chunk
 
-                blob_url = f'{ollama_url}/api/blobs/sha256:{hashed}'
-                async with session.post(
-                    blob_url,
-                    data=blob_data,
-                    ssl=AIOHTTP_CLIENT_SESSION_SSL,
-                    timeout=aiohttp.ClientTimeout(total=30),
-                ) as blob_resp:
-                    if blob_resp.ok:
-                        os.remove(file_path)
-                        yield f'data: {json.dumps({"done": done, "blob": f"sha256:{hashed}", "name": file_name})}\n\n'
-                    else:
-                        raise RuntimeError('Ollama: Could not create blob, Please try again.')
+            async with session.post(
+                blob_url,
+                data=blob_chunks(),
+                headers={'Content-Length': str(blob_size)},
+                ssl=AIOHTTP_CLIENT_SESSION_SSL,
+                timeout=aiohttp.ClientTimeout(total=30),
+            ) as blob_resp:
+                if blob_resp.ok:
+                    await asyncio.to_thread(os.remove, file_path)
+                    yield f'data: {json.dumps({"done": done, "blob": f"sha256:{hashed}", "name": file_name})}\n\n'
+                else:
+                    raise RuntimeError('Ollama: Could not create blob, Please try again.')
 
 
 @router.post('/models/download')
@@ -1621,17 +1624,11 @@ async def upload_model(
     os.makedirs(UPLOAD_DIR, exist_ok=True)
 
     # Stage 1: persist the uploaded file to disk
-    chunk_size = 1024 * 1024 * 2  # 2 MiB
+    chunk_size = AIOHTTP_FILE_STREAM_CHUNK_SIZE
 
-    def _persist_upload():
-        with open(file_path, 'wb') as out_f:
-            while True:
-                chunk = file.file.read(chunk_size)
-                if not chunk:
-                    break
-                out_f.write(chunk)
-
-    await asyncio.to_thread(_persist_upload)
+    async with aiofiles.open(file_path, 'wb') as out_f:
+        while chunk := await file.read(chunk_size):
+            await out_f.write(chunk)
 
     async def file_process_stream():
         nonlocal ollama_url
@@ -1643,25 +1640,26 @@ async def upload_model(
         log.info(f'Model Hash: {file_hash}')
 
         try:
-            with open(file_path, 'rb') as f:
-                bytes_read = 0
-                while chunk := f.read(chunk_size):
+            bytes_read = 0
+            async with aiofiles.open(file_path, 'rb') as f:
+                while chunk := await f.read(chunk_size):
                     bytes_read += len(chunk)
                     progress = round(bytes_read / total_size * 100, 2)
-                    yield f'data: {json.dumps({"progress": progress, "total": total_size, "completed": bytes_read})}\n\n'
-
-            # Stage 3: push blob to Ollama
-            def _read_blob():
-                with open(file_path, 'rb') as f:
-                    return f.read()
-
-            blob_data = await asyncio.to_thread(_read_blob)
+                    event = json.dumps({'progress': progress, 'total': total_size, 'completed': bytes_read})
+                    yield f'data: {event}\n\n'
 
             session = await get_session()
             blob_url = f'{ollama_url}/api/blobs/sha256:{file_hash}'
+
+            async def blob_chunks():
+                async with aiofiles.open(file_path, 'rb') as blob_file:
+                    while chunk := await blob_file.read(chunk_size):
+                        yield chunk
+
             async with session.post(
                 blob_url,
-                data=blob_data,
+                data=blob_chunks(),
+                headers={'Content-Length': str(total_size)},
                 ssl=AIOHTTP_CLIENT_SESSION_SSL,
                 timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
             ) as resp:
@@ -1669,7 +1667,7 @@ async def upload_model(
                     raise Exception('Ollama: Could not create blob, Please try again.')
 
             log.info('Uploaded to /api/blobs')
-            os.remove(file_path)
+            await asyncio.to_thread(os.remove, file_path)
 
             # Stage 4: create the model
             model, _ext = os.path.splitext(filename)
@@ -1690,7 +1688,10 @@ async def upload_model(
             ) as create_resp:
                 if create_resp.ok:
                     log.info('API SUCCESS!')
-                    yield f'data: {json.dumps({"done": True, "blob": f"sha256:{file_hash}", "name": filename, "model_created": model})}\n\n'
+                    event = json.dumps(
+                        {'done': True, 'blob': f'sha256:{file_hash}', 'name': filename, 'model_created': model}
+                    )
+                    yield f'data: {event}\n\n'
                 else:
                     resp_text = await create_resp.text()
                     raise Exception(f'Failed to create model in Ollama. {resp_text}')

backend/open_webui/routers/openai.py

@@ -8,6 +8,7 @@ import re
 from typing import Optional
 from urllib.parse import quote, urlparse
 
+import aiofiles
 import aiohttp
 from aiocache import cached
 from azure.identity import DefaultAzureCredential, get_bearer_token_provider
@@ -491,13 +492,12 @@ async def speech(request: Request, user=Depends(get_verified_user)):
 
             r.raise_for_status()
 
-            # Save the streaming content to a file
-            with open(file_path, 'wb') as f:
+            async with aiofiles.open(file_path, 'wb') as f:
                 async for chunk in r.content.iter_chunked(8192):
-                    f.write(chunk)
+                    await f.write(chunk)
 
-            with open(file_body_path, 'w') as f:
-                json.dump(json.loads(body.decode('utf-8')), f)
+            async with aiofiles.open(file_body_path, 'w') as f:
+                await f.write(json.dumps(json.loads(body.decode('utf-8'))))
 
             # Return the saved file
             return FileResponse(file_path)

backend/open_webui/routers/pipelines.py

@@ -1,9 +1,9 @@
 import asyncio
 import logging
 import os
-import shutil
 from typing import Optional
 
+import aiofiles
 import aiohttp
 from fastapi import (
     APIRouter,
@@ -18,7 +18,7 @@ from fastapi import (
 )
 from open_webui.config import CACHE_DIR
 from open_webui.constants import ERROR_MESSAGES
-from open_webui.env import AIOHTTP_CLIENT_SESSION_SSL
+from open_webui.env import AIOHTTP_CLIENT_SESSION_SSL, AIOHTTP_FILE_STREAM_CHUNK_SIZE
 from open_webui.events import EVENTS, publish_event
 from open_webui.models.config import Config
 from open_webui.routers.openai import get_all_models_responses
@@ -237,35 +237,37 @@ async def upload_pipeline(
 
     response = None
     try:
-        # Save the uploaded file off the event loop (uploads can be large).
-        def _save_upload():
-            with open(file_path, 'wb') as buffer:
-                shutil.copyfileobj(file.file, buffer)
-
-        await asyncio.to_thread(_save_upload)
+        async with aiofiles.open(file_path, 'wb') as buffer:
+            while chunk := await file.read(AIOHTTP_FILE_STREAM_CHUNK_SIZE):
+                await buffer.write(chunk)
 
         url, key = await get_openai_connection(urlIdx)
 
         headers = {'Authorization': f'Bearer {key}'}
 
         async with aiohttp.ClientSession(trust_env=True) as session:
-            with open(file_path, 'rb') as f:
-                form_data = aiohttp.FormData()
-                form_data.add_field(
-                    'file',
-                    f,
-                    filename=filename,
-                    content_type='application/octet-stream',
-                )
+            form_data = aiohttp.FormData()
 
-                async with session.post(
-                    f'{url}/pipelines/upload',
-                    headers=headers,
-                    data=form_data,
-                    ssl=AIOHTTP_CLIENT_SESSION_SSL,
-                ) as response:
-                    response.raise_for_status()
-                    data = await response.json()
+            async def pipeline_chunks():
+                async with aiofiles.open(file_path, 'rb') as pipeline_file:
+                    while chunk := await pipeline_file.read(AIOHTTP_FILE_STREAM_CHUNK_SIZE):
+                        yield chunk
+
+            form_data.add_field(
+                'file',
+                pipeline_chunks(),
+                filename=filename,
+                content_type='application/octet-stream',
+            )
+
+            async with session.post(
+                f'{url}/pipelines/upload',
+                headers=headers,
+                data=form_data,
+                ssl=AIOHTTP_CLIENT_SESSION_SSL,
+            ) as response:
+                response.raise_for_status()
+                data = await response.json()
 
         await publish_event(
             request,
@@ -297,7 +299,7 @@ async def upload_pipeline(
     finally:
         # Ensure the file is deleted after the upload is completed or on failure
         if os.path.exists(file_path):
-            os.remove(file_path)
+            await asyncio.to_thread(os.remove, file_path)
 
 
 class AddPipelineForm(BaseModel):

backend/open_webui/routers/retrieval.py

@@ -2910,21 +2910,24 @@ async def reset_upload_dir(request: Request, user=Depends(get_admin_user)) -> bo
     folder = f'{UPLOAD_DIR}'
     try:
         # Check if the directory exists
-        if os.path.exists(folder):
+        if await asyncio.to_thread(os.path.exists, folder):
             # Iterate over all the files and directories in the specified directory
-            for filename in os.listdir(folder):
+            for filename in await asyncio.to_thread(os.listdir, folder):
                 file_path = os.path.join(folder, filename)
                 try:
-                    if os.path.isfile(file_path) or os.path.islink(file_path):
-                        os.unlink(file_path)  # Remove the file or link
-                    elif os.path.isdir(file_path):
-                        shutil.rmtree(file_path)  # Remove the directory
+                    if await asyncio.to_thread(os.path.isfile, file_path) or await asyncio.to_thread(
+                        os.path.islink, file_path
+                    ):
+                        await asyncio.to_thread(os.unlink, file_path)  # Remove the file or link
+                    elif await asyncio.to_thread(os.path.isdir, file_path):
+                        await asyncio.to_thread(shutil.rmtree, file_path)  # Remove the directory
                 except Exception as e:
                     log.exception(f'Failed to delete {file_path}. Reason: {e}')
         else:
             log.warning(f'The directory {folder} does not exist')
     except Exception as e:
         log.exception(f'Failed to process the directory {folder}. Reason: {e}')
+
     await publish_event(
         request,
         EVENTS.RETRIEVAL_UPLOADS_RESET,

backend/open_webui/utils/files.py

@@ -6,6 +6,7 @@ import re
 from pathlib import Path
 from typing import Optional
 
+import aiofiles
 from fastapi import (
     APIRouter,
     Depends,
@@ -75,7 +76,7 @@ async def get_image_base64_from_url(url: str, user=None) -> Optional[str]:
             # file-ID resolver which enforces ownership/access checks.
             return await get_image_base64_from_file_id(url, user=user)
 
-    except Exception as e:
+    except Exception:
         return None
 
 
@@ -200,15 +201,15 @@ async def get_image_base64_from_file_id(id: str, user=None) -> Optional[str]:
 
         # Check if the file already exists in the cache
         if file_path.is_file():
-            with open(file_path, 'rb') as image_file:
-                encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
-                content_type = mimetypes.guess_type(file_path.name)[0] or (file.meta or {}).get('content_type')
-                if not content_type and ENABLE_IMAGE_CONTENT_TYPE_EXTENSION_FALLBACK:
-                    content_type = _IMAGE_MIME_FALLBACK.get(file_path.suffix.lower())
-                if not content_type:
-                    return None
-                return f'data:{content_type};base64,{encoded_string}'
+            async with aiofiles.open(file_path, 'rb') as image_file:
+                encoded_string = base64.b64encode(await image_file.read()).decode('utf-8')
+            content_type = mimetypes.guess_type(file_path.name)[0] or (file.meta or {}).get('content_type')
+            if not content_type and ENABLE_IMAGE_CONTENT_TYPE_EXTENSION_FALLBACK:
+                content_type = _IMAGE_MIME_FALLBACK.get(file_path.suffix.lower())
+            if not content_type:
+                return None
+            return f'data:{content_type};base64,{encoded_string}'
         else:
             return None
-    except Exception as e:
+    except Exception:
         return None