diff --git a/backend/open_webui/utils/context_compaction.py b/backend/open_webui/utils/context_compaction.py index 19d44f417d..c70c7785b7 100644 --- a/backend/open_webui/utils/context_compaction.py +++ b/backend/open_webui/utils/context_compaction.py @@ -305,8 +305,13 @@ 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 + cut_idx = None + kept_output = None for idx, message in enumerate(messages): + if checkpoint := _find_tool_loop_checkpoint(messages, idx): + summary, cut_idx, kept_output = checkpoint + summary_idx = cut_idx - 1 value = message.get('contextSummary') or message.get('context_summary') if isinstance(value, str) and value.strip(): summary = value @@ -314,19 +319,46 @@ def _apply_latest_summary_checkpoint(messages: list[dict]) -> tuple[list[dict], if summary_idx is None: return messages, None - return messages[summary_idx:], summary + return [ + {**message, 'output': kept_output} if idx == cut_idx else message + for idx, message in enumerate(messages) + if idx >= summary_idx + ], summary + + +def _find_tool_loop_checkpoint(messages: list[dict], idx: int) -> tuple[str, int, list[dict]] | None: + """Tool loop summary saved in this reply, with the reply it cut and that reply's remaining output.""" + output = messages[idx].get('output') or [] + marker = next((item for item in reversed(output) if item.get('type') == 'open_webui:context_compaction'), None) + if not marker: + return None + # The loop can cut into an earlier reply when this one started near the threshold. + for cut_idx in range(idx, -1, -1): + cut_output = messages[cut_idx].get('output') or [] + start = next( + ( + position + for position, item in enumerate(cut_output) + if item.get('type') == 'function_call' and item.get('call_id') == marker['call_id'] + ), + None, + ) + if start is not None: + while start and cut_output[start - 1].get('type') != 'function_call_output': + start -= 1 + kept_output = [item for item in cut_output[start:] if item.get('type') != 'open_webui:context_compaction'] + return marker['summary'], cut_idx, kept_output + return 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 57ac2e7517..76b1e8950d 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -5917,6 +5917,7 @@ async def streaming_chat_response_handler(response, ctx): ) tool_call_sources = [] # Track citation sources from tool results all_tool_call_sources = [] # Accumulated sources across all iterations + attached_usage = None user_message = get_last_user_message(form_data['messages']) messages = list(form_data['messages']) @@ -6280,6 +6281,11 @@ async def streaming_chat_response_handler(response, ctx): image_urls.append(part.get('image_url', '')) message['content'] = ''.join(text_parts) + # A count reported for the last request covers everything before these results. + if usage is not attached_usage and tool_messages: + tool_messages[0]['usage'] = usage + attached_usage = usage + # Refresh citation context without restoring compacted history. user_prompt = get_last_user_message_item(form_data['messages']) for message in messages: @@ -6320,6 +6326,20 @@ 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} + kept_call_id = next( + message['tool_calls'][0]['id'] + for message in messages + if message.get('tool_calls') + ) + output.append( + { + 'type': 'open_webui:context_compaction', + 'id': output_id('cc'), + 'call_id': kept_call_id, + 'summary': context_summary, + } + ) + await emit_output() except Exception: log.exception('Tool loop context compaction failed; keeping the current context')