From c47dd7b7717c4186e0f0549ca3c8cb4d9bb38135 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Sun, 12 Apr 2026 17:22:06 -0500 Subject: [PATCH] refac --- backend/open_webui/env.py | 30 ++++++ backend/open_webui/main.py | 4 + backend/open_webui/routers/ollama.py | 13 +-- backend/open_webui/routers/openai.py | 46 ++++----- backend/open_webui/utils/middleware.py | 5 +- backend/open_webui/utils/misc.py | 6 +- backend/open_webui/utils/session_pool.py | 114 +++++++++++++++++++++++ 7 files changed, 182 insertions(+), 36 deletions(-) create mode 100644 backend/open_webui/utils/session_pool.py diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py index 180926801..57b2bd564 100644 --- a/backend/open_webui/env.py +++ b/backend/open_webui/env.py @@ -806,6 +806,36 @@ else: AIOHTTP_CLIENT_TIMEOUT_TOOL_SERVER = AIOHTTP_CLIENT_TIMEOUT +#################################### +# AIOHTTP Connection Pool +#################################### + +AIOHTTP_POOL_CONNECTIONS = os.environ.get('AIOHTTP_POOL_CONNECTIONS', '') +if AIOHTTP_POOL_CONNECTIONS == '': + AIOHTTP_POOL_CONNECTIONS = None +else: + try: + AIOHTTP_POOL_CONNECTIONS = int(AIOHTTP_POOL_CONNECTIONS) + except ValueError: + AIOHTTP_POOL_CONNECTIONS = None + +AIOHTTP_POOL_CONNECTIONS_PER_HOST = os.environ.get('AIOHTTP_POOL_CONNECTIONS_PER_HOST', '') +if AIOHTTP_POOL_CONNECTIONS_PER_HOST == '': + AIOHTTP_POOL_CONNECTIONS_PER_HOST = None +else: + try: + AIOHTTP_POOL_CONNECTIONS_PER_HOST = int(AIOHTTP_POOL_CONNECTIONS_PER_HOST) + except ValueError: + AIOHTTP_POOL_CONNECTIONS_PER_HOST = None + +AIOHTTP_POOL_DNS_TTL = os.environ.get('AIOHTTP_POOL_DNS_TTL', '300') +try: + AIOHTTP_POOL_DNS_TTL = int(AIOHTTP_POOL_DNS_TTL) + if AIOHTTP_POOL_DNS_TTL < 0: + AIOHTTP_POOL_DNS_TTL = 300 +except ValueError: + AIOHTTP_POOL_DNS_TTL = 300 + RAG_EMBEDDING_TIMEOUT = os.environ.get('RAG_EMBEDDING_TIMEOUT', '') if RAG_EMBEDDING_TIMEOUT == '': diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 620a8aa67..b0d0b3eab 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -719,6 +719,10 @@ async def lifespan(app: FastAPI): yield + # Shutdown: clean up shared resources + from open_webui.utils.session_pool import close_session + await close_session() + if hasattr(app.state, 'redis_task_command_listener'): app.state.redis_task_command_listener.cancel() diff --git a/backend/open_webui/routers/ollama.py b/backend/open_webui/routers/ollama.py index 11c916846..93b34f5e6 100644 --- a/backend/open_webui/routers/ollama.py +++ b/backend/open_webui/routers/ollama.py @@ -50,7 +50,10 @@ from open_webui.models.groups import Groups from open_webui.utils.access_control import check_model_access from open_webui.utils.misc import ( calculate_sha256, +) +from open_webui.utils.session_pool import ( cleanup_response, + get_session, stream_wrapper, ) from open_webui.utils.payload import ( @@ -122,10 +125,7 @@ async def send_request( r = None streaming = False try: - session = aiohttp.ClientSession( - trust_env=True, - timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), - ) + session = await get_session() headers = { 'Content-Type': 'application/json', @@ -140,6 +140,7 @@ async def send_request( r = await session.request( method, url, data=payload, headers=headers, ssl=AIOHTTP_CLIENT_SESSION_SSL, + timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), ) if not r.ok: @@ -165,7 +166,7 @@ async def send_request( streaming = True return StreamingResponse( - stream_wrapper(r, session), + stream_wrapper(r), status_code=r.status, headers=response_headers, ) @@ -184,7 +185,7 @@ async def send_request( ) finally: if not streaming: - await cleanup_response(r, session) + await cleanup_response(r) def get_api_key(idx, url, configs): diff --git a/backend/open_webui/routers/openai.py b/backend/open_webui/routers/openai.py index 51d4267c3..0bfddf47d 100644 --- a/backend/open_webui/routers/openai.py +++ b/backend/open_webui/routers/openai.py @@ -52,9 +52,12 @@ from open_webui.utils.payload import ( apply_system_prompt_to_body, ) from open_webui.utils.misc import ( - cleanup_response, convert_logit_bias_input_to_json, stream_chunks_handler, +) +from open_webui.utils.session_pool import ( + cleanup_response, + get_session, stream_wrapper, ) @@ -1174,12 +1177,11 @@ async def generate_chat_completion( payload = json.dumps(payload) r = None - session = None streaming = False response = None try: - session = aiohttp.ClientSession(trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)) + session = await get_session() r = await session.request( method='POST', @@ -1188,13 +1190,14 @@ async def generate_chat_completion( headers=headers, cookies=cookies, ssl=AIOHTTP_CLIENT_SESSION_SSL, + timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), ) # Check if response is SSE if 'text/event-stream' in r.headers.get('Content-Type', ''): streaming = True return StreamingResponse( - stream_wrapper(r, session, stream_chunks_handler), + stream_wrapper(r, content_handler=stream_chunks_handler), status_code=r.status, headers=dict(r.headers), ) @@ -1225,7 +1228,7 @@ async def generate_chat_completion( ) finally: if not streaming: - await cleanup_response(r, session) + await cleanup_response(r) async def embeddings(request: Request, form_data: dict, user): @@ -1261,27 +1264,24 @@ async def embeddings(request: Request, form_data: dict, user): ) r = None - session = None streaming = False headers, cookies = await get_headers_and_cookies(request, url, key, api_config, user=user) try: - session = aiohttp.ClientSession( - trust_env=True, - timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), - ) + session = await get_session() r = await session.request( method='POST', url=f'{url}/embeddings', data=body, headers=headers, cookies=cookies, + timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), ) if 'text/event-stream' in r.headers.get('Content-Type', ''): streaming = True return StreamingResponse( - stream_wrapper(r, session), + stream_wrapper(r), status_code=r.status, headers=dict(r.headers), ) @@ -1306,7 +1306,7 @@ async def embeddings(request: Request, form_data: dict, user): ) finally: if not streaming: - await cleanup_response(r, session) + await cleanup_response(r) class ResponsesForm(BaseModel): @@ -1365,7 +1365,6 @@ async def responses( ) r = None - session = None streaming = False try: @@ -1388,10 +1387,7 @@ async def responses( else: request_url = f'{url}/responses' - session = aiohttp.ClientSession( - trust_env=True, - timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), - ) + session = await get_session() r = await session.request( method='POST', url=request_url, @@ -1399,13 +1395,14 @@ async def responses( headers=headers, cookies=cookies, ssl=AIOHTTP_CLIENT_SESSION_SSL, + timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), ) # Check if response is SSE if 'text/event-stream' in r.headers.get('Content-Type', ''): streaming = True return StreamingResponse( - stream_wrapper(r, session), + stream_wrapper(r), status_code=r.status, headers=dict(r.headers), ) @@ -1433,7 +1430,7 @@ async def responses( ) finally: if not streaming: - await cleanup_response(r, session) + await cleanup_response(r) @router.api_route('/{path:path}', methods=['GET', 'POST', 'PUT', 'DELETE']) @@ -1479,7 +1476,6 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)): ) r = None - session = None streaming = False try: @@ -1508,10 +1504,7 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)): else: request_url = f'{url}/{path}' - session = aiohttp.ClientSession( - trust_env=True, - timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), - ) + session = await get_session() r = await session.request( method=request.method, url=request_url, @@ -1519,13 +1512,14 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)): headers=headers, cookies=cookies, ssl=AIOHTTP_CLIENT_SESSION_SSL, + timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), ) # Check if response is SSE if 'text/event-stream' in r.headers.get('Content-Type', ''): streaming = True return StreamingResponse( - stream_wrapper(r, session), + stream_wrapper(r), status_code=r.status, headers=dict(r.headers), ) @@ -1553,4 +1547,4 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)): ) finally: if not streaming: - await cleanup_response(r, session) + await cleanup_response(r) diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index fb4912bef..ce75cd15f 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -4060,11 +4060,12 @@ async def streaming_chat_response_handler(response, ctx): if responses_api_tool_calls: tool_calls.append(_split_tool_calls(responses_api_tool_calls)) + try: + await stream_body_handler(response, form_data) + finally: if response.background: await response.background() - await stream_body_handler(response, form_data) - tool_call_retries = 0 tool_call_sources = [] # Track citation sources from tool results all_tool_call_sources = [] # Accumulated sources across all iterations diff --git a/backend/open_webui/utils/misc.py b/backend/open_webui/utils/misc.py index 787ef4d6e..7b52f0afd 100644 --- a/backend/open_webui/utils/misc.py +++ b/backend/open_webui/utils/misc.py @@ -905,9 +905,11 @@ async def cleanup_response( session: Optional[aiohttp.ClientSession], ): if response: - response.close() + if not response.closed: + await response.close() if session: - await session.close() + if not session.closed: + await session.close() async def stream_wrapper(response, session, content_handler=None): diff --git a/backend/open_webui/utils/session_pool.py b/backend/open_webui/utils/session_pool.py new file mode 100644 index 000000000..86ffa6cd9 --- /dev/null +++ b/backend/open_webui/utils/session_pool.py @@ -0,0 +1,114 @@ +"""Shared aiohttp ClientSession pool. + +Instead of creating a new ClientSession (and TCPConnector) per request, +callers acquire a long-lived session from this module. The pool manages +a single TCPConnector with configurable limits, enabling TCP/SSL connection +reuse, shared DNS cache, and bounded concurrency. + +All pool parameters are configurable via environment variables: + - AIOHTTP_POOL_CONNECTIONS (default 100) — max total connections + - AIOHTTP_POOL_CONNECTIONS_PER_HOST (default 30) — per-host limit + - AIOHTTP_POOL_DNS_TTL (default 300) — DNS cache TTL in seconds + +Usage: + from open_webui.utils.session_pool import get_session, cleanup_response + + session = await get_session() + r = await session.request(...) + # When done with the *response* (not the session): + await cleanup_response(r) + +IMPORTANT: Callers must NOT close the shared session. Only the response +needs cleanup. The session is closed once during application shutdown +via ``close_session()``. +""" + +import logging +from typing import Optional + +import aiohttp + +from open_webui.env import ( + AIOHTTP_CLIENT_TIMEOUT, + AIOHTTP_POOL_CONNECTIONS, + AIOHTTP_POOL_CONNECTIONS_PER_HOST, + AIOHTTP_POOL_DNS_TTL, +) + +log = logging.getLogger(__name__) + +_session: Optional[aiohttp.ClientSession] = None + + +async def get_session() -> aiohttp.ClientSession: + """Return the shared aiohttp ClientSession, creating it lazily.""" + global _session + if _session is None or _session.closed: + connector_kwargs = { + 'ttl_dns_cache': AIOHTTP_POOL_DNS_TTL, + 'enable_cleanup_closed': True, + } + if AIOHTTP_POOL_CONNECTIONS is not None: + connector_kwargs['limit'] = AIOHTTP_POOL_CONNECTIONS + else: + connector_kwargs['limit'] = 0 # aiohttp: 0 = unlimited + if AIOHTTP_POOL_CONNECTIONS_PER_HOST is not None: + connector_kwargs['limit_per_host'] = AIOHTTP_POOL_CONNECTIONS_PER_HOST + else: + connector_kwargs['limit_per_host'] = 0 # aiohttp: 0 = unlimited + connector = aiohttp.TCPConnector(**connector_kwargs) + timeout = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT) + _session = aiohttp.ClientSession( + connector=connector, + timeout=timeout, + trust_env=True, + ) + log.info( + 'Created shared aiohttp session pool ' + '(limit=%s, per_host=%s, dns_ttl=%d)', + AIOHTTP_POOL_CONNECTIONS or 'unlimited', + AIOHTTP_POOL_CONNECTIONS_PER_HOST or 'unlimited', + AIOHTTP_POOL_DNS_TTL, + ) + return _session + + +async def close_session(): + """Close the shared session. Called during application shutdown.""" + global _session + if _session and not _session.closed: + await _session.close() + log.info('Closed shared aiohttp session pool') + _session = None + + +async def cleanup_response( + response: Optional[aiohttp.ClientResponse], + session: Optional[aiohttp.ClientSession] = None, +): + """Release and close an aiohttp response, optionally closing the session. + + When using the shared pool, ``session`` should be ``None`` (the pool + session is never closed per-request). When a caller creates its own + one-off session, pass it here to close it after the response. + """ + if response: + if not response.closed: + await response.close() + if session: + if not session.closed: + await session.close() + + +async def stream_wrapper(response, session=None, content_handler=None): + """Wrap a stream to ensure cleanup happens even if streaming is interrupted. + + This is more reliable than BackgroundTask which may not run if the client + disconnects. When using the shared pool, ``session`` should be ``None``. + """ + try: + stream = content_handler(response.content) if content_handler else response.content + async for chunk in stream: + yield chunk + finally: + await cleanup_response(response, session)