refac

backend/open_webui/utils/chat_fork.py

@@ -0,0 +1,41 @@
+from copy import deepcopy
+
+
+def build_fork_history(messages_map: dict, source_message_id: str) -> tuple[dict, list[dict]]:
+    if not messages_map:
+        raise ValueError('chat has no messages to fork')
+
+    branch: list[tuple[str, dict]] = []
+    seen: set[str] = set()
+    message_id = source_message_id
+
+    while message_id:
+        if message_id in seen:
+            raise ValueError('message branch contains a cycle')
+        seen.add(message_id)
+
+        message = messages_map.get(message_id)
+        if not isinstance(message, dict):
+            raise ValueError('message not found')
+
+        branch.append((message_id, message))
+        message_id = message.get('parentId')
+
+    fork_messages: dict[str, dict] = {}
+    ordered_messages: list[dict] = []
+    parent_id = None
+
+    for message_id, message in reversed(branch):
+        copied = deepcopy(message)
+        copied['id'] = message_id
+        copied['parentId'] = parent_id
+        copied['childrenIds'] = []
+
+        if parent_id:
+            fork_messages[parent_id]['childrenIds'] = [message_id]
+
+        fork_messages[message_id] = copied
+        ordered_messages.append(copied)
+        parent_id = message_id
+
+    return {'messages': fork_messages, 'currentId': source_message_id}, ordered_messages

backend/open_webui/utils/context_compaction.py

@@ -148,8 +148,13 @@ async def compact_chat_branch(request, user, chat: Any, model_id: str, models: d
     if not config['enable']:
         return {'ok': True, 'compacted': False, 'reason': 'disabled'}
 
-    history = (chat.chat or {}).get('history') or {}
-    current_id = history.get('currentId')
+    chat_data = chat.chat or {}
+    history = chat_data.get('history') or {}
+    current_id = getattr(chat, 'current_message_id', None) or history.get('currentId')
+    if not current_id:
+        current_id = chat_data.get('currentId') or chat_data.get('branchPointMessageId')
+    if not current_id and isinstance(chat_data.get('messages'), list) and chat_data['messages']:
+        current_id = chat_data['messages'][-1].get('id')
     if not current_id:
         return {'ok': True, 'compacted': False, 'reason': 'empty'}
 
@@ -216,8 +221,13 @@ def _resolve_token_threshold(global_threshold: int, global_cap: int, metadata: d
 
 
 async def get_chat_context_usage(chat: Any, model_id: str | None = None) -> dict | None:
-    history = (chat.chat or {}).get('history') or {}
-    current_id = history.get('currentId')
+    chat_data = chat.chat or {}
+    history = chat_data.get('history') or {}
+    current_id = getattr(chat, 'current_message_id', None) or history.get('currentId')
+    if not current_id:
+        current_id = chat_data.get('currentId') or chat_data.get('branchPointMessageId')
+    if not current_id and isinstance(chat_data.get('messages'), list) and chat_data['messages']:
+        current_id = chat_data['messages'][-1].get('id')
     if not current_id:
         return None