scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-11-15 09:31:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

refac

Browse source
This commit is contained in:
Timothy Jaeryang Baek 2026-11-03 17:17:59 +09:00
parent 8e61cc0210
commit a8b478ed7b
2 changed files with 212 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

117
backend/open_webui/utils/context_compaction.py
View file

@ -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 = []

145
backend/open_webui/utils/middleware.py
View file

@ -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: