From 7eefeef4f17118f81c87acb464ad562803a6f26c Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Sun, 6 Sep 2026 19:32:31 -0400 Subject: [PATCH] refac --- backend/open_webui/main.py | 8 +++- backend/open_webui/utils/middleware.py | 63 ++++++++++++++++++++++---- 2 files changed, 61 insertions(+), 10 deletions(-) diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index ec7827be10..ac7b35c2ac 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -1624,6 +1624,9 @@ async def chat_completion( async def process_chat(request, form_data, user, metadata, model, tasks=None): try: + ctx = None + if not form_data.get('stream') and metadata.get('assistant_message_id'): + ctx = await build_chat_response_context(request, form_data, user, model, metadata, tasks, []) form_data, metadata, events = await process_chat_payload(request, form_data, user, metadata, model) if await drain_approved_tool_calls(request, form_data, user, model, metadata): @@ -1639,7 +1642,10 @@ async def chat_completion( if isinstance(response, JSONResponse) and response.status_code >= 400: raise Exception(get_response_error_detail(response)) - ctx = await build_chat_response_context(request, form_data, user, model, metadata, tasks, events) + if ctx is None: + ctx = await build_chat_response_context(request, form_data, user, model, metadata, tasks, events) + else: + ctx.update(form_data=form_data, metadata=metadata, events=events) return await process_chat_response(response, ctx) except asyncio.CancelledError: diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index eb4d1e742e..a076c2c4e9 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -3106,6 +3106,15 @@ async def get_event_emitter_and_caller(metadata): async def build_chat_response_context(request, form_data, user, model, metadata, tasks, events): event_emitter, event_caller = await get_event_emitter_and_caller(metadata) + assistant_message = None + if not form_data.get('stream') and metadata.get('assistant_message_id'): + # Preserve the original output before request conversion strips it from temporary chats. + if is_saved_chat_id(metadata.get('chat_id')): + assistant_message = await Chats.get_message_by_id_and_message_id( + metadata['chat_id'], metadata['assistant_message_id'] + ) + elif form_data.get('messages') and form_data['messages'][-1].get('role') == 'assistant': + assistant_message = form_data['messages'][-1] return { 'request': request, 'form_data': form_data, @@ -3116,6 +3125,7 @@ async def build_chat_response_context(request, form_data, user, model, metadata, 'events': events, 'event_emitter': event_emitter, 'event_caller': event_caller, + 'assistant_message': copy.deepcopy(assistant_message), } @@ -4014,6 +4024,7 @@ async def non_streaming_chat_response_handler(response, ctx): chat_id = metadata.get('chat_id') or '' save_to_chat = is_saved_chat_id(chat_id) + continuing = bool(metadata.get('assistant_message_id')) if event_emitter: try: @@ -4057,21 +4068,24 @@ async def non_streaming_chat_response_handler(response, ctx): response_output = response_data.get('output') content = choices[0].get('message', {}).get('content') if choices else '' - if choices and (content or response_output): - if content or response_output: - await event_emitter( - { - 'type': 'chat:completion', - 'data': response_data, - } - ) + if (continuing and 'error' not in response_data) or ( + not continuing and choices and (content or response_output) + ): + if content or response_output or continuing: + if not continuing: + await event_emitter( + { + 'type': 'chat:completion', + 'data': response_data, + } + ) title = await Chats.get_chat_title_by_id(metadata['chat_id']) if save_to_chat else '' # Use output from backend if provided (OR-compliant backends), # otherwise generate from response content if not response_output: - choice_message = choices[0].get('message', {}) + choice_message = choices[0].get('message', {}) if choices else {} reasoning_content = choice_message.get('reasoning_content') or choice_message.get('reasoning') reasoning_details = get_reasoning_details(choice_message) response_output = [] @@ -4103,10 +4117,41 @@ async def non_streaming_chat_response_handler(response, ctx): } ) + if continuing: + message = ctx.get('assistant_message') or {} + previous = message.get('output') or [] + if not previous and message.get('content'): + previous = [ + { + 'type': 'message', + 'id': message.get('id') or output_id('msg'), + 'role': 'assistant', + 'content': [{'type': 'output_text', 'text': message['content']}], + } + ] + response_output = list(response_output) + if ( + previous + and response_output + and previous[-1].get('type') == response_output[0].get('type') == 'message' + and previous[-1].get('role', 'assistant') == 'assistant' + and response_output[0].get('role', 'assistant') == 'assistant' + ): + response_output[0] = { + **previous[-1], + **response_output[0], + 'id': previous[-1].get('id'), + 'content': [*previous[-1].get('content', []), *response_output[0].get('content', [])], + } + previous = previous[:-1] + response_output = previous + response_output + content = get_output_text(response_output) + await event_emitter( { 'type': 'chat:completion', 'data': { + **(response_data if continuing else {}), 'done': True, 'output': response_output, 'title': title,