ewilliams28/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-08-13 01:02:25 -06:00
Code Issues 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
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,
backend/open_webui/routers/images.py
+81 -69
View File
@@ -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))
backend/open_webui/routers/openai.py
+7 -6
View File
@@ -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',
)
backend/open_webui/utils/files.py
+9 -7
View File
@@ -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)