From cf887b68ea58bcd1d8b035842c4ea35a5113ed8b Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Thu, 23 Jul 2026 02:54:56 -0400 Subject: [PATCH] refac --- ...2c3d4e5f_add_current_message_id_to_chat.py | 139 ++++++++++++++++++ backend/open_webui/models/chats.py | 43 +++++- backend/open_webui/routers/chats.py | 94 +++++++++++- backend/open_webui/utils/chat_fork.py | 41 ++++++ .../open_webui/utils/context_compaction.py | 18 ++- 5 files changed, 326 insertions(+), 9 deletions(-) create mode 100644 backend/open_webui/migrations/versions/9a1b2c3d4e5f_add_current_message_id_to_chat.py create mode 100644 backend/open_webui/utils/chat_fork.py diff --git a/backend/open_webui/migrations/versions/9a1b2c3d4e5f_add_current_message_id_to_chat.py b/backend/open_webui/migrations/versions/9a1b2c3d4e5f_add_current_message_id_to_chat.py new file mode 100644 index 0000000000..771cb55f52 --- /dev/null +++ b/backend/open_webui/migrations/versions/9a1b2c3d4e5f_add_current_message_id_to_chat.py @@ -0,0 +1,139 @@ +"""add current_message_id to chat + +Revision ID: 9a1b2c3d4e5f +Revises: 856c5b02fb54 +Create Date: 2026-07-23 00:00:00.000000 + +""" + +import json +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = '9a1b2c3d4e5f' +down_revision: Union[str, None] = '856c5b02fb54' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + conn = op.get_bind() + columns = [col['name'] for col in sa.inspect(conn).get_columns('chat')] + if 'current_message_id' not in columns: + op.add_column('chat', sa.Column('current_message_id', sa.Text(), nullable=True)) + + chat = sa.table( + 'chat', + sa.column('id', sa.String()), + sa.column('chat', sa.Text()), + sa.column('current_message_id', sa.Text()), + ) + chat_message = sa.table( + 'chat_message', + sa.column('id', sa.Text()), + sa.column('chat_id', sa.Text()), + sa.column('parent_id', sa.Text()), + sa.column('created_at', sa.BigInteger()), + ) + + messages_by_chat: dict[str, dict[str, dict]] = {} + for row in conn.execute( + sa.select( + chat_message.c.chat_id, + chat_message.c.id, + chat_message.c.parent_id, + chat_message.c.created_at, + ) + ): + values = row._mapping + chat_id = values['chat_id'] + prefix = f'{chat_id}-' + message_id = values['id'] + if message_id and message_id.startswith(prefix): + message_id = message_id[len(prefix) :] + if not message_id: + continue + parent_id = values['parent_id'] + if parent_id and parent_id.startswith(prefix): + parent_id = parent_id[len(prefix) :] + messages_by_chat.setdefault(chat_id, {})[message_id] = { + 'parent_id': parent_id, + 'created_at': values['created_at'] or 0, + } + + for row in conn.execute(sa.select(chat.c.id, chat.c.chat, chat.c.current_message_id)): + values = row._mapping + chat_id = values['id'] + prefix = f'{chat_id}-' + chat_messages = messages_by_chat.get(chat_id, {}) + + chat_data = {} + if isinstance(values['chat'], dict): + chat_data = values['chat'] + elif isinstance(values['chat'], str): + try: + parsed = json.loads(values['chat']) + chat_data = parsed if isinstance(parsed, dict) else {} + except (TypeError, ValueError, json.JSONDecodeError): + chat_data = {} + + history = chat_data.get('history') if isinstance(chat_data.get('history'), dict) else {} + messages = history.get('messages') if isinstance(history.get('messages'), dict) else {} + if not messages and isinstance(chat_data.get('messages'), list): + messages = { + message['id']: message + for message in chat_data['messages'] + if isinstance(message, dict) and message.get('id') + } + json_messages = { + message_id: { + 'parent_id': message.get('parentId') if isinstance(message, dict) else None, + 'created_at': message.get('timestamp', 0) if isinstance(message, dict) else 0, + } + for message_id, message in messages.items() + } + available_messages = chat_messages or json_messages + candidates = [ + values['current_message_id'], + history.get('currentId'), + chat_data.get('currentId'), + chat_data.get('branchPointMessageId'), + ] + + current_message_id = next( + ( + candidate[len(prefix) :] if candidate.startswith(prefix) else candidate + for candidate in candidates + if candidate + and ( + not available_messages + or (candidate[len(prefix) :] if candidate.startswith(prefix) else candidate) in available_messages + ) + ), + None, + ) + + if not current_message_id and available_messages: + parent_ids = { + message['parent_id'] + for message in available_messages.values() + if message.get('parent_id') in available_messages + } + leaf_ids = [message_id for message_id in available_messages if message_id not in parent_ids] + current_message_id = max( + leaf_ids or list(available_messages), + key=lambda message_id: available_messages[message_id].get('created_at') or 0, + ) + + if current_message_id and current_message_id != values['current_message_id']: + conn.execute( + sa.update(chat) + .where(chat.c.id == chat_id) + .values(current_message_id=current_message_id) + ) + + +def downgrade() -> None: + op.drop_column('chat', 'current_message_id') diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py index 1d959f0bc6..4f2213b84d 100644 --- a/backend/open_webui/models/chats.py +++ b/backend/open_webui/models/chats.py @@ -62,6 +62,7 @@ class Chat(Base): # database table mapping for chat entity tasks = Column(JSON, nullable=True) summary = Column(Text, nullable=True) + current_message_id = Column(Text, nullable=True) last_read_at = Column(BigInteger, nullable=True) @@ -98,6 +99,7 @@ class ChatModel(BaseModel): tasks: list | None = None summary: str | None = None + current_message_id: str | None = None last_read_at: int | None = None @@ -145,6 +147,7 @@ class ChatForm(BaseModel): class ChatImportForm(ChatForm): meta: dict | None = {} pinned: bool | None = False + current_message_id: str | None = None created_at: int | None = None updated_at: int | None = None @@ -177,6 +180,7 @@ class ChatResponse(BaseModel): tasks: list | None = None summary: str | None = None + current_message_id: str | None = None context_usage: dict | None = None @@ -280,6 +284,21 @@ class ChatTable: """Recursively remove null bytes from strings in dict/list structures.""" return sanitize_data_for_db(obj) + def get_current_message_id(self, chat: dict | None) -> str | None: + chat = chat or {} + history = chat.get('history') if isinstance(chat.get('history'), dict) else {} + current_id = history.get('currentId') or chat.get('currentId') or chat.get('branchPointMessageId') + if current_id: + return current_id + + messages = chat.get('messages') + if isinstance(messages, list): + for message in reversed(messages): + if isinstance(message, dict) and message.get('id'): + return message['id'] + + return None + def _sanitize_chat_row(self, chat_item): """ Clean a Chat SQLAlchemy model's title + chat JSON, @@ -375,6 +394,7 @@ class ChatTable: 'chat': self._clean_null_bytes(form_data.chat), 'folder_id': form_data.folder_id, 'meta': internal_meta or {}, + 'current_message_id': self.get_current_message_id(form_data.chat), 'created_at': int(time.time()), 'updated_at': int(time.time()), 'last_read_at': int(time.time()), @@ -388,8 +408,14 @@ class ChatTable: # Dual-write initial messages to chat_message table try: - history = form_data.chat.get('history', {}) - messages = history.get('messages', {}) + history = form_data.chat.get('history') if isinstance(form_data.chat.get('history'), dict) else {} + messages = history.get('messages') if isinstance(history.get('messages'), dict) else {} + if not messages and isinstance(form_data.chat.get('messages'), list): + messages = { + message.get('id'): message + for message in form_data.chat['messages'] + if isinstance(message, dict) and message.get('id') + } for message_id, message in messages.items(): if isinstance(message, dict) and message.get('role'): await ChatMessages.upsert_message( @@ -458,6 +484,7 @@ class ChatTable: 'meta': form_data.meta, 'pinned': form_data.pinned, 'folder_id': form_data.folder_id, + 'current_message_id': form_data.current_message_id or self.get_current_message_id(form_data.chat), 'created_at': (form_data.created_at if form_data.created_at else int(time.time())), 'updated_at': (form_data.updated_at if form_data.updated_at else int(time.time())), } @@ -497,8 +524,14 @@ class ChatTable: # Dual-write messages to chat_message table for form_data, chat_obj in zip(chat_import_forms, chats): - history = form_data.chat.get('history', {}) - messages = history.get('messages', {}) + history = form_data.chat.get('history') if isinstance(form_data.chat.get('history'), dict) else {} + messages = history.get('messages') if isinstance(history.get('messages'), dict) else {} + if not messages and isinstance(form_data.chat.get('messages'), list): + messages = { + message.get('id'): message + for message in form_data.chat['messages'] + if isinstance(message, dict) and message.get('id') + } for message_id, message in messages.items(): if isinstance(message, dict) and message.get('role'): try: @@ -530,6 +563,8 @@ class ChatTable: chat_item.chat = self._clean_null_bytes(chat) chat_item.title = self._clean_null_bytes(chat['title']) if 'title' in chat else 'New Chat' + if any(key in chat for key in ('history', 'messages', 'currentId', 'branchPointMessageId')): + chat_item.current_message_id = self.get_current_message_id(chat) if touch: chat_item.updated_at = int(time.time()) diff --git a/backend/open_webui/routers/chats.py b/backend/open_webui/routers/chats.py index 5a565407a6..8747c65a4f 100644 --- a/backend/open_webui/routers/chats.py +++ b/backend/open_webui/routers/chats.py @@ -38,6 +38,7 @@ from open_webui.tasks import has_active_tasks, stop_item_tasks from open_webui.utils.access_control import filter_allowed_access_grants, has_permission from open_webui.utils.access_control.folders import has_folder_access from open_webui.utils.auth import get_admin_user, get_verified_user +from open_webui.utils.chat_fork import build_fork_history from open_webui.utils.context_compaction import compact_chat_branch, get_chat_context_usage from open_webui.utils.misc import get_message_list from open_webui.utils.models import get_all_models @@ -1167,7 +1168,8 @@ async def compact_chat_by_id( history = (chat.chat or {}).get('history') or {} messages_map = await Chats.get_messages_map_by_chat_id(id) - message_list = get_message_list(messages_map or history.get('messages') or {}, history.get('currentId')) + current_message_id = chat.current_message_id or history.get('currentId') + message_list = get_message_list(messages_map or history.get('messages') or {}, current_message_id) model_id = (form_data.model if form_data else None) or next( (message.get('model') for message in reversed(message_list) if message.get('model')), None, @@ -1564,6 +1566,96 @@ class CloneForm(BaseModel): title: str | None = None +class ForkForm(BaseModel): + message_id: str | None = None + + +@router.post('/{id}/fork', response_model=ChatResponse | None) +async def fork_chat_by_id( + request: Request, + id: str, + form_data: ForkForm | None = None, + user=Depends(get_verified_user), + db: AsyncSession = Depends(get_async_session), +): + await require_chat_import_permission(request, user, db) + + chat = await Chats.get_chat_by_id_and_user_id(id, user.id, db=db) + if not chat: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()) + + if await has_active_tasks(request.app.state.redis, id): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail='Wait for the current response to finish before forking.', + ) + + history = (chat.chat or {}).get('history') or {} + messages_map = await Chats.get_messages_map_by_chat_id(id) or history.get('messages') or {} + if any( + message.get('role') == 'assistant' and message.get('done') is False + for message in messages_map.values() + if isinstance(message, dict) + ): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail='Wait for the current response to finish before forking.', + ) + + source_message_id = (form_data.message_id if form_data else None) or chat.current_message_id or history.get('currentId') + if not source_message_id: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='chat has no messages to fork') + + try: + fork_history, fork_messages = build_fork_history(messages_map, source_message_id) + except ValueError as exc: + detail = str(exc) + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND if detail == 'message not found' else status.HTTP_400_BAD_REQUEST, + detail=detail, + ) from exc + + updated_chat = {**(chat.chat or {})} + updated_chat.pop('currentId', None) + updated_chat.update( + { + 'originalChatId': chat.id, + 'branchPointMessageId': source_message_id, + 'title': f'{chat.title} (fork)', + 'history': fork_history, + 'messages': fork_messages, + } + ) + meta = { + **(chat.meta or {}), + 'forked_from': chat.id, + 'forked_from_message_id': source_message_id, + } + + fork = await Chats.insert_new_chat( + str(uuid4()), + user.id, + ChatForm(chat=updated_chat, folder_id=chat.folder_id), + db=db, + internal_meta=meta, + ) + + if not fork: + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=ERROR_MESSAGES.DEFAULT()) + + if chat.pinned: + fork = await Chats.toggle_chat_pinned_by_id(fork.id, db=db) or fork + + await publish_event( + request, + EVENTS.CHAT_CLONED, + actor=user, + subject_id=fork.id, + data={'original_chat_id': id, 'forked_from_message_id': source_message_id}, + ) + return ChatResponse(**fork.model_dump()) + + @router.post('/{id}/clone', response_model=ChatResponse | None) async def clone_chat_by_id( request: Request, diff --git a/backend/open_webui/utils/chat_fork.py b/backend/open_webui/utils/chat_fork.py new file mode 100644 index 0000000000..4963af0609 --- /dev/null +++ b/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 diff --git a/backend/open_webui/utils/context_compaction.py b/backend/open_webui/utils/context_compaction.py index ca0e7a4f9d..910fa5b7b6 100644 --- a/backend/open_webui/utils/context_compaction.py +++ b/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