bhetherman/open-webui
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 committed 2026-07-27 04:17:00 -04:00
1 parent 8ab44ed3b1
commit 4493b56e42
3 files changed
+29 -49
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

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.events import EVENTS, publish_event, publish_model_provider_request_failed
from open_webui.env import ( from open_webui.env import (
AIOHTTP_CLIENT_SESSION_SSL, AIOHTTP_CLIENT_SESSION_SSL,
AIOHTTP_CLIENT_TIMEOUT,
AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST, AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST,
AIOHTTP_FILE_STREAM_CHUNK_SIZE, AIOHTTP_FILE_STREAM_CHUNK_SIZE,
BYPASS_MODEL_ACCESS_CONTROL, BYPASS_MODEL_ACCESS_CONTROL,
@@ -58,6 +57,7 @@ log = logging.getLogger(__name__)
# clients to attempt decompression of an already-decoded payload, resulting # clients to attempt decompression of an already-decoded payload, resulting
# in ZlibError. See https://github.com/aio-libs/aiohttp/issues/4462. # in ZlibError. See https://github.com/aio-libs/aiohttp/issues/4462.
_STRIP_PROXY_HEADERS = frozenset({'Content-Encoding', 'Content-Length', 'Transfer-Encoding'}) _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: def _clean_proxy_headers(raw_headers) -> dict:
@@ -85,7 +85,7 @@ async def send_get_request(
url, url,
headers=headers, headers=headers,
ssl=AIOHTTP_CLIENT_SESSION_SSL, ssl=AIOHTTP_CLIENT_SESSION_SSL,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST), timeout=_MODEL_LIST_TIMEOUT,
) as r: ) as r:
return await r.json(loads=JSONCodec.loads) return await r.json(loads=JSONCodec.loads)
except Exception as exc: except Exception as exc:
@@ -267,7 +267,7 @@ async def verify_connection(
f'{form_data.url}/api/version', f'{form_data.url}/api/version',
headers=headers, headers=headers,
ssl=AIOHTTP_CLIENT_SESSION_SSL, ssl=AIOHTTP_CLIENT_SESSION_SSL,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST), timeout=_MODEL_LIST_TIMEOUT,
) as r: ) as r:
if r.status != 200: if r.status != 200:
detail = f'HTTP Error: {r.status}' detail = f'HTTP Error: {r.status}'
@@ -1117,7 +1117,7 @@ async def generate_chat_completion(
raise HTTPException(status_code=400, detail=str(exc)) raise HTTPException(status_code=400, detail=str(exc))
if isinstance(form_data, BaseModel): 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) payload.pop('metadata', None)
@@ -1221,8 +1221,7 @@ async def generate_openai_completion(
log.exception(exc) log.exception(exc)
raise HTTPException(status_code=400, detail=str(exc)) raise HTTPException(status_code=400, detail=str(exc))
payload = {**form_data.model_dump(exclude_none=True, exclude=['metadata'])} payload = form_data.model_dump(exclude_none=True, exclude=['metadata'])
payload.pop('metadata', None)
model_id = form_data.model model_id = form_data.model
model_info = await Models.get_model_by_id(model_id) model_info = await Models.get_model_by_id(model_id)
@@ -1276,7 +1275,7 @@ async def generate_openai_embeddings(
log.exception(exc) log.exception(exc)
raise HTTPException(status_code=400, detail=str(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) payload.pop('metadata', None)
model_id = form_data.model model_id = form_data.model
@@ -1326,8 +1325,7 @@ async def generate_openai_chat_completion(
log.exception(exc) log.exception(exc)
raise HTTPException(status_code=400, detail=str(exc)) raise HTTPException(status_code=400, detail=str(exc))
payload = {**form_data.model_dump(exclude_none=True, exclude=['metadata'])} payload = form_data.model_dump(exclude_none=True, exclude=['metadata'])
payload.pop('metadata', None)
model_id = form_data.model model_id = form_data.model
model_info = await Models.get_model_by_id(model_id) model_info = await Models.get_model_by_id(model_id)
@@ -1669,7 +1667,7 @@ async def upload_model(
data=blob_chunks(), data=blob_chunks(),
headers={'Content-Length': str(total_size)}, headers={'Content-Length': str(total_size)},
ssl=AIOHTTP_CLIENT_SESSION_SSL, ssl=AIOHTTP_CLIENT_SESSION_SSL,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), timeout=get_client_timeout(),
) as resp: ) as resp:
if not resp.ok: if not resp.ok:
raise Exception('Ollama: Could not create blob, Please try again.') raise Exception('Ollama: Could not create blob, Please try again.')
@@ -1692,7 +1690,7 @@ async def upload_model(
headers={'Content-Type': 'application/json'}, headers={'Content-Type': 'application/json'},
data=json.dumps(create_payload), data=json.dumps(create_payload),
ssl=AIOHTTP_CLIENT_SESSION_SSL, ssl=AIOHTTP_CLIENT_SESSION_SSL,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), timeout=get_client_timeout(),
) as create_resp: ) as create_resp:
if create_resp.ok: if create_resp.ok:
log.info('API SUCCESS!') 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.events import EVENTS, publish_event, publish_model_provider_request_failed
from open_webui.env import ( from open_webui.env import (
AIOHTTP_CLIENT_SESSION_SSL, AIOHTTP_CLIENT_SESSION_SSL,
AIOHTTP_CLIENT_TIMEOUT,
AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST, AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST,
BYPASS_MODEL_ACCESS_CONTROL, BYPASS_MODEL_ACCESS_CONTROL,
ENABLE_FORWARD_USER_INFO_HEADERS, ENABLE_FORWARD_USER_INFO_HEADERS,
@@ -79,6 +78,8 @@ log = logging.getLogger(__name__)
# clients to attempt decompression of an already-decoded payload, resulting # clients to attempt decompression of an already-decoded payload, resulting
# in ZlibError. See https://github.com/aio-libs/aiohttp/issues/4462. # in ZlibError. See https://github.com/aio-libs/aiohttp/issues/4462.
_STRIP_PROXY_HEADERS = frozenset({'Content-Encoding', 'Content-Length', 'Transfer-Encoding'}) _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: def _clean_proxy_headers(raw_headers) -> dict:
@@ -93,9 +94,8 @@ async def send_get_request(
user: UserModel = None, user: UserModel = None,
config=None, config=None,
): ):
timeout = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST)
try: 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: if request and config:
headers, cookies = await get_headers_and_cookies(request, url, key, config, user=user) headers, cookies = await get_headers_and_cookies(request, url, key, config, user=user)
else: else:
@@ -354,7 +354,7 @@ async def count_anthropic_tokens(request: Request, form_data: dict, user: UserMo
headers=headers, headers=headers,
cookies=cookies, cookies=cookies,
ssl=AIOHTTP_CLIENT_SESSION_SSL, ssl=AIOHTTP_CLIENT_SESSION_SSL,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), timeout=get_client_timeout(),
) )
try: try:
@@ -661,19 +661,7 @@ async def get_all_models(request: Request, user: UserModel) -> dict[str, list]:
return None return None
def is_supported_openai_models(model_id): def is_supported_openai_models(model_id):
if any( return not any(name in model_id for name in _UNSUPPORTED_OPENAI_MODEL_KEYWORDS)
name in model_id
for name in [
'babbage',
'dall-e',
'davinci',
'embedding',
'tts',
'whisper',
]
):
return False
return True
def get_merged_models(model_lists): def get_merged_models(model_lists):
log.debug(f'merge_models_lists {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): for idx, model_list in enumerate(model_lists):
if model_list is not None and 'error' not in model_list: 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: for model in model_list:
model_id = model.get('id') or model.get('name') 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): if hostname == 'api.openai.com' and not is_supported_openai_models(model_id):
# Skip unwanted OpenAI models # Skip unwanted OpenAI models
continue continue
if model_id and model_id not in models: 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', '') provider = model.get('provider', '')
merged = { merged = {
**model, **model,
@@ -740,7 +729,7 @@ async def get_models(request: Request, url_idx: int | None = None, user=Depends(
r = None r = None
async with aiohttp.ClientSession( async with aiohttp.ClientSession(
trust_env=True, trust_env=True,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST), timeout=_MODEL_LIST_TIMEOUT,
) as session: ) as session:
try: try:
headers, cookies = await get_headers_and_cookies(request, url, key, api_config, user=user) 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'] = [ response_data['data'] = [
model model
for model in response_data.get('data', []) for model in response_data.get('data', [])
if not any( if not any(name in model['id'] for name in _UNSUPPORTED_OPENAI_MODEL_KEYWORDS)
name in model['id']
for name in [
'babbage',
'dall-e',
'davinci',
'embedding',
'tts',
'whisper',
]
)
] ]
models = response_data models = response_data
@@ -826,7 +805,7 @@ async def verify_connection(
async with aiohttp.ClientSession( async with aiohttp.ClientSession(
trust_env=True, trust_env=True,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST), timeout=_MODEL_LIST_TIMEOUT,
) as session: ) as session:
try: try:
headers, cookies = await get_headers_and_cookies(request, url, key, api_config, user=user) 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, data=body,
headers=headers, headers=headers,
cookies=cookies, cookies=cookies,
timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), timeout=get_client_timeout(),
ssl=AIOHTTP_CLIENT_SESSION_SSL, 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 _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: def get_client_timeout(stream: bool = False) -> aiohttp.ClientTimeout:
return aiohttp.ClientTimeout( return _CLIENT_STREAM_TIMEOUT if stream else _CLIENT_TIMEOUT
total=AIOHTTP_CLIENT_TIMEOUT,
sock_read=AIOHTTP_CLIENT_STREAM_IDLE_TIMEOUT if stream else None,
)
async def get_session() -> aiohttp.ClientSession: async def get_session() -> aiohttp.ClientSession: