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-12 17:22:06 -05:00
parent 4498e6faf2
commit c47dd7b771
7 changed files with 182 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/env.py
+30
View File
@@ -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 == '':
backend/open_webui/main.py
+4
View File
@@ -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()
backend/open_webui/routers/ollama.py
+7 -6
View File
@@ -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):
backend/open_webui/routers/openai.py
+20 -26
View File
@@ -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)
backend/open_webui/utils/middleware.py
+3 -2
View File
@@ -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
backend/open_webui/utils/misc.py
+4 -2
View File
@@ -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):
backend/open_webui/utils/session_pool.py
+114
View File
@@ -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)