diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 0bb731188b..fce771208f 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -21,7 +21,7 @@ from typing import Optional from aiocache import cached import aiohttp import anyio.to_thread -import requests + from redis import Redis @@ -60,6 +60,7 @@ from starsessions.stores.redis import RedisStore from open_webui.utils import logger from open_webui.utils.audit import AuditLevel, AuditLoggingMiddleware from open_webui.utils.logger import start_logger +from open_webui.utils.session_pool import get_session from open_webui.socket.main import ( MODELS, app as socket_app, @@ -2512,7 +2513,13 @@ async def oauth_backchannel_logout( @app.get('/manifest.json') async def get_manifest_json(): if app.state.EXTERNAL_PWA_MANIFEST_URL: - return requests.get(app.state.EXTERNAL_PWA_MANIFEST_URL).json() + session = await get_session() + async with session.get( + app.state.EXTERNAL_PWA_MANIFEST_URL, + ssl=AIOHTTP_CLIENT_SESSION_SSL, + ) as r: + r.raise_for_status() + return await r.json() else: return { 'name': app.state.WEBUI_NAME, diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index 0d534db7f6..832022e92d 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -10,6 +10,7 @@ from pathlib import Path from typing import Optional from urllib.parse import quote +import aiohttp import requests from fastapi import APIRouter, Depends, HTTPException, Request, UploadFile from fastapi.responses import FileResponse @@ -21,7 +22,8 @@ from open_webui.config import ( ) from open_webui.constants import ERROR_MESSAGES from open_webui.retrieval.web.utils import validate_url -from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS +from open_webui.env import AIOHTTP_CLIENT_SESSION_SSL, ENABLE_FORWARD_USER_INFO_HEADERS +from open_webui.utils.session_pool import get_session from open_webui.models.chats import Chats from open_webui.routers.files import upload_file_handler, get_file_content_by_id @@ -313,12 +315,14 @@ def get_automatic1111_api_auth(request: Request): async def verify_url(request: Request, user=Depends(get_admin_user)): if request.app.state.config.IMAGE_GENERATION_ENGINE == 'automatic1111': try: - r = requests.get( + session = await get_session() + async with session.get( url=f'{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options', headers={'authorization': get_automatic1111_api_auth(request)}, - ) - r.raise_for_status() - return True + ssl=AIOHTTP_CLIENT_SESSION_SSL, + ) as r: + r.raise_for_status() + return True except Exception: request.app.state.config.ENABLE_IMAGE_GENERATION = False raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL) @@ -327,12 +331,14 @@ async def verify_url(request: Request, user=Depends(get_admin_user)): if request.app.state.config.COMFYUI_API_KEY: headers = {'Authorization': f'Bearer {request.app.state.config.COMFYUI_API_KEY}'} try: - r = requests.get( + session = await get_session() + async with session.get( url=f'{request.app.state.config.COMFYUI_BASE_URL}/object_info', headers=headers, - ) - r.raise_for_status() - return True + ssl=AIOHTTP_CLIENT_SESSION_SSL, + ) as r: + r.raise_for_status() + return True except Exception: request.app.state.config.ENABLE_IMAGE_GENERATION = False raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL) @@ -357,11 +363,13 @@ async def get_models(request: Request, user=Depends(get_verified_user)): elif request.app.state.config.IMAGE_GENERATION_ENGINE == 'comfyui': # TODO - get models from comfyui headers = {'Authorization': f'Bearer {request.app.state.config.COMFYUI_API_KEY}'} - r = requests.get( + session = await get_session() + async with session.get( url=f'{request.app.state.config.COMFYUI_BASE_URL}/object_info', headers=headers, - ) - info = r.json() + ssl=AIOHTTP_CLIENT_SESSION_SSL, + ) as r: + info = await r.json() workflow = json.loads(request.app.state.config.COMFYUI_WORKFLOW) model_node_id = None @@ -399,11 +407,13 @@ async def get_models(request: Request, user=Depends(get_verified_user)): request.app.state.config.IMAGE_GENERATION_ENGINE == 'automatic1111' or request.app.state.config.IMAGE_GENERATION_ENGINE == '' ): - r = requests.get( + session = await get_session() + async with session.get( url=f'{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/sd-models', headers={'authorization': get_automatic1111_api_auth(request)}, - ) - models = r.json() + ssl=AIOHTTP_CLIENT_SESSION_SSL, + ) as r: + models = await r.json() return list( map( lambda model: {'id': model['title'], 'name': model['model_name']}, @@ -533,7 +543,7 @@ async def image_generations( model = get_image_model(request) - r = None + try: if request.app.state.config.IMAGE_GENERATION_ENGINE == 'openai': headers = { @@ -568,16 +578,15 @@ async def image_generations( ), } - # Use asyncio.to_thread for the requests.post call - r = await asyncio.to_thread( - requests.post, + session = await get_session() + async with session.post( url=url, json=data, headers=headers, - ) - - r.raise_for_status() - res = r.json() + ssl=AIOHTTP_CLIENT_SESSION_SSL, + ) as r: + r.raise_for_status() + res = await r.json() images = [] @@ -619,16 +628,15 @@ async def image_generations( model = f'{model}:generateContent' data = {'contents': [{'parts': [{'text': form_data.prompt}]}]} - # Use asyncio.to_thread for the requests.post call - r = await asyncio.to_thread( - requests.post, + session = await get_session() + async with session.post( url=f'{request.app.state.config.IMAGES_GEMINI_API_BASE_URL}/models/{model}', json=data, headers=headers, - ) - - r.raise_for_status() - res = r.json() + ssl=AIOHTTP_CLIENT_SESSION_SSL, + ) as r: + r.raise_for_status() + res = await r.json() images = [] @@ -727,15 +735,14 @@ async def image_generations( if request.app.state.config.AUTOMATIC1111_PARAMS: data = {**data, **request.app.state.config.AUTOMATIC1111_PARAMS} - # Use asyncio.to_thread for the requests.post call - r = await asyncio.to_thread( - requests.post, + session = await get_session() + async with session.post( url=f'{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/txt2img', json=data, headers={'authorization': get_automatic1111_api_auth(request)}, - ) - - res = r.json() + ssl=AIOHTTP_CLIENT_SESSION_SSL, + ) as r: + res = await r.json() log.debug(f'res: {res}') images = [] @@ -753,10 +760,8 @@ async def image_generations( return images except Exception as e: error = e - if r != None: - data = r.json() - if 'error' in data: - error = data['error']['message'] + if isinstance(e, aiohttp.ClientResponseError): + error = e.message raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error)) @@ -798,11 +803,12 @@ async def image_edits( if data.startswith('http://') or data.startswith('https://'): # Validate URL to prevent SSRF attacks against local/private networks validate_url(data) - r = await asyncio.to_thread(requests.get, data) - r.raise_for_status() + session = await get_session() + async with session.get(data, ssl=AIOHTTP_CLIENT_SESSION_SSL) as r: + r.raise_for_status() - image_data = base64.b64encode(r.content).decode('utf-8') - return f'data:{r.headers["content-type"]};base64,{image_data}' + image_data = base64.b64encode(await r.read()).decode('utf-8') + return f'data:{r.headers["content-type"]};base64,{image_data}' else: file_id = None @@ -846,7 +852,7 @@ async def image_edits( ), ) - r = None + try: if request.app.state.config.IMAGE_EDIT_ENGINE == 'openai': headers = { @@ -883,17 +889,30 @@ async def image_edits( if request.app.state.config.IMAGES_EDIT_OPENAI_API_VERSION: url_search_params += f'?api-version={request.app.state.config.IMAGES_EDIT_OPENAI_API_VERSION}' - # Use asyncio.to_thread for the requests.post call - r = await asyncio.to_thread( - requests.post, + # Build multipart form data for aiohttp + form = aiohttp.FormData() + for key, value in data.items(): + if isinstance(value, dict): + form.add_field(key, json.dumps(value)) + else: + form.add_field(key, str(value)) + for param_name, (filename, file_obj, content_type_val) in files: + form.add_field( + param_name, + file_obj, + filename=filename, + content_type=content_type_val, + ) + + session = await get_session() + async with session.post( url=f'{request.app.state.config.IMAGES_EDIT_OPENAI_API_BASE_URL}/images/edits{url_search_params}', headers=headers, - files=files, - data=data, - ) - - r.raise_for_status() - res = r.json() + data=form, + ssl=AIOHTTP_CLIENT_SESSION_SSL, + ) as r: + r.raise_for_status() + res = await r.json() images = [] for image in res['data']: @@ -940,16 +959,15 @@ async def image_edits( ] ) - # Use asyncio.to_thread for the requests.post call - r = await asyncio.to_thread( - requests.post, + session = await get_session() + async with session.post( url=f'{request.app.state.config.IMAGES_EDIT_GEMINI_API_BASE_URL}/models/{model}', json=data, headers=headers, - ) - - r.raise_for_status() - res = r.json() + ssl=AIOHTTP_CLIENT_SESSION_SSL, + ) as r: + r.raise_for_status() + res = await r.json() images = [] for image in res['candidates']: @@ -1048,13 +1066,7 @@ async def image_edits( return images except Exception as e: error = e - if r != None: - data = r.text - try: - data = json.loads(data) - if 'error' in data: - error = data['error']['message'] - except Exception: - error = data + if isinstance(e, aiohttp.ClientResponseError): + error = e.message raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error)) diff --git a/backend/open_webui/routers/openai.py b/backend/open_webui/routers/openai.py index 652c910fea..82c844afba 100644 --- a/backend/open_webui/routers/openai.py +++ b/backend/open_webui/routers/openai.py @@ -8,7 +8,7 @@ from urllib.parse import quote, urlparse import aiohttp from aiocache import cached -import requests + from azure.identity import DefaultAzureCredential, get_bearer_token_provider @@ -312,19 +312,20 @@ async def speech(request: Request, user=Depends(get_verified_user)): r = None try: - r = requests.post( + session = await get_session() + r = await session.post( url=f'{url}/audio/speech', data=body, headers=headers, cookies=cookies, - stream=True, + ssl=AIOHTTP_CLIENT_SESSION_SSL, ) r.raise_for_status() # Save the streaming content to a file with open(file_path, 'wb') as f: - for chunk in r.iter_content(chunk_size=8192): + async for chunk in r.content.iter_chunked(8192): f.write(chunk) with open(file_body_path, 'w') as f: @@ -339,14 +340,14 @@ async def speech(request: Request, user=Depends(get_verified_user)): detail = None if r is not None: try: - res = r.json() + res = await r.json() if 'error' in res: detail = f'External: {res["error"]}' except Exception: detail = f'External: {e}' raise HTTPException( - status_code=r.status_code if r else 500, + status_code=r.status if r else 500, detail=detail if detail else 'Open WebUI: Server Connection Error', ) diff --git a/backend/open_webui/utils/files.py b/backend/open_webui/utils/files.py index 21e3af5752..9eec22a5c3 100644 --- a/backend/open_webui/utils/files.py +++ b/backend/open_webui/utils/files.py @@ -25,7 +25,8 @@ import base64 import io import re -import requests +from open_webui.env import AIOHTTP_CLIENT_SESSION_SSL +from open_webui.utils.session_pool import get_session BASE64_IMAGE_URL_PREFIX = re.compile(r'data:image/\w+;base64,', re.IGNORECASE) MARKDOWN_IMAGE_URL_PATTERN = re.compile(r'!\[(.*?)\]\((.+?)\)', re.IGNORECASE) @@ -37,12 +38,13 @@ async def get_image_base64_from_url(url: str) -> Optional[str]: # Validate URL to prevent SSRF attacks against local/private networks validate_url(url) # Download the image from the URL - response = requests.get(url) - response.raise_for_status() - image_data = response.content - encoded_string = base64.b64encode(image_data).decode('utf-8') - content_type = response.headers.get('Content-Type', 'image/png') - return f'data:{content_type};base64,{encoded_string}' + session = await get_session() + async with session.get(url, ssl=AIOHTTP_CLIENT_SESSION_SSL) as response: + response.raise_for_status() + image_data = await response.read() + encoded_string = base64.b64encode(image_data).decode('utf-8') + content_type = response.headers.get('Content-Type', 'image/png') + return f'data:{content_type};base64,{encoded_string}' else: file = await Files.get_file_by_id(url)