From bcf540a98935566537b05911a6bb4ace03c5648d Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Wed, 14 Oct 2026 18:35:56 +0400 Subject: [PATCH] refac --- .../open_webui/utils/context_compaction.py | 59 ++++++++++++++----- backend/open_webui/utils/middleware.py | 25 +++++++- 2 files changed, 67 insertions(+), 17 deletions(-) diff --git a/backend/open_webui/utils/context_compaction.py b/backend/open_webui/utils/context_compaction.py index 19d44f417d..7f9649bdd7 100644 --- a/backend/open_webui/utils/context_compaction.py +++ b/backend/open_webui/utils/context_compaction.py @@ -303,30 +303,59 @@ def _build_context_usage(tokens: int, threshold: int) -> dict: def _apply_latest_summary_checkpoint(messages: list[dict]) -> tuple[list[dict], str | None]: - summary = None - summary_idx = None - - for idx, message in enumerate(messages): + for idx in range(len(messages) - 1, -1, -1): + message = messages[idx] value = message.get('contextSummary') or message.get('context_summary') if isinstance(value, str) and value.strip(): - summary = value - summary_idx = idx + return messages[idx:], value - if summary_idx is None: - return messages, None - return messages[summary_idx:], summary + marker = next( + ( + item + for item in reversed(message.get('output') or []) + if item.get('type') == 'open_webui:context_compaction' + ), + None, + ) + if not marker: + continue + + # A tool-loop checkpoint can retain output from an earlier reply. + for cut_idx in range(idx, -1, -1): + output = messages[cut_idx].get('output') or [] + start = next( + ( + position + for position, item in enumerate(output) + if item.get('type') == 'function_call' and item.get('call_id') == marker['call_id'] + ), + None, + ) + if start is not None: + while start and output[start - 1].get('type') != 'function_call_output': + start -= 1 + return [ + *messages[max(0, cut_idx - 1) : cut_idx], + { + **messages[cut_idx], + 'output': [ + item for item in output[start:] if item.get('type') != 'open_webui:context_compaction' + ], + }, + *messages[cut_idx + 1 :], + ], marker['summary'] + + return messages, None def _exceeds_token_threshold(messages: list[dict], system_prompt: str, summary: str | None, threshold: int) -> bool: if threshold <= 0: return False - # Expanded tool histories include fresh results and rebuilt prompts absent from prior usage. - if not any(message.get('role') == 'tool' for message in messages): - for idx in range(len(messages) - 1, -1, -1): - usage = messages[idx].get('usage') or (messages[idx].get('info') or {}).get('usage') - if isinstance(usage, dict) and (tokens := _usage_token_count(usage)): - return tokens + _estimate_messages_tokens(messages[idx + 1 :]) > threshold + for idx in range(len(messages) - 1, -1, -1): + usage = messages[idx].get('usage') or (messages[idx].get('info') or {}).get('usage') + if isinstance(usage, dict) and (tokens := _usage_token_count(usage)): + return tokens + _estimate_messages_tokens(messages[idx + 1 :]) > threshold estimated = _estimate_tokens(system_prompt) + _estimate_tokens(summary or '') + _estimate_messages_tokens(messages) return estimated > threshold diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 866cd94c82..0afeffa277 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -4966,6 +4966,7 @@ async def streaming_chat_response_handler(response, ctx): nonlocal last_response_id response_tool_calls = [] + response_usage = {} delta_count = 0 delta_chunk_size = max( @@ -5228,6 +5229,7 @@ async def streaming_chat_response_handler(response, ctx): # Normalize and capture usage for DB persistence if response_metadata.get('usage'): + response_usage = response_metadata['usage'] usage = merge_usage(usage, response_metadata['usage']) response_metadata['usage'] = usage @@ -5270,6 +5272,7 @@ async def streaming_chat_response_handler(response, ctx): raw_usage = data.get('usage', {}) or {} raw_usage.update(data.get('timings', {})) # llama.cpp if raw_usage: + response_usage = raw_usage usage = merge_usage(usage, raw_usage) await event_emitter( { @@ -5906,9 +5909,11 @@ async def streaming_chat_response_handler(response, ctx): if responses_api_tool_calls: tool_calls.append(_split_tool_calls(responses_api_tool_calls)) + return normalize_usage(response_usage) + output_start = len(prior_output) try: - await stream_body_handler(response, form_data) + response_usage = await stream_body_handler(response, form_data) finally: if response.background: await response.background() @@ -6284,6 +6289,9 @@ async def streaming_chat_response_handler(response, ctx): image_urls.append(part.get('image_url', '')) message['content'] = ''.join(text_parts) + if response_usage and tool_messages: + tool_messages[0]['usage'] = response_usage + # Refresh citation context without restoring compacted history. user_prompt = get_last_user_message_item(form_data['messages']) for message in messages: @@ -6324,6 +6332,19 @@ 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} + output.append( + { + 'type': 'open_webui:context_compaction', + 'id': output_id('cc'), + 'call_id': next( + message['tool_calls'][0]['id'] + for message in messages + if message.get('tool_calls') + ), + 'summary': context_summary, + } + ) + await emit_output() except Exception: log.exception('Tool loop context compaction failed; keeping the current context') @@ -6384,7 +6405,7 @@ async def streaming_chat_response_handler(response, ctx): prior_output.pop() output = [] output_start = len(prior_output) - await stream_body_handler(res, new_form_data) + response_usage = await stream_body_handler(res, new_form_data) elif getattr(res, 'status_code', 200) >= 400: await emit_message_error(get_message_error_content(get_response_error_detail(res))) break