Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-04-13 14:29:27 -05:00
parent 96265cf042
commit 2ddcb30b9a
4 changed files with 106 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/main.py
+9 -2
View File
@@ -21,7 +21,7 @@ from typing import Optional
from aiocache import cached from aiocache import cached
import aiohttp import aiohttp
import anyio.to_thread import anyio.to_thread
import requests
from redis import Redis from redis import Redis
@@ -60,6 +60,7 @@ from starsessions.stores.redis import RedisStore
from open_webui.utils import logger from open_webui.utils import logger
from open_webui.utils.audit import AuditLevel, AuditLoggingMiddleware from open_webui.utils.audit import AuditLevel, AuditLoggingMiddleware
from open_webui.utils.logger import start_logger from open_webui.utils.logger import start_logger
from open_webui.utils.session_pool import get_session
from open_webui.socket.main import ( from open_webui.socket.main import (
MODELS, MODELS,
app as socket_app, app as socket_app,
@@ -2512,7 +2513,13 @@ async def oauth_backchannel_logout(
@app.get('/manifest.json') @app.get('/manifest.json')
async def get_manifest_json(): async def get_manifest_json():
if app.state.EXTERNAL_PWA_MANIFEST_URL: 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: else:
return { return {
'name': app.state.WEBUI_NAME, 'name': app.state.WEBUI_NAME,
backend/open_webui/routers/images.py
+81 -69
View File
@@ -10,6 +10,7 @@ from pathlib import Path
from typing import Optional from typing import Optional
from urllib.parse import quote from urllib.parse import quote
import aiohttp
import requests import requests
from fastapi import APIRouter, Depends, HTTPException, Request, UploadFile from fastapi import APIRouter, Depends, HTTPException, Request, UploadFile
from fastapi.responses import FileResponse from fastapi.responses import FileResponse
@@ -21,7 +22,8 @@ from open_webui.config import (
) )
from open_webui.constants import ERROR_MESSAGES from open_webui.constants import ERROR_MESSAGES
from open_webui.retrieval.web.utils import validate_url 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.models.chats import Chats
from open_webui.routers.files import upload_file_handler, get_file_content_by_id 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)): async def verify_url(request: Request, user=Depends(get_admin_user)):
if request.app.state.config.IMAGE_GENERATION_ENGINE == 'automatic1111': if request.app.state.config.IMAGE_GENERATION_ENGINE == 'automatic1111':
try: try:
r = requests.get( session = await get_session()
async with session.get(
url=f'{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options', url=f'{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options',
headers={'authorization': get_automatic1111_api_auth(request)}, headers={'authorization': get_automatic1111_api_auth(request)},
) ssl=AIOHTTP_CLIENT_SESSION_SSL,
r.raise_for_status() ) as r:
return True r.raise_for_status()
return True
except Exception: except Exception:
request.app.state.config.ENABLE_IMAGE_GENERATION = False request.app.state.config.ENABLE_IMAGE_GENERATION = False
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL) 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: if request.app.state.config.COMFYUI_API_KEY:
headers = {'Authorization': f'Bearer {request.app.state.config.COMFYUI_API_KEY}'} headers = {'Authorization': f'Bearer {request.app.state.config.COMFYUI_API_KEY}'}
try: try:
r = requests.get( session = await get_session()
async with session.get(
url=f'{request.app.state.config.COMFYUI_BASE_URL}/object_info', url=f'{request.app.state.config.COMFYUI_BASE_URL}/object_info',
headers=headers, headers=headers,
) ssl=AIOHTTP_CLIENT_SESSION_SSL,
r.raise_for_status() ) as r:
return True r.raise_for_status()
return True
except Exception: except Exception:
request.app.state.config.ENABLE_IMAGE_GENERATION = False request.app.state.config.ENABLE_IMAGE_GENERATION = False
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL) 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': elif request.app.state.config.IMAGE_GENERATION_ENGINE == 'comfyui':
# TODO - get models from comfyui # TODO - get models from comfyui
headers = {'Authorization': f'Bearer {request.app.state.config.COMFYUI_API_KEY}'} 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', url=f'{request.app.state.config.COMFYUI_BASE_URL}/object_info',
headers=headers, headers=headers,
) ssl=AIOHTTP_CLIENT_SESSION_SSL,
info = r.json() ) as r:
info = await r.json()
workflow = json.loads(request.app.state.config.COMFYUI_WORKFLOW) workflow = json.loads(request.app.state.config.COMFYUI_WORKFLOW)
model_node_id = None 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' request.app.state.config.IMAGE_GENERATION_ENGINE == 'automatic1111'
or request.app.state.config.IMAGE_GENERATION_ENGINE == '' 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', url=f'{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/sd-models',
headers={'authorization': get_automatic1111_api_auth(request)}, headers={'authorization': get_automatic1111_api_auth(request)},
) ssl=AIOHTTP_CLIENT_SESSION_SSL,
models = r.json() ) as r:
models = await r.json()
return list( return list(
map( map(
lambda model: {'id': model['title'], 'name': model['model_name']}, lambda model: {'id': model['title'], 'name': model['model_name']},
@@ -533,7 +543,7 @@ async def image_generations(
model = get_image_model(request) model = get_image_model(request)
r = None
try: try:
if request.app.state.config.IMAGE_GENERATION_ENGINE == 'openai': if request.app.state.config.IMAGE_GENERATION_ENGINE == 'openai':
headers = { headers = {
@@ -568,16 +578,15 @@ async def image_generations(
), ),
} }
# Use asyncio.to_thread for the requests.post call session = await get_session()
r = await asyncio.to_thread( async with session.post(
requests.post,
url=url, url=url,
json=data, json=data,
headers=headers, headers=headers,
) ssl=AIOHTTP_CLIENT_SESSION_SSL,
) as r:
r.raise_for_status() r.raise_for_status()
res = r.json() res = await r.json()
images = [] images = []
@@ -619,16 +628,15 @@ async def image_generations(
model = f'{model}:generateContent' model = f'{model}:generateContent'
data = {'contents': [{'parts': [{'text': form_data.prompt}]}]} data = {'contents': [{'parts': [{'text': form_data.prompt}]}]}
# Use asyncio.to_thread for the requests.post call session = await get_session()
r = await asyncio.to_thread( async with session.post(
requests.post,
url=f'{request.app.state.config.IMAGES_GEMINI_API_BASE_URL}/models/{model}', url=f'{request.app.state.config.IMAGES_GEMINI_API_BASE_URL}/models/{model}',
json=data, json=data,
headers=headers, headers=headers,
) ssl=AIOHTTP_CLIENT_SESSION_SSL,
) as r:
r.raise_for_status() r.raise_for_status()
res = r.json() res = await r.json()
images = [] images = []
@@ -727,15 +735,14 @@ async def image_generations(
if request.app.state.config.AUTOMATIC1111_PARAMS: if request.app.state.config.AUTOMATIC1111_PARAMS:
data = {**data, **request.app.state.config.AUTOMATIC1111_PARAMS} data = {**data, **request.app.state.config.AUTOMATIC1111_PARAMS}
# Use asyncio.to_thread for the requests.post call session = await get_session()
r = await asyncio.to_thread( async with session.post(
requests.post,
url=f'{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/txt2img', url=f'{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/txt2img',
json=data, json=data,
headers={'authorization': get_automatic1111_api_auth(request)}, headers={'authorization': get_automatic1111_api_auth(request)},
) ssl=AIOHTTP_CLIENT_SESSION_SSL,
) as r:
res = r.json() res = await r.json()
log.debug(f'res: {res}') log.debug(f'res: {res}')
images = [] images = []
@@ -753,10 +760,8 @@ async def image_generations(
return images return images
except Exception as e: except Exception as e:
error = e error = e
if r != None: if isinstance(e, aiohttp.ClientResponseError):
data = r.json() error = e.message
if 'error' in data:
error = data['error']['message']
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error)) 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://'): if data.startswith('http://') or data.startswith('https://'):
# Validate URL to prevent SSRF attacks against local/private networks # Validate URL to prevent SSRF attacks against local/private networks
validate_url(data) validate_url(data)
r = await asyncio.to_thread(requests.get, data) session = await get_session()
r.raise_for_status() 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') image_data = base64.b64encode(await r.read()).decode('utf-8')
return f'data:{r.headers["content-type"]};base64,{image_data}' return f'data:{r.headers["content-type"]};base64,{image_data}'
else: else:
file_id = None file_id = None
@@ -846,7 +852,7 @@ async def image_edits(
), ),
) )
r = None
try: try:
if request.app.state.config.IMAGE_EDIT_ENGINE == 'openai': if request.app.state.config.IMAGE_EDIT_ENGINE == 'openai':
headers = { headers = {
@@ -883,17 +889,30 @@ async def image_edits(
if request.app.state.config.IMAGES_EDIT_OPENAI_API_VERSION: 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}' url_search_params += f'?api-version={request.app.state.config.IMAGES_EDIT_OPENAI_API_VERSION}'
# Use asyncio.to_thread for the requests.post call # Build multipart form data for aiohttp
r = await asyncio.to_thread( form = aiohttp.FormData()
requests.post, 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}', url=f'{request.app.state.config.IMAGES_EDIT_OPENAI_API_BASE_URL}/images/edits{url_search_params}',
headers=headers, headers=headers,
files=files, data=form,
data=data, ssl=AIOHTTP_CLIENT_SESSION_SSL,
) ) as r:
r.raise_for_status()
r.raise_for_status() res = await r.json()
res = r.json()
images = [] images = []
for image in res['data']: for image in res['data']:
@@ -940,16 +959,15 @@ async def image_edits(
] ]
) )
# Use asyncio.to_thread for the requests.post call session = await get_session()
r = await asyncio.to_thread( async with session.post(
requests.post,
url=f'{request.app.state.config.IMAGES_EDIT_GEMINI_API_BASE_URL}/models/{model}', url=f'{request.app.state.config.IMAGES_EDIT_GEMINI_API_BASE_URL}/models/{model}',
json=data, json=data,
headers=headers, headers=headers,
) ssl=AIOHTTP_CLIENT_SESSION_SSL,
) as r:
r.raise_for_status() r.raise_for_status()
res = r.json() res = await r.json()
images = [] images = []
for image in res['candidates']: for image in res['candidates']:
@@ -1048,13 +1066,7 @@ async def image_edits(
return images return images
except Exception as e: except Exception as e:
error = e error = e
if r != None: if isinstance(e, aiohttp.ClientResponseError):
data = r.text error = e.message
try:
data = json.loads(data)
if 'error' in data:
error = data['error']['message']
except Exception:
error = data
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error)) raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error))
backend/open_webui/routers/openai.py
+7 -6
View File
@@ -8,7 +8,7 @@ from urllib.parse import quote, urlparse
import aiohttp import aiohttp
from aiocache import cached from aiocache import cached
import requests
from azure.identity import DefaultAzureCredential, get_bearer_token_provider 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 r = None
try: try:
r = requests.post( session = await get_session()
r = await session.post(
url=f'{url}/audio/speech', url=f'{url}/audio/speech',
data=body, data=body,
headers=headers, headers=headers,
cookies=cookies, cookies=cookies,
stream=True, ssl=AIOHTTP_CLIENT_SESSION_SSL,
) )
r.raise_for_status() r.raise_for_status()
# Save the streaming content to a file # Save the streaming content to a file
with open(file_path, 'wb') as f: 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) f.write(chunk)
with open(file_body_path, 'w') as f: with open(file_body_path, 'w') as f:
@@ -339,14 +340,14 @@ async def speech(request: Request, user=Depends(get_verified_user)):
detail = None detail = None
if r is not None: if r is not None:
try: try:
res = r.json() res = await r.json()
if 'error' in res: if 'error' in res:
detail = f'External: {res["error"]}' detail = f'External: {res["error"]}'
except Exception: except Exception:
detail = f'External: {e}' detail = f'External: {e}'
raise HTTPException( 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', detail=detail if detail else 'Open WebUI: Server Connection Error',
) )
backend/open_webui/utils/files.py
+9 -7
View File
@@ -25,7 +25,8 @@ import base64
import io import io
import re 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) BASE64_IMAGE_URL_PREFIX = re.compile(r'data:image/\w+;base64,', re.IGNORECASE)
MARKDOWN_IMAGE_URL_PATTERN = re.compile(r'!\[(.*?)\]\((.+?)\)', 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 to prevent SSRF attacks against local/private networks
validate_url(url) validate_url(url)
# Download the image from the URL # Download the image from the URL
response = requests.get(url) session = await get_session()
response.raise_for_status() async with session.get(url, ssl=AIOHTTP_CLIENT_SESSION_SSL) as response:
image_data = response.content response.raise_for_status()
encoded_string = base64.b64encode(image_data).decode('utf-8') image_data = await response.read()
content_type = response.headers.get('Content-Type', 'image/png') encoded_string = base64.b64encode(image_data).decode('utf-8')
return f'data:{content_type};base64,{encoded_string}' content_type = response.headers.get('Content-Type', 'image/png')
return f'data:{content_type};base64,{encoded_string}'
else: else:
file = await Files.get_file_by_id(url) file = await Files.get_file_by_id(url)