From c0fb36c9b833a85a3a7364e195cf54f3d6c7a787 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Sat, 12 Sep 2026 20:23:50 -0400 Subject: [PATCH] refac --- backend/open_webui/utils/middleware.py | 106 +++++++++++++------------ 1 file changed, 56 insertions(+), 50 deletions(-) diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index e497b27196..4523e842ad 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -6435,62 +6435,68 @@ async def streaming_chat_response_handler(response, ctx): def wrap_item(item): return f'data: {item}\n\n' - assistant_message = {} - filter_context = FilterContext() - has_api_outlet_filters = ENABLE_API_OUTLET_FILTERS and bool(filter_functions) - if ENABLE_API_OUTLET_FILTERS and not has_api_outlet_filters: - try: - model_id = model.get('id') if isinstance(model, dict) else model - has_api_outlet_filters = bool( - (isinstance(model, dict) and 'pipeline' in model) - or get_sorted_filters(model_id, request.app.state.MODELS) + try: + assistant_message = {} + filter_context = FilterContext() + has_api_outlet_filters = ENABLE_API_OUTLET_FILTERS and bool(filter_functions) + if ENABLE_API_OUTLET_FILTERS and not has_api_outlet_filters: + try: + model_id = model.get('id') if isinstance(model, dict) else model + has_api_outlet_filters = bool( + (isinstance(model, dict) and 'pipeline' in model) + or get_sorted_filters(model_id, request.app.state.MODELS) + ) + except Exception: + has_api_outlet_filters = True + + for event in events: + event, _ = await process_filter_functions( + request=request, + filter_context=filter_context, + filter_functions=filter_functions, + filter_type='stream', + form_data=event, + extra_params=extra_params, ) - except Exception: - has_api_outlet_filters = True - for event in events: - event, _ = await process_filter_functions( - request=request, - filter_context=filter_context, - filter_functions=filter_functions, - filter_type='stream', - form_data=event, - extra_params=extra_params, - ) + if event: + yield wrap_item(JSONCodec.dumps(event)) - if event: - yield wrap_item(JSONCodec.dumps(event)) + async for data in original_generator: + if filter_functions: + line = data.decode('utf-8', 'replace') if isinstance(data, bytes) else data + if isinstance(line, str) and line.startswith('data:'): + payload = line.removeprefix('data:').strip() + if payload and payload != '[DONE]': + try: + event = JSONCodec.loads(payload) + except JSONCodec.JSONDecodeError: + event = None - async for data in original_generator: - if filter_functions: - line = data.decode('utf-8', 'replace') if isinstance(data, bytes) else data - if isinstance(line, str) and line.startswith('data:'): - payload = line.removeprefix('data:').strip() - if payload and payload != '[DONE]': - try: - event = JSONCodec.loads(payload) - except JSONCodec.JSONDecodeError: - event = None + if isinstance(event, dict): + event, _ = await process_filter_functions( + request=request, + filter_context=filter_context, + filter_functions=filter_functions, + filter_type='stream', + form_data=event, + extra_params=extra_params, + ) + data = wrap_item(JSONCodec.dumps(event)) if event else None - if isinstance(event, dict): - event, _ = await process_filter_functions( - request=request, - filter_context=filter_context, - filter_functions=filter_functions, - filter_type='stream', - form_data=event, - extra_params=extra_params, - ) - data = wrap_item(JSONCodec.dumps(event)) if event else None + if data: + if has_api_outlet_filters: + update_assistant_message_from_stream(assistant_message, data) + yield data - if data: - if has_api_outlet_filters: - update_assistant_message_from_stream(assistant_message, data) - yield data - - if has_api_outlet_filters and assistant_message: - ctx['assistant_message'] = assistant_message - await outlet_filter_handler(ctx) + if has_api_outlet_filters and assistant_message: + ctx['assistant_message'] = assistant_message + await outlet_filter_handler(ctx) + except Exception as e: + log.exception('Chat completion stream failed mid-response: %s', e) + # Separate the error frame from any unfinished upstream event. + yield f'\n\ndata: {JSONCodec.dumps({"error": {"message": "Chat completion stream failed"}})}\n\n' + yield 'data: [DONE]\n\n' return StreamingResponse( stream_wrapper(response.body_iterator, events),