From a8b478ed7b8d22fe5b8f96be3c4cfd08cd30affb Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Tue, 3 Nov 2026 17:17:59 +0900 Subject: [PATCH] refac --- .../open_webui/utils/context_compaction.py | 117 ++++++++++++-- backend/open_webui/utils/middleware.py | 145 ++++++++++++++---- 2 files changed, 212 insertions(+), 50 deletions(-) diff --git a/backend/open_webui/utils/context_compaction.py b/backend/open_webui/utils/context_compaction.py index 87aa9335ca..da2c03c6e4 100644 --- a/backend/open_webui/utils/context_compaction.py +++ b/backend/open_webui/utils/context_compaction.py @@ -1,5 +1,6 @@ from __future__ import annotations +import copy import logging import re from typing import Any @@ -44,16 +45,18 @@ Summarize the conversation history that will be compacted out of the active chat async def compact_messages_for_request( request, + form_data: dict | None, user, messages: list[dict], metadata: dict, model_id: str, models: dict, system_prompt: str = '', -) -> tuple[list[dict], str | None, bool]: + bypass_system_prompt: bool = False, +) -> tuple[list[dict], str | None, list[dict]]: config = await _load_config() if not config['enable']: - return messages, None, False + return messages, None, [] system_messages = [messages[0]] if messages and messages[0].get('role') == 'system' else [] messages = messages[1:] if system_messages else messages @@ -62,13 +65,13 @@ async def compact_messages_for_request( system_prompt = system_prompt or (get_content_from_message(system_messages[0]) if system_messages else '') token_threshold = _resolve_token_threshold(config['token_threshold'], config['token_cap'], metadata) if not _exceeds_token_threshold(messages, system_prompt, previous_summary, token_threshold): - return [*system_messages, *messages], previous_summary, False + return [*system_messages, *messages], previous_summary, [] boundary = _find_compaction_boundary(messages, config['retention_percentage']) compacted_messages = messages[:boundary] recent_messages = messages[boundary:] if not compacted_messages or not recent_messages: - return [*system_messages, *messages], previous_summary, False + return [*system_messages, *messages], previous_summary, [] event_emitter = None if metadata.get('chat_id') and metadata.get('message_id'): @@ -91,6 +94,7 @@ async def compact_messages_for_request( try: summary = await _generate_summary( request, + form_data, user, model_id, models, @@ -98,6 +102,7 @@ async def compact_messages_for_request( recent_messages, previous_summary, config['prompt_template'], + bypass_system_prompt, ) except Exception: if event_emitter: @@ -153,7 +158,18 @@ async def compact_messages_for_request( } ) - return [*system_messages, *recent_messages], summary, True + return [*system_messages, *recent_messages], summary, compacted_messages + + +async def apply_summary_checkpoint(messages: list[dict]) -> tuple[list[dict], str | None]: + if not (await _load_config())['enable']: + return messages, None + + system_messages = [messages[0]] if messages and messages[0].get('role') == 'system' else [] + messages = messages[1:] if system_messages else messages + + messages, summary = _apply_latest_summary_checkpoint(messages) + return [*system_messages, *messages], summary async def compact_chat_branch(request, user, chat: Any, model_id: str, models: dict) -> dict: @@ -183,6 +199,7 @@ async def compact_chat_branch(request, user, chat: Any, model_id: str, models: d summary = await _generate_summary( request, + None, user, model_id, models, @@ -385,6 +402,7 @@ def _find_compaction_boundary(messages: list[dict], retention_percentage: int = async def _generate_summary( request, + form_data: dict | None, user, model_id: str, models: dict, @@ -392,6 +410,7 @@ async def _generate_summary( recent_messages: list[dict], previous_summary: str | None, summary_prompt_template: str, + bypass_system_prompt: bool = False, ) -> str: from open_webui.utils.chat import generate_chat_completion @@ -408,6 +427,14 @@ async def _generate_summary( raise ValueError('No available model for context compaction') summary_prompt_template = summary_prompt_template.strip() or DEFAULT_CONTEXT_COMPACTION_PROMPT + recent_message_prefix = _message_prefix(recent_messages[0]) + has_unique_prefix = ( + recent_message_prefix + and sum( + _message_prefix(message) == recent_message_prefix for message in [*compacted_messages, *recent_messages] + ) + == 1 + ) # The text-only prompt formatter otherwise omits tool names and arguments. compacted_messages, recent_messages = list(compacted_messages), list(recent_messages) for messages in (compacted_messages, recent_messages): @@ -417,13 +444,15 @@ async def _generate_summary( **message, 'content': f'{get_content_from_message(message) or ""}\n[TOOL CALLS] {JSONCodec.dumps(message["tool_calls"])}', } - all_messages = [*compacted_messages, *recent_messages] - prompt = replace_prompt_variable(summary_prompt_template, get_last_user_message(all_messages) or '') - prompt = replace_messages_variable(prompt, all_messages) - prompt = replace_messages_variable(prompt, compacted_messages, 'COMPACTED_MESSAGES') - prompt = replace_messages_variable(prompt, recent_messages, 'RECENT_MESSAGES') - prompt = prompt_variables_template(prompt, {'{{PREVIOUS_SUMMARY}}': previous_summary or ''}) - prompt = await prompt_template(prompt, user) + + async def render_prompt(template: str) -> str: + all_messages = [*compacted_messages, *recent_messages] + prompt = replace_prompt_variable(template, get_last_user_message(all_messages) or '') + prompt = replace_messages_variable(prompt, all_messages) + prompt = replace_messages_variable(prompt, compacted_messages, 'COMPACTED_MESSAGES') + prompt = replace_messages_variable(prompt, recent_messages, 'RECENT_MESSAGES') + prompt = prompt_variables_template(prompt, {'{{PREVIOUS_SUMMARY}}': previous_summary or ''}) + return await prompt_template(prompt, user) task_model_params = task_config.get('task.model.params') or {} if not isinstance(task_model_params, dict): @@ -433,14 +462,60 @@ async def _generate_summary( 'max_tokens': models[task_model_id].get('info', {}).get('params', {}).get('max_tokens', 1000) } + metadata = { + **(request.state.metadata if hasattr(request.state, 'metadata') else {}), + 'task': 'context_compaction', + } + + # Replaying the chat request unchanged lets the provider read the history from its prompt cache. + # A repeated or transformed boundary cannot safely identify the history to discard. + if ( + form_data + and task_model_id == model_id + and has_unique_prefix + and sum(_message_prefix(message) == recent_message_prefix for message in form_data['messages']) == 1 + ): + prompt = await render_prompt( + prompt_variables_template( + summary_prompt_template, + { + '{{COMPACTED_MESSAGES}}': ( + 'Every non-system message above before (excluding) the message beginning with ' + f'"{recent_message_prefix}".' + ), + '{{RECENT_MESSAGES}}': ( + f'The message above beginning with "{recent_message_prefix}" and every message after it.' + ), + }, + ) + ) + payload = { + 'model': task_model_id, + 'messages': [*copy.deepcopy(form_data['messages']), {'role': 'user', 'content': prompt}], + 'stream': False, + 'metadata': metadata, + } + if form_data.get('tools'): + payload['tools'] = copy.deepcopy(form_data['tools']) + try: + response = await generate_chat_completion( + request, + form_data=apply_params_to_form_data(payload, models[task_model_id], task_model_params), + user=user, + bypass_system_prompt=bypass_system_prompt, + ) + except Exception: + response = None + summary = _response_text(response).strip() + if summary: + return summary + + prompt = await render_prompt(summary_prompt_template) payload = { 'model': task_model_id, 'messages': [{'role': 'user', 'content': prompt}], 'stream': False, - 'metadata': { - **(request.state.metadata if hasattr(request.state, 'metadata') else {}), - 'task': 'context_compaction', - }, + 'metadata': metadata, } payload = apply_params_to_form_data(payload, models[task_model_id], task_model_params) @@ -457,6 +532,14 @@ async def _generate_summary( return '\n'.join(parts)[:4000] +def _message_prefix(message: dict) -> str: + content = get_content_from_message(message) or ' '.join( + f'{tool_call["function"]["name"]}({tool_call["function"]["arguments"]})' + for tool_call in message.get('tool_calls') or [] + ) + return content.strip()[:200] + + def _response_text(response: Any) -> str: if isinstance(response, list) and len(response) == 1: response = response[0] @@ -473,6 +556,8 @@ def _response_text(response: Any) -> str: choices = response.get('choices') or [] if choices: message = choices[0].get('message') or {} + if message.get('tool_calls'): + return '' return message.get('content') or message.get('reasoning_content') or '' parts = [] diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 1db1b76ede..e2284b1118 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -89,7 +89,7 @@ from open_webui.utils.audit import log_model_usage from open_webui.utils.chat import generate_chat_completion from open_webui.utils.chat_id import is_saved_chat_id, is_temporary_chat_id from open_webui.utils.code_interpreter import execute_code_jupyter -from open_webui.utils.context_compaction import compact_messages_for_request +from open_webui.utils.context_compaction import apply_summary_checkpoint, compact_messages_for_request from open_webui.utils.files import ( convert_markdown_base64_images, get_file_url_from_base64, @@ -2178,27 +2178,44 @@ async def process_chat_payload(request, form_data, user, metadata, model): if regeneration_prompt: form_data['messages'].append({'role': 'user', 'content': regeneration_prompt}) + filter_functions = ( + await get_filter_functions(request, model, metadata.get('filter_ids', [])) if ENABLE_FUNCTIONS else [] + ) + compaction_messages = None + context_summary = None if (is_saved_chat_id(chat_id) or is_temporary_chat_id(chat_id)) and user_message_id: system_message = get_system_message(form_data.get('messages', [])) chat_system_prompt = get_content_from_message(system_message) if system_message else '' try: - form_data['messages'], context_summary, _ = await compact_messages_for_request( - request, - user, - form_data.get('messages', []), - metadata, - form_data.get('model'), - request.app.state.MODELS, - chat_system_prompt, - ) - if context_summary: - form_data['messages'] = add_or_update_system_message( - f'[CONVERSATION SUMMARY]\n{context_summary}', - form_data['messages'], - append=True, + # Filters and resumed tool calls must still receive compacted history. + if ( + filter_functions + or 'pipeline' in model + or model.get('pipe') + or get_sorted_filters(form_data['model'], request.app.state.MODELS) + or pending_assistant_message is not None + or metadata.get('params', {}).get('function_calling') == 'legacy' + ): + form_data['messages'], context_summary, _ = await compact_messages_for_request( + request, + None, + user, + form_data.get('messages', []), + metadata, + form_data['model'], + request.app.state.MODELS, + chat_system_prompt, ) + if context_summary: + form_data['messages'] = add_or_update_system_message( + f'[CONVERSATION SUMMARY]\n{context_summary}', form_data['messages'], append=True + ) + else: + compaction_messages = form_data.get('messages', []) + form_data['messages'], context_summary = await apply_summary_checkpoint(compaction_messages) except Exception: + compaction_messages = None log.exception('Context compaction failed; continuing with full chat history') # Compact the full chain, but prepare history separately from the response being resumed. @@ -2223,6 +2240,7 @@ async def process_chat_payload(request, form_data, user, metadata, model): include_file_context=metadata['include_file_context'], ) form_data['messages'] = sanitize_tool_pairs(form_data['messages']) + history_messages = [message for message in form_data['messages'] if message.get('role') != 'system'] system_message = get_system_message(form_data.get('messages', [])) if system_message: # Chat Controls/User Settings @@ -2349,12 +2367,9 @@ async def process_chat_payload(request, form_data, user, metadata, model): except Exception as e: raise e - filter_functions = [] filter_context = get_filter_context(request) if ENABLE_FUNCTIONS else None if ENABLE_FUNCTIONS: try: - filter_functions = await get_filter_functions(request, model, metadata.get('filter_ids', [])) - form_data, flags = await process_filter_functions( request=request, filter_context=filter_context, @@ -2983,6 +2998,51 @@ async def process_chat_payload(request, form_data, user, metadata, model): form_data = normalize_messages_for_model(form_data) + if compaction_messages is not None: + summary_form_data = {**form_data} + if context_summary: + summary_form_data['messages'] = add_or_update_system_message( + f'[CONVERSATION SUMMARY]\n{context_summary}', + [dict(message) for message in form_data['messages']], + append=True, + ) + try: + _, context_summary, compacted = await compact_messages_for_request( + request, + summary_form_data, + user, + compaction_messages, + metadata, + form_data['model'], + request.app.state.MODELS, + chat_system_prompt, + ) + if compacted: + # Compacted messages lead the history this request was built from. + compacted_count = len( + sanitize_tool_pairs( + process_messages_with_output( + compacted, + reasoning_format=get_reasoning_format(model), + include_file_context=metadata['include_file_context'], + ) + ) + ) + compacted_ids = {id(message) for message in history_messages[:compacted_count]} + form_data['messages'] = [ + message for message in form_data['messages'] if id(message) not in compacted_ids + ] + except Exception: + log.exception('Context compaction failed; continuing with full chat history') + + if context_summary: + summary_content = f'[CONVERSATION SUMMARY]\n{context_summary}' + form_data['messages'] = add_or_update_system_message(summary_content, form_data['messages'], append=True) + base_system_prompt = metadata['base_system_prompt'] + metadata['base_system_prompt'] = ( + f'{base_system_prompt}\n{summary_content}' if base_system_prompt else summary_content + ) + return form_data, metadata, events, False @@ -5943,6 +6003,25 @@ async def streaming_chat_response_handler(response, ctx): } ) + # Use the existing checkpoint format internally and strip it from the provider payload. + def get_request_messages(messages: list[dict]) -> list[dict]: + request_messages = process_messages_with_output(messages) + context_summary = next( + ( + message['contextSummary'] + for message in messages + if message.get('contextSummary') + ), + None, + ) + if context_summary: + request_messages.insert( + 1 if request_messages and request_messages[0].get('role') == 'system' else 0, + {'role': 'system', 'content': f'[CONVERSATION SUMMARY]\n{context_summary}'}, + ) + return request_messages + + new_form_data['messages'] = get_request_messages(messages) try: user_messages = [ (message, message.get('content')) @@ -5951,6 +6030,11 @@ async def streaming_chat_response_handler(response, ctx): and message.get('role') == 'user' and message.get('content') == user_prompt.get('content') ] + summary_form_data = None + if not filter_functions: + summary_form_data = normalize_messages_for_model( + await convert_url_images_to_base64({**new_form_data}, user=user) + ) try: # Summarize the user's text, then restore source context for the model. for message, _ in user_messages: @@ -5959,7 +6043,14 @@ async def streaming_chat_response_handler(response, ctx): metadata.get('base_user_prompt') or user_message, [message] ) messages, context_summary, compacted = await compact_messages_for_request( - request, user, messages, metadata, form_data['model'], request.app.state.MODELS + request, + summary_form_data, + user, + messages, + metadata, + form_data['model'], + request.app.state.MODELS, + bypass_system_prompt=True, ) finally: for message, content in user_messages: @@ -5975,6 +6066,7 @@ async def streaming_chat_response_handler(response, ctx): idx for idx, message in enumerate(messages) if message.get('role') != 'system' ) messages[checkpoint] = {**messages[checkpoint], 'contextSummary': context_summary} + new_form_data['messages'] = get_request_messages(messages) output.append( { 'type': 'open_webui:context_compaction', @@ -5991,21 +6083,6 @@ async def streaming_chat_response_handler(response, ctx): except Exception: log.exception('Tool loop context compaction failed; keeping the current context') - # Use the existing checkpoint format internally and strip it from the provider payload. - new_form_data['messages'] = process_messages_with_output(messages) - context_summary = next( - (message['contextSummary'] for message in messages if message.get('contextSummary')), - None, - ) - if context_summary: - new_form_data['messages'].insert( - 1 - if new_form_data['messages'] - and new_form_data['messages'][0].get('role') == 'system' - else 0, - {'role': 'system', 'content': f'[CONVERSATION SUMMARY]\n{context_summary}'}, - ) - new_form_data = await convert_url_images_to_base64(new_form_data, user=user) if filter_functions: