diff --git a/backend/open_webui/utils/context_compaction.py b/backend/open_webui/utils/context_compaction.py index 19d44f417d..d3955148cb 100644 --- a/backend/open_webui/utils/context_compaction.py +++ b/backend/open_webui/utils/context_compaction.py @@ -321,12 +321,10 @@ def _exceeds_token_threshold(messages: list[dict], system_prompt: str, summary: 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..1adfba67fd 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: