ewilliams28/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-08-12 16:52:25 -06:00
Code Issues Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-07-27 04:17:00 -04:00
parent 8ab44ed3b1
commit 4493b56e42
3 changed files with 29 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/routers/ollama.py
+9 -11
View File
@@ -24,7 +24,6 @@ from open_webui.constants import ERROR_MESSAGES
from open_webui.events import EVENTS, publish_event, publish_model_provider_request_failed
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,
@@ -58,6 +57,7 @@ log = logging.getLogger(__name__)
# clients to attempt decompression of an already-decoded payload, resulting
# in ZlibError. See https://github.com/aio-libs/aiohttp/issues/4462.
_STRIP_PROXY_HEADERS = frozenset({'Content-Encoding', 'Content-Length', 'Transfer-Encoding'})
_MODEL_LIST_TIMEOUT = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST)
def _clean_proxy_headers(raw_headers) -> dict:
@@ -85,7 +85,7 @@ async def send_get_request(
url,
headers=headers,
ssl=AIOHTTP_CLIENT_SESSION_SSL,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST),
timeout=_MODEL_LIST_TIMEOUT,
) as r:
return await r.json(loads=JSONCodec.loads)
except Exception as exc:
@@ -267,7 +267,7 @@ async def verify_connection(
f'{form_data.url}/api/version',
headers=headers,
ssl=AIOHTTP_CLIENT_SESSION_SSL,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST),
timeout=_MODEL_LIST_TIMEOUT,
) as r:
if r.status != 200:
detail = f'HTTP Error: {r.status}'
@@ -1117,7 +1117,7 @@ async def generate_chat_completion(
raise HTTPException(status_code=400, detail=str(exc))
if isinstance(form_data, BaseModel):
payload = {**form_data.model_dump(exclude_none=True)}
payload = form_data.model_dump(exclude_none=True)
payload.pop('metadata', None)
@@ -1221,8 +1221,7 @@ async def generate_openai_completion(
log.exception(exc)
raise HTTPException(status_code=400, detail=str(exc))
payload = {**form_data.model_dump(exclude_none=True, exclude=['metadata'])}
payload.pop('metadata', None)
payload = form_data.model_dump(exclude_none=True, exclude=['metadata'])
model_id = form_data.model
model_info = await Models.get_model_by_id(model_id)
@@ -1276,7 +1275,7 @@ async def generate_openai_embeddings(
log.exception(exc)
raise HTTPException(status_code=400, detail=str(exc))
payload = {**form_data.model_dump(exclude_none=True)}
payload = form_data.model_dump(exclude_none=True)
payload.pop('metadata', None)
model_id = form_data.model
@@ -1326,8 +1325,7 @@ async def generate_openai_chat_completion(
log.exception(exc)
raise HTTPException(status_code=400, detail=str(exc))
payload = {**form_data.model_dump(exclude_none=True, exclude=['metadata'])}
payload.pop('metadata', None)
payload = form_data.model_dump(exclude_none=True, exclude=['metadata'])
model_id = form_data.model
model_info = await Models.get_model_by_id(model_id)
@@ -1669,7 +1667,7 @@ async def upload_model(
data=blob_chunks(),
headers={'Content-Length': str(total_size)},
ssl=AIOHTTP_CLIENT_SESSION_SSL,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
timeout=get_client_timeout(),
) as resp:
if not resp.ok:
raise Exception('Ollama: Could not create blob, Please try again.')
@@ -1692,7 +1690,7 @@ async def upload_model(
headers={'Content-Type': 'application/json'},
data=json.dumps(create_payload),
ssl=AIOHTTP_CLIENT_SESSION_SSL,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
timeout=get_client_timeout(),
) as create_resp:
if create_resp.ok:
log.info('API SUCCESS!')
backend/open_webui/routers/openai.py
+13 -34
View File
@@ -26,7 +26,6 @@ from open_webui.constants import ERROR_MESSAGES
from open_webui.events import EVENTS, publish_event, publish_model_provider_request_failed
from open_webui.env import (
AIOHTTP_CLIENT_SESSION_SSL,
AIOHTTP_CLIENT_TIMEOUT,
AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST,
BYPASS_MODEL_ACCESS_CONTROL,
ENABLE_FORWARD_USER_INFO_HEADERS,
@@ -79,6 +78,8 @@ log = logging.getLogger(__name__)
# clients to attempt decompression of an already-decoded payload, resulting
# in ZlibError. See https://github.com/aio-libs/aiohttp/issues/4462.
_STRIP_PROXY_HEADERS = frozenset({'Content-Encoding', 'Content-Length', 'Transfer-Encoding'})
_MODEL_LIST_TIMEOUT = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST)
_UNSUPPORTED_OPENAI_MODEL_KEYWORDS = ('babbage', 'dall-e', 'davinci', 'embedding', 'tts', 'whisper')
def _clean_proxy_headers(raw_headers) -> dict:
@@ -93,9 +94,8 @@ async def send_get_request(
user: UserModel = None,
config=None,
):
timeout = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST)
try:
async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session:
async with aiohttp.ClientSession(timeout=_MODEL_LIST_TIMEOUT, trust_env=True) as session:
if request and config:
headers, cookies = await get_headers_and_cookies(request, url, key, config, user=user)
else:
@@ -354,7 +354,7 @@ async def count_anthropic_tokens(request: Request, form_data: dict, user: UserMo
headers=headers,
cookies=cookies,
ssl=AIOHTTP_CLIENT_SESSION_SSL,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
timeout=get_client_timeout(),
)
try:
@@ -661,19 +661,7 @@ async def get_all_models(request: Request, user: UserModel) -> dict[str, list]:
return None
def is_supported_openai_models(model_id):
if any(
name in model_id
for name in [
'babbage',
'dall-e',
'davinci',
'embedding',
'tts',
'whisper',
]
):
return False
return True
return not any(name in model_id for name in _UNSUPPORTED_OPENAI_MODEL_KEYWORDS)
def get_merged_models(model_lists):
log.debug(f'merge_models_lists {model_lists}')
@@ -681,17 +669,18 @@ async def get_all_models(request: Request, user: UserModel) -> dict[str, list]:
for idx, model_list in enumerate(model_lists):
if model_list is not None and 'error' not in model_list:
base_url = api_base_urls[idx]
hostname = urlparse(base_url).hostname if base_url else None
api_config = api_configs.get(str(idx), api_configs.get(base_url, {}))
for model in model_list:
model_id = model.get('id') or model.get('name')
base_url = api_base_urls[idx]
hostname = urlparse(base_url).hostname if base_url else None
if hostname == 'api.openai.com' and not is_supported_openai_models(model_id):
# Skip unwanted OpenAI models
continue
if model_id and model_id not in models:
api_config = api_configs.get(str(idx), api_configs.get(base_url, {}))
provider = model.get('provider', '')
merged = {
**model,
@@ -740,7 +729,7 @@ async def get_models(request: Request, url_idx: int | None = None, user=Depends(
r = None
async with aiohttp.ClientSession(
trust_env=True,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST),
timeout=_MODEL_LIST_TIMEOUT,
) as session:
try:
headers, cookies = await get_headers_and_cookies(request, url, key, api_config, user=user)
@@ -777,17 +766,7 @@ async def get_models(request: Request, url_idx: int | None = None, user=Depends(
response_data['data'] = [
model
for model in response_data.get('data', [])
if not any(
name in model['id']
for name in [
'babbage',
'dall-e',
'davinci',
'embedding',
'tts',
'whisper',
]
)
if not any(name in model['id'] for name in _UNSUPPORTED_OPENAI_MODEL_KEYWORDS)
]
models = response_data
@@ -826,7 +805,7 @@ async def verify_connection(
async with aiohttp.ClientSession(
trust_env=True,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST),
timeout=_MODEL_LIST_TIMEOUT,
) as session:
try:
headers, cookies = await get_headers_and_cookies(request, url, key, api_config, user=user)
@@ -1517,7 +1496,7 @@ async def embeddings(request: Request, form_data: dict, user):
data=body,
headers=headers,
cookies=cookies,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
timeout=get_client_timeout(),
ssl=AIOHTTP_CLIENT_SESSION_SSL,
)
backend/open_webui/utils/session_pool.py
+7 -4
View File
@@ -39,12 +39,15 @@ log = logging.getLogger(__name__)
_session: Optional[aiohttp.ClientSession] = None
_CLIENT_TIMEOUT = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)
_CLIENT_STREAM_TIMEOUT = aiohttp.ClientTimeout(
total=AIOHTTP_CLIENT_TIMEOUT,
sock_read=AIOHTTP_CLIENT_STREAM_IDLE_TIMEOUT,
)
def get_client_timeout(stream: bool = False) -> aiohttp.ClientTimeout:
return aiohttp.ClientTimeout(
total=AIOHTTP_CLIENT_TIMEOUT,
sock_read=AIOHTTP_CLIENT_STREAM_IDLE_TIMEOUT if stream else None,
)
return _CLIENT_STREAM_TIMEOUT if stream else _CLIENT_TIMEOUT
async def get_session() -> aiohttp.ClientSession: