diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py index da99fbb94d..f20cb2df53 100644 --- a/backend/open_webui/env.py +++ b/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. diff --git a/backend/open_webui/routers/audio.py b/backend/open_webui/routers/audio.py index 47dec5fa4c..1599852f3c 100644 --- a/backend/open_webui/routers/audio.py +++ b/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 diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index d26734da64..6820e96ab6 100644 --- a/backend/open_webui/routers/images.py +++ b/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 diff --git a/backend/open_webui/routers/ollama.py b/backend/open_webui/routers/ollama.py index eec46d6bc8..ff1bc3a19a 100644 --- a/backend/open_webui/routers/ollama.py +++ b/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}') diff --git a/backend/open_webui/routers/openai.py b/backend/open_webui/routers/openai.py index ba94d11a65..5d0f9062ed 100644 --- a/backend/open_webui/routers/openai.py +++ b/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) diff --git a/backend/open_webui/routers/pipelines.py b/backend/open_webui/routers/pipelines.py index 604e45375d..bffc7dd012 100644 --- a/backend/open_webui/routers/pipelines.py +++ b/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): diff --git a/backend/open_webui/routers/retrieval.py b/backend/open_webui/routers/retrieval.py index 3f62af4af4..2b9ec9d956 100644 --- a/backend/open_webui/routers/retrieval.py +++ b/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, diff --git a/backend/open_webui/utils/files.py b/backend/open_webui/utils/files.py index e43057c5d8..74527a40b6 100644 --- a/backend/open_webui/utils/files.py +++ b/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