refac

Author: tjbck

backend/open_webui/main.py

@@ -21,7 +21,7 @@
 from aiocache import cached
 import aiohttp
 import anyio.to_thread
-import requests
+
 from redis import Redis
 
 
@@ -60,6 +60,7 @@
 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,

backend/open_webui/routers/images.py

@@ -10,6 +10,7 @@
 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.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 load_url_image(data):
             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 @@ def get_image_file_item(base64_string, param_name='image'):
             ),
         )
 
-    r = None
+
     try:
         if request.app.state.config.IMAGE_EDIT_ENGINE == 'openai':
             headers = {
@@ -883,17 +889,30 @@ def get_image_file_item(base64_string, param_name='image'):
             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 @@ def get_image_file_item(base64_string, param_name='image'):
                     ]
                 )
 
-            # 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 @@ def get_image_file_item(base64_string, param_name='image'):
             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))

backend/open_webui/routers/openai.py

@@ -8,7 +8,7 @@
 
 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',
             )
 

backend/open_webui/utils/files.py

@@ -25,7 +25,8 @@
 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)