refac

backend/open_webui/utils/context_compaction.py

@@ -0,0 +1,316 @@
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any
+
+from fastapi.responses import JSONResponse
+
+from open_webui.models.chats import Chats
+from open_webui.models.config import Config
+from open_webui.utils.misc import get_content_from_message, get_last_user_message, get_message_list
+from open_webui.utils.task import (
+    get_task_model_id,
+    prompt_template,
+    prompt_variables_template,
+    replace_messages_variable,
+    replace_prompt_variable,
+)
+
+log = logging.getLogger(__name__)
+
+DEFAULT_CONTEXT_COMPACTION_PROMPT = """### Task:
+Summarize the conversation history that will be compacted out of the active chat context.
+
+### Instructions:
+- Preserve key decisions, user preferences, and constraints.
+- Preserve files, artifacts, tool results, and code changes that matter going forward.
+- Preserve the current task state, unresolved questions, and next steps.
+- Be factual and specific. Do not invent details.
+- Keep the summary concise, but complete enough for the assistant to continue without the removed messages.
+
+### Previous Summary:
+{{PREVIOUS_SUMMARY}}
+
+### Messages Being Compacted:
+{{COMPACTED_MESSAGES}}
+
+### Recent Messages Kept In Context:
+{{RECENT_MESSAGES}}"""
+
+
+async def compact_messages_for_request(
+    request,
+    user,
+    messages: list[dict],
+    metadata: dict,
+    model_id: str,
+    models: dict,
+    system_prompt: str = '',
+) -> tuple[list[dict], str | None, bool]:
+    config = await _load_config()
+    if not config['enable']:
+        return messages, None, False
+
+    messages, previous_summary = _apply_latest_summary_checkpoint(messages)
+    if (
+        not _exceeds_token_threshold(messages, system_prompt, previous_summary, config['token_threshold'])
+        or len(messages) <= 3
+    ):
+        return messages, previous_summary, False
+
+    boundary = _find_compaction_boundary(messages)
+    compacted_messages = messages[:boundary]
+    recent_messages = messages[boundary:]
+    if not compacted_messages or not recent_messages:
+        return messages, previous_summary, False
+
+    summary = await _generate_summary(
+        request,
+        user,
+        model_id,
+        models,
+        compacted_messages,
+        recent_messages,
+        previous_summary,
+        config['prompt_template'],
+    )
+
+    chat_id = metadata.get('chat_id')
+    message_id = metadata.get('message_id')
+    if chat_id and message_id and not chat_id.startswith(('local:', 'channel:')):
+        await Chats.upsert_message_to_chat_by_id_and_message_id(chat_id, message_id, {'contextSummary': summary})
+
+    log.info(
+        'Compacted chat context for chat=%s message=%s dropped=%d kept=%d summary_chars=%d',
+        chat_id,
+        message_id,
+        len(compacted_messages),
+        len(recent_messages),
+        len(summary),
+    )
+
+    return recent_messages, summary, True
+
+
+async def compact_chat_branch(request, user, chat: Any, model_id: str, models: dict) -> dict:
+    config = await _load_config()
+    if not config['enable']:
+        return {'ok': True, 'compacted': False, 'reason': 'disabled'}
+
+    history = (chat.chat or {}).get('history') or {}
+    current_id = history.get('currentId')
+    if not current_id:
+        return {'ok': True, 'compacted': False, 'reason': 'empty'}
+
+    messages_map = await Chats.get_messages_map_by_chat_id(chat.id)
+    if not messages_map:
+        messages_map = history.get('messages') or {}
+
+    messages, previous_summary = _apply_latest_summary_checkpoint(get_message_list(messages_map, current_id))
+    if len(messages) <= 2:
+        return {'ok': True, 'compacted': False, 'reason': 'too_short'}
+
+    compacted_messages = messages[:-1]
+    recent_messages = messages[-1:]
+    summary = await _generate_summary(
+        request,
+        user,
+        model_id,
+        models,
+        compacted_messages,
+        recent_messages,
+        previous_summary,
+        config['prompt_template'],
+    )
+    await Chats.upsert_message_to_chat_by_id_and_message_id(chat.id, current_id, {'contextSummary': summary})
+
+    return {
+        'ok': True,
+        'compacted': True,
+        'dropped_messages': len(compacted_messages),
+        'kept_messages': len(recent_messages),
+        'summary_chars': len(summary),
+    }
+
+
+async def _load_config() -> dict:
+    values = await Config.get_many(
+        'chat.context_compaction.enable',
+        'chat.context_compaction.token_threshold',
+        'chat.context_compaction.prompt_template',
+    )
+    return {
+        'enable': bool(values.get('chat.context_compaction.enable', False)),
+        'token_threshold': int(values.get('chat.context_compaction.token_threshold', 80000) or 80000),
+        'prompt_template': values.get('chat.context_compaction.prompt_template', '') or '',
+    }
+
+
+def _apply_latest_summary_checkpoint(messages: list[dict]) -> tuple[list[dict], str | None]:
+    summary = None
+    summary_idx = None
+
+    for idx, message in enumerate(messages):
+        value = message.get('contextSummary') or message.get('context_summary')
+        if isinstance(value, str) and value.strip():
+            summary = value
+            summary_idx = idx
+
+    if summary_idx is None:
+        return messages, None
+    return messages[summary_idx:], summary
+
+
+def _exceeds_token_threshold(messages: list[dict], system_prompt: str, summary: str | None, threshold: int) -> bool:
+    if threshold <= 0:
+        return False
+
+    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 usage.get('input_tokens'):
+            total = int(usage.get('input_tokens') or 0) + int(usage.get('output_tokens') or 0)
+            return total + _estimate_messages_tokens(messages[idx + 1 :]) > threshold
+
+    estimated = _estimate_tokens(system_prompt) + _estimate_tokens(summary or '') + _estimate_messages_tokens(messages)
+    return estimated > threshold
+
+
+def _find_compaction_boundary(messages: list[dict]) -> int:
+    keep_count = max(2, len(messages) * 2 // 5)
+    split = max(1, len(messages) - keep_count)
+
+    while split < len(messages) - 1:
+        previous = messages[split - 1] if split > 0 else {}
+        current = messages[split]
+        if current.get('role') == 'tool' or previous.get('tool_calls') or previous.get('output'):
+            split += 1
+            continue
+        break
+
+    return min(split, len(messages) - 2)
+
+
+async def _generate_summary(
+    request,
+    user,
+    model_id: str,
+    models: dict,
+    compacted_messages: list[dict],
+    recent_messages: list[dict],
+    previous_summary: str | None,
+    summary_prompt_template: str,
+) -> str:
+    from open_webui.utils.chat import generate_chat_completion
+
+    task_model_id = get_task_model_id(
+        model_id,
+        await Config.get('task.model.default'),
+        await Config.get('task.model.external'),
+        models,
+    )
+    if task_model_id not in models:
+        task_model_id = model_id
+    if task_model_id not in models:
+        raise ValueError('No available model for context compaction')
+
+    summary_prompt_template = summary_prompt_template.strip() or DEFAULT_CONTEXT_COMPACTION_PROMPT
+    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)
+
+    max_tokens = models[task_model_id].get('info', {}).get('params', {}).get('max_tokens', 1000)
+    payload = {
+        'model': task_model_id,
+        'messages': [{'role': 'user', 'content': prompt}],
+        'stream': False,
+        **(
+            {'max_tokens': max_tokens}
+            if models[task_model_id].get('owned_by') == 'ollama'
+            else {'max_completion_tokens': max_tokens}
+        ),
+        'metadata': {
+            **(request.state.metadata if hasattr(request.state, 'metadata') else {}),
+            'task': 'context_compaction',
+        },
+    }
+
+    response = await generate_chat_completion(request, form_data=payload, user=user)
+    summary = _response_text(response).strip()
+    if summary:
+        return summary
+
+    parts = [previous_summary] if previous_summary else []
+    for message in compacted_messages:
+        content = get_content_from_message(message)
+        if content:
+            parts.append(f'- {message.get("role", "unknown")}: {content[:500]}')
+    return '\n'.join(parts)[:4000]
+
+
+def _response_text(response: Any) -> str:
+    if isinstance(response, list) and len(response) == 1:
+        response = response[0]
+
+    if isinstance(response, JSONResponse):
+        try:
+            response = json.loads(response.body.decode('utf-8', 'replace'))
+        except Exception:
+            return ''
+
+    if not isinstance(response, dict):
+        return ''
+
+    choices = response.get('choices') or []
+    if choices:
+        message = choices[0].get('message') or {}
+        return message.get('content') or message.get('reasoning_content') or ''
+
+    parts = []
+    for item in response.get('output') or []:
+        for content in item.get('content') or []:
+            if isinstance(content, dict):
+                parts.append(content.get('text') or content.get('content') or '')
+    return '\n'.join(part for part in parts if part)
+
+
+def _estimate_messages_tokens(messages: list[dict]) -> int:
+    total = 0
+    for message in messages:
+        total += 4
+        content = message.get('content')
+        if isinstance(content, list):
+            for item in content:
+                if not isinstance(item, dict):
+                    total += _estimate_tokens(item)
+                elif item.get('type') in {'image', 'image_url'}:
+                    total += 1000
+                else:
+                    total += _estimate_tokens(item.get('text') or item.get('content') or item)
+        else:
+            total += _estimate_tokens(content)
+
+        total += _estimate_tokens(message.get('output'))
+        total += _estimate_tokens(message.get('tool_calls'))
+        total += _estimate_tokens(message.get('files'))
+    return total
+
+
+def _estimate_tokens(value: Any) -> int:
+    if value is None:
+        return 0
+
+    if not isinstance(value, str):
+        try:
+            value = json.dumps(value, ensure_ascii=False)
+        except Exception:
+            value = str(value)
+
+    if not value:
+        return 0
+
+    return max(1, len(value) // 4)

backend/open_webui/utils/middleware.py

@@ -77,6 +77,7 @@ from open_webui.utils.access_control import has_connection_access, has_permissio
 from open_webui.utils.access_control.files import get_accessible_folder_files
 from open_webui.utils.chat import generate_chat_completion
 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.files import (
     convert_markdown_base64_images,
     get_file_url_from_base64,
@@ -2170,7 +2171,10 @@ async def load_messages_from_db(chat_id: str, message_id: str) -> Optional[list[
     if not db_messages:
         return None
 
-    return [{k: v for k, v in msg.items() if k in ('role', 'content', 'output', 'files')} for msg in db_messages]
+    return [
+        {k: v for k, v in msg.items() if k in ('role', 'content', 'output', 'files', 'contextSummary')}
+        for msg in db_messages
+    ]
 
 
 def get_reasoning_format(model: dict) -> str | None:
@@ -2221,6 +2225,50 @@ def process_messages_with_output(
     return processed
 
 
+def strip_compaction_fields(messages: list[dict]) -> list[dict]:
+    stripped = []
+    for message in messages:
+        clean = dict(message)
+        clean.pop('contextSummary', None)
+        clean.pop('context_summary', None)
+        stripped.append(clean)
+    return stripped
+
+
+def sanitize_tool_pairs(messages: list[dict]) -> list[dict]:
+    tool_result_ids = {
+        message.get('tool_call_id')
+        for message in messages
+        if message.get('role') == 'tool' and message.get('tool_call_id')
+    }
+
+    tool_call_ids = {
+        tool_call.get('id')
+        for message in messages
+        for tool_call in (message.get('tool_calls') or [])
+        if message.get('role') == 'assistant' and tool_call.get('id')
+    }
+
+    sanitized = []
+    for message in messages:
+        if message.get('role') == 'assistant' and message.get('tool_calls'):
+            kept = [
+                tool_call for tool_call in message.get('tool_calls') or [] if tool_call.get('id') in tool_result_ids
+            ]
+            if kept:
+                sanitized.append({**message, 'tool_calls': kept})
+            else:
+                clean = dict(message)
+                clean.pop('tool_calls', None)
+                clean.pop('reasoning_items', None)
+                if clean.get('content'):
+                    sanitized.append(clean)
+        elif message.get('role') != 'tool' or message.get('tool_call_id') in tool_call_ids:
+            sanitized.append(message)
+
+    return sanitized
+
+
 SKILL_MENTION_RE = re.compile(r'<\$([^|>]+)\|?[^>]*>')
 
 
@@ -2368,7 +2416,11 @@ async def process_chat_payload(request, form_data, user, metadata, model):
                 assistant_message = await Chats.get_message_by_id_and_message_id(chat_id, assistant_message_id)
                 if assistant_message and (assistant_message.get('content') or assistant_message.get('output')):
                     db_messages.append(
-                        {k: v for k, v in assistant_message.items() if k in ('role', 'content', 'output', 'files')}
+                        {
+                            k: v
+                            for k, v in assistant_message.items()
+                            if k in ('role', 'content', 'output', 'files', 'contextSummary')
+                        }
                     )
 
             system_message = get_system_message(form_data.get('messages', []))
@@ -2401,11 +2453,44 @@ async def process_chat_payload(request, form_data, user, metadata, model):
     if regeneration_prompt:
         form_data['messages'].append({'role': 'user', 'content': regeneration_prompt})
 
+    if chat_id and user_message_id and not chat_id.startswith('local:') and not chat_id.startswith('channel:'):
+        if getattr(request.state, 'direct', False) and hasattr(request.state, 'model'):
+            compaction_models = {
+                request.state.model['id']: request.state.model,
+            }
+        else:
+            compaction_models = request.app.state.MODELS
+
+        system_message = get_system_message(form_data.get('messages', []))
+        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'),
+                compaction_models,
+                system_prompt,
+            )
+            if context_summary:
+                form_data['messages'] = add_or_update_system_message(
+                    f'[CONVERSATION SUMMARY]\n{context_summary}',
+                    form_data['messages'],
+                    append=True,
+                )
+        except Exception:
+            log.exception('Context compaction failed; continuing with full chat history')
+
+    form_data['messages'] = strip_compaction_fields(form_data.get('messages', []))
+
     # Process messages with OR-aligned output items for clean LLM messages
     form_data['messages'] = process_messages_with_output(
         form_data.get('messages', []),
         reasoning_format=get_reasoning_format(model),
     )
+    form_data['messages'] = sanitize_tool_pairs(form_data['messages'])
 
     system_message = get_system_message(form_data.get('messages', []))
     if system_message:  # Chat Controls/User Settings

backend/open_webui/utils/task.py

@@ -194,7 +194,9 @@ def apply_content_filter(messages: list[dict], filter_str: str) -> list[dict]:
     return result
 
 
-def replace_messages_variable(template: str, messages: Optional[list[dict]] = None) -> str:
+def replace_messages_variable(
+    template: str, messages: Optional[list[dict]] = None, variable_name: str = 'MESSAGES'
+) -> str:
     def replacement_function(match):
         # Groups: (1) filter for bare MESSAGES
         #         (2) START count, (3) filter for START
@@ -240,12 +242,13 @@ def replace_messages_variable(template: str, messages: Optional[list[dict]] = No
 
         return get_messages_content(selected)
 
+    variable_pattern = re.escape(variable_name)
     template = re.sub(
         r'(?:'
-        r'\{\{MESSAGES(?:\|(\w+:\d+))?\}\}'
-        r'|\{\{MESSAGES:START:(\d+)(?:\|(\w+:\d+))?\}\}'
-        r'|\{\{MESSAGES:END:(\d+)(?:\|(\w+:\d+))?\}\}'
-        r'|\{\{MESSAGES:MIDDLETRUNCATE:(\d+)(?:\|(\w+:\d+))?\}\}'
+        rf'\{{\{{{variable_pattern}(?:\|(\w+:\d+))?\}}\}}'
+        rf'|\{{\{{{variable_pattern}:START:(\d+)(?:\|(\w+:\d+))?\}}\}}'
+        rf'|\{{\{{{variable_pattern}:END:(\d+)(?:\|(\w+:\d+))?\}}\}}'
+        rf'|\{{\{{{variable_pattern}:MIDDLETRUNCATE:(\d+)(?:\|(\w+:\d+))?\}}\}}'
         r')',
         replacement_function,
         template,