From 4493b56e424db29fa9e72310b1ca6b025c3e5f8b Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Mon, 27 Jul 2026 04:17:00 -0400 Subject: [PATCH] refac --- backend/open_webui/routers/ollama.py | 20 +++++----- backend/open_webui/routers/openai.py | 47 +++++++----------------- backend/open_webui/utils/session_pool.py | 11 ++++-- 3 files changed, 29 insertions(+), 49 deletions(-) diff --git a/backend/open_webui/routers/ollama.py b/backend/open_webui/routers/ollama.py index 08bd947d50..b91aed2ac4 100644 --- a/backend/open_webui/routers/ollama.py +++ b/backend/open_webui/routers/ollama.py @@ -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!') diff --git a/backend/open_webui/routers/openai.py b/backend/open_webui/routers/openai.py index 9794f52265..d5cad77bef 100644 --- a/backend/open_webui/routers/openai.py +++ b/backend/open_webui/routers/openai.py @@ -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, ) diff --git a/backend/open_webui/utils/session_pool.py b/backend/open_webui/utils/session_pool.py index 980435e033..fb66c3e2cf 100644 --- a/backend/open_webui/utils/session_pool.py +++ b/backend/open_webui/utils/session_pool.py @@ -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: