Merge 106f8ef47e into 93d25dbe3b

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

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')