From 2eb632853c6db4690fa0c64965708176282aedee Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 17 Apr 2026 21:49:33 +0000 Subject: [PATCH] refac: Re-normalize chat tags into chat_tag association table --- backend/open_webui/internal/db.py | 65 +++ .../17a6d37e23d2_add_chat_tag_table.py | 363 ++++++++++++++ backend/open_webui/models/chats.py | 470 ++++++++++++------ backend/open_webui/models/tags.py | 61 ++- backend/open_webui/routers/analytics.py | 31 +- backend/open_webui/routers/chats.py | 79 +-- 6 files changed, 848 insertions(+), 221 deletions(-) create mode 100644 backend/open_webui/migrations/versions/17a6d37e23d2_add_chat_tag_table.py diff --git a/backend/open_webui/internal/db.py b/backend/open_webui/internal/db.py index d75ca11d04..cf53a04be0 100644 --- a/backend/open_webui/internal/db.py +++ b/backend/open_webui/internal/db.py @@ -25,6 +25,7 @@ from open_webui.env import ( ) from peewee_migrate import Router from sqlalchemy import Dialect, create_engine, MetaData, event, types +from sqlalchemy.dialects import postgresql, sqlite from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import scoped_session, sessionmaker, Session @@ -300,3 +301,67 @@ async def get_async_db_context(db: Optional[AsyncSession] = None): else: async with get_async_db() as session: yield session + + +def _insert_for_dialect(dialect_name: str): + if dialect_name == 'postgresql': + return postgresql.insert + if dialect_name == 'sqlite': + return sqlite.insert + raise NotImplementedError( + f'insert_on_conflict_nothing: unsupported dialect {dialect_name!r}; only postgresql and sqlite are supported' + ) + + +# PG caps statements at 65,535 binds; SQLite < 3.32 (May 2020) caps at 999. +_PG_MAX_BIND_PARAMS = 65_000 +_SQLITE_MAX_BIND_PARAMS = 900 + + +def sql_param_batch(dialect_name: str, cols_per_row: int = 1) -> int: + cols_per_row = max(1, cols_per_row) + if dialect_name == 'postgresql': + budget = _PG_MAX_BIND_PARAMS + elif dialect_name == 'sqlite': + budget = _SQLITE_MAX_BIND_PARAMS + else: + raise NotImplementedError( + f'sql_param_batch: unsupported dialect {dialect_name!r}; only postgresql and sqlite are supported' + ) + return max(1, budget // cols_per_row) + + +async def insert_on_conflict_nothing( + db: AsyncSession, + target, # mapped ORM class or sa.Table + values: dict, + index_elements: list[str], +): + """Single-row INSERT ... ON CONFLICT (index_elements) DO NOTHING on + postgresql or sqlite. Caller is responsible for committing.""" + insert = _insert_for_dialect(db.get_bind().dialect.name) + await db.execute( + insert(target).values(**values).on_conflict_do_nothing(index_elements=index_elements) + ) + + +async def insert_all_on_conflict_nothing( + db: AsyncSession, + target, # mapped ORM class or sa.Table + values_list: list[dict], + index_elements: list[str], +): + """Bulk INSERT ... ON CONFLICT (index_elements) DO NOTHING on postgresql + or sqlite. Caller is responsible for committing.""" + if not values_list: + return + dialect_name = db.get_bind().dialect.name + insert = _insert_for_dialect(dialect_name) + # Bind-count budget uses the widest dict. Callers must still pass rows + # with matching keys - SA won't infer a uniform column set otherwise. + batch_size = sql_param_batch(dialect_name, cols_per_row=max(len(v) for v in values_list)) + for start in range(0, len(values_list), batch_size): + batch = values_list[start:start + batch_size] + await db.execute( + insert(target).values(batch).on_conflict_do_nothing(index_elements=index_elements) + ) diff --git a/backend/open_webui/migrations/versions/17a6d37e23d2_add_chat_tag_table.py b/backend/open_webui/migrations/versions/17a6d37e23d2_add_chat_tag_table.py new file mode 100644 index 0000000000..be2db0993b --- /dev/null +++ b/backend/open_webui/migrations/versions/17a6d37e23d2_add_chat_tag_table.py @@ -0,0 +1,363 @@ +"""Add chat_tag table + +Revision ID: 17a6d37e23d2 +Revises: c1d2e3f4a5b6 +Create Date: 2026-04-17 00:00:00.000000 + +""" + +import json +import logging +from itertools import islice +from typing import Iterable, Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql, sqlite + +log = logging.getLogger(__name__) + +revision: str = '17a6d37e23d2' +down_revision: Union[str, None] = 'c1d2e3f4a5b6' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +# Pages are drained with .fetchall() so no server-side cursor stays open +# across pages (large PG deployments OOM'd on yield_per in prior migrations). +CHAT_PAGE_SIZE = 1000 + +# Mirrors open_webui.internal.db budgets - keep in sync. +_PG_MAX_BIND_PARAMS = 65_000 +_SQLITE_MAX_BIND_PARAMS = 900 + + +def _row_batch_size(dialect: str, cols_per_row: int) -> int: + budget = _SQLITE_MAX_BIND_PARAMS if dialect == 'sqlite' else _PG_MAX_BIND_PARAMS + return max(1, budget // max(1, cols_per_row)) + +LOG_EVERY_CHATS = 50_000 + + +def _normalize_tag_id(raw: str) -> str: + # Must stay in sync with open_webui.models.tags.normalize_tag_id. + return raw.replace(' ', '_').lower() + + +def _chunked(source: Iterable, size: int) -> Iterable[list]: + it = iter(source) + while True: + batch = list(islice(it, size)) + if not batch: + return + yield batch + + +def _bulk_insert_on_conflict_nothing(conn, table, rows, index_elements): + """Bulk INSERT ... ON CONFLICT DO NOTHING. PG and SQLite only.""" + if not rows: + return + dialect = conn.dialect.name + batch_size = _row_batch_size(dialect, cols_per_row=len(rows[0])) + for batch in _chunked(rows, batch_size): + if dialect == 'postgresql': + stmt = postgresql.insert(table).values(batch).on_conflict_do_nothing(index_elements=index_elements) + elif dialect == 'sqlite': + stmt = sqlite.insert(table).values(batch).on_conflict_do_nothing(index_elements=index_elements) + else: + raise NotImplementedError( + f'_bulk_insert_on_conflict_nothing: unsupported dialect {dialect!r}; only postgresql and sqlite are supported' + ) + conn.execute(stmt) + + +def upgrade() -> None: + dialect = op.get_bind().dialect.name + if dialect not in ('postgresql', 'sqlite'): + raise NotImplementedError( + f'chat_tag migration: unsupported dialect {dialect!r}; only postgresql and sqlite are supported' + ) + + op.create_table( + 'chat_tag', + sa.Column('chat_id', sa.String(), nullable=False), + sa.Column('tag_id', sa.String(), nullable=False), + sa.Column('user_id', sa.String(), nullable=False), + sa.PrimaryKeyConstraint('chat_id', 'tag_id', 'user_id', name='pk_chat_tag'), + sa.ForeignKeyConstraint( + ['chat_id'], + ['chat.id'], + name='fk_chat_tag_chat_id', + ondelete='CASCADE', + ), + sa.ForeignKeyConstraint( + ['tag_id', 'user_id'], + ['tag.id', 'tag.user_id'], + name='fk_chat_tag_tag', + ondelete='CASCADE', + ), + ) + + conn = op.get_bind() + + chat = sa.table( + 'chat', + sa.column('id', sa.String()), + sa.column('user_id', sa.String()), + sa.column('meta', sa.JSON()), + ) + tag = sa.table( + 'tag', + sa.column('id', sa.String()), + sa.column('name', sa.String()), + sa.column('user_id', sa.String()), + sa.column('meta', sa.JSON()), + ) + chat_tag = sa.table( + 'chat_tag', + sa.column('chat_id', sa.String()), + sa.column('tag_id', sa.String()), + sa.column('user_id', sa.String()), + ) + + last_chat_id: Union[str, None] = None + chats_processed = 0 + chat_tag_rows_submitted = 0 + tag_rows_submitted = 0 + meta_rows_stripped = 0 + next_log_threshold = 0 # log the first page unconditionally + + strip_meta_update = ( + sa.update(chat) + .where(chat.c.id == sa.bindparam('target_chat_id')) + .values(meta=sa.bindparam('new_meta', type_=sa.JSON())) + ) + + while True: + chat_page_query = sa.select(chat.c.id, chat.c.user_id, chat.c.meta).order_by(chat.c.id) + if last_chat_id is not None: + chat_page_query = chat_page_query.where(chat.c.id > last_chat_id) + chat_page_query = chat_page_query.limit(CHAT_PAGE_SIZE) + + chat_rows = conn.execute(chat_page_query).fetchall() + if not chat_rows: + break + + # First raw display name seen per (tag_id, user_id) wins for new rows; + # pre-existing tag rows are never overwritten. + display_name_by_tag_key: dict[tuple[str, str], str] = {} + chat_tag_payload: list[dict] = [] + meta_strip_payload: list[dict] = [] + + for chat_row in chat_rows: + meta = chat_row.meta + if isinstance(meta, str): + try: + meta = json.loads(meta) + except (TypeError, ValueError): + meta = None + if not isinstance(meta, dict): + continue + + # Shared snapshots (user_id='shared-...') historically leaked tags + # via the meta blob; don't promote them to real associations. + is_shared_snapshot = isinstance(chat_row.user_id, str) and chat_row.user_id.startswith('shared-') + + raw_tag_names = meta.get('tags') + if is_shared_snapshot or not isinstance(raw_tag_names, list) or not raw_tag_names: + # Still strip the 'tags' key if present, so meta is + # consistently tag-free post-upgrade. + if 'tags' in meta: + stripped_meta = {k: v for k, v in meta.items() if k != 'tags'} + meta_strip_payload.append( + {'target_chat_id': chat_row.id, 'new_meta': stripped_meta} + ) + continue + + seen_tag_ids_in_chat: set[str] = set() + for raw_tag_name in raw_tag_names: + if not isinstance(raw_tag_name, str): + continue + tag_id = _normalize_tag_id(raw_tag_name) + # 'none' is the search sentinel; don't promote it to a real association. + if not tag_id or tag_id == 'none' or tag_id in seen_tag_ids_in_chat: + continue + seen_tag_ids_in_chat.add(tag_id) + + display_name_by_tag_key.setdefault((tag_id, chat_row.user_id), raw_tag_name) + chat_tag_payload.append( + {'chat_id': chat_row.id, 'tag_id': tag_id, 'user_id': chat_row.user_id} + ) + + stripped_meta = {k: v for k, v in meta.items() if k != 'tags'} + meta_strip_payload.append( + {'target_chat_id': chat_row.id, 'new_meta': stripped_meta} + ) + + if display_name_by_tag_key: + # Per-page reset is safe: the existing-tag filter below never + # overwrites a pre-existing (tag_id, user_id) row. + tag_keys = list(display_name_by_tag_key.keys()) + existing_tag_keys: set[tuple[str, str]] = set() + # tuple_(id, user_id) = 2 binds per row. + key_batch_size = _row_batch_size(dialect, cols_per_row=2) + for key_batch in _chunked(tag_keys, key_batch_size): + existing_tag_query = sa.select(tag.c.id, tag.c.user_id).where( + sa.tuple_(tag.c.id, tag.c.user_id).in_(key_batch) + ) + for existing_row in conn.execute(existing_tag_query).fetchall(): + existing_tag_keys.add((existing_row.id, existing_row.user_id)) + + new_tag_rows = [ + {'id': tid, 'name': raw_name, 'user_id': uid} + for (tid, uid), raw_name in display_name_by_tag_key.items() + if (tid, uid) not in existing_tag_keys + ] + if new_tag_rows: + _bulk_insert_on_conflict_nothing(conn, tag, new_tag_rows, index_elements=['id', 'user_id']) + tag_rows_submitted += len(new_tag_rows) + + if chat_tag_payload: + _bulk_insert_on_conflict_nothing( + conn, chat_tag, chat_tag_payload, + index_elements=['chat_id', 'tag_id', 'user_id'], + ) + chat_tag_rows_submitted += len(chat_tag_payload) + + if meta_strip_payload: + # Paginated; a single full-table UPDATE holds write locks too long. + # (Also: meta is sa.JSON, so PG's json - 'tags' needs a jsonb cast.) + conn.execute(strip_meta_update, meta_strip_payload) + meta_rows_stripped += len(meta_strip_payload) + + last_chat_id = chat_rows[-1].id + chats_processed += len(chat_rows) + + if chats_processed >= next_log_threshold: + log.info( + f'chat_tag backfill progress: {chats_processed} chats processed, ' + f'{chat_tag_rows_submitted} associations submitted, ' + f'{tag_rows_submitted} tags submitted, ' + f'{meta_rows_stripped} meta rows stripped, last_chat_id={last_chat_id}' + ) + next_log_threshold = chats_processed + LOG_EVERY_CHATS + + log.info( + f'chat_tag backfill complete: {chats_processed} chats processed, ' + f'{chat_tag_rows_submitted} associations submitted, {tag_rows_submitted} tags submitted, ' + f'{meta_rows_stripped} meta rows stripped' + ) + + # Index built after backfill so bulk inserts don't pay index-maintenance + # cost per row. + op.create_index('chat_tag_user_tag_idx', 'chat_tag', ['user_id', 'tag_id']) + + +def downgrade() -> None: + # Reserialize chat_tag into meta['tags'] before the drop (post-upgrade + # writes only hit chat_tag). Joins tag to recover tag.name so a + # round-trip upgrade -> downgrade preserves user-visible casing. + # Still lossy on row order (no user-meaningful order survives chat_tag). + conn = op.get_bind() + dialect = conn.dialect.name + if dialect not in ('postgresql', 'sqlite'): + raise NotImplementedError( + f'chat_tag migration: unsupported dialect {dialect!r}; only postgresql and sqlite are supported' + ) + + chat = sa.table( + 'chat', + sa.column('id', sa.String()), + sa.column('user_id', sa.String()), + sa.column('meta', sa.JSON()), + ) + tag = sa.table( + 'tag', + sa.column('id', sa.String()), + sa.column('name', sa.String()), + sa.column('user_id', sa.String()), + ) + chat_tag = sa.table( + 'chat_tag', + sa.column('chat_id', sa.String()), + sa.column('tag_id', sa.String()), + sa.column('user_id', sa.String()), + ) + + last_chat_id: Union[str, None] = None + chats_rewritten = 0 + bulk_update = ( + sa.update(chat) + .where(chat.c.id == sa.bindparam('target_chat_id')) + .values(meta=sa.bindparam('new_meta', type_=sa.JSON())) + ) + + # Paginate by chat.id so a single heavily-tagged chat can never span + # page boundaries (avoids the "truncated tag list" edge case). + while True: + chat_page_query = sa.select(chat.c.id, chat.c.meta).order_by(chat.c.id) + if last_chat_id is not None: + chat_page_query = chat_page_query.where(chat.c.id > last_chat_id) + chat_page_query = chat_page_query.limit(CHAT_PAGE_SIZE) + + page_rows = conn.execute(chat_page_query).fetchall() + if not page_rows: + break + + chat_ids_in_page = [row.id for row in page_rows] + existing_meta_by_chat_id = {row.id: row.meta for row in page_rows} + + # ORDER BY tag.name gives deterministic meta['tags'] ordering across + # downgrade runs (row order from chat_tag is otherwise undefined). + # IN chunked so CHAT_PAGE_SIZE > SQLite's 900-bind cap can't blow up. + tag_rows: list = [] + chat_id_batch_size = _row_batch_size(dialect, cols_per_row=1) + for id_batch in _chunked(chat_ids_in_page, chat_id_batch_size): + tag_rows.extend( + conn.execute( + sa.select(chat_tag.c.chat_id, tag.c.name) + .select_from( + chat_tag.join( + tag, + sa.and_( + chat_tag.c.tag_id == tag.c.id, + chat_tag.c.user_id == tag.c.user_id, + ), + ) + ) + .where(chat_tag.c.chat_id.in_(id_batch)) + .order_by(chat_tag.c.chat_id, tag.c.name) + ).fetchall() + ) + tag_names_by_chat_id: dict[str, list[str]] = {cid: [] for cid in chat_ids_in_page} + for tag_row in tag_rows: + tag_names_by_chat_id[tag_row.chat_id].append(tag_row.name) + + update_params = [] + for chat_id in chat_ids_in_page: + tag_names = tag_names_by_chat_id[chat_id] + existing_meta = existing_meta_by_chat_id.get(chat_id) + if isinstance(existing_meta, str): + try: + existing_meta = json.loads(existing_meta) + except (TypeError, ValueError): + existing_meta = {} + # Lossy for originally-non-dict meta; upgrade skipped those rows. + if not isinstance(existing_meta, dict): + existing_meta = {} + # Skip chats that had no tags pre-upgrade and still have none: + # don't grow their meta with an empty 'tags' key. + if not tag_names and 'tags' not in existing_meta: + continue + merged_meta = {**existing_meta, 'tags': tag_names} + update_params.append({'target_chat_id': chat_id, 'new_meta': merged_meta}) + + if update_params: + conn.execute(bulk_update, update_params) + chats_rewritten += len(update_params) + + last_chat_id = chat_ids_in_page[-1] + + log.info(f'chat_tag downgrade: serialized tags back into meta for {chats_rewritten} chats') + + op.drop_index('chat_tag_user_tag_idx', table_name='chat_tag') + op.drop_table('chat_tag') diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py index ba6611a811..9642683637 100644 --- a/backend/open_webui/models/chats.py +++ b/backend/open_webui/models/chats.py @@ -8,8 +8,15 @@ from sqlalchemy import select, delete, update, func, or_, and_, text from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql import exists from sqlalchemy.sql.expression import bindparam -from open_webui.internal.db import Base, JSONField, get_async_db_context -from open_webui.models.tags import TagModel, Tag, Tags +from open_webui.internal.db import ( + Base, + JSONField, + get_async_db_context, + insert_all_on_conflict_nothing, + insert_on_conflict_nothing, + sql_param_batch, +) +from open_webui.models.tags import TagModel, Tag, Tags, normalize_tag_id, RESERVED_TAG_ID_NONE from open_webui.models.folders import Folders from open_webui.models.chat_messages import ChatMessage, ChatMessages from open_webui.models.automations import AutomationRun @@ -21,6 +28,8 @@ from sqlalchemy import ( Boolean, Column, ForeignKey, + ForeignKeyConstraint, + PrimaryKeyConstraint, String, Text, JSON, @@ -70,6 +79,35 @@ class Chat(Base): ) +# Writer-enforced invariant: chat_tag.user_id == chat(chat_id).user_id. +# user_id is in the PK so the composite FK to tag(id, user_id) is declarable. +class ChatTag(Base): + __tablename__ = 'chat_tag' + + chat_id = Column(String, nullable=False) + tag_id = Column(String, nullable=False) + user_id = Column(String, nullable=False) + + __table_args__ = ( + PrimaryKeyConstraint('chat_id', 'tag_id', 'user_id', name='pk_chat_tag'), + ForeignKeyConstraint( + ['chat_id'], + ['chat.id'], + name='fk_chat_tag_chat_id', + ondelete='CASCADE', + ), + ForeignKeyConstraint( + ['tag_id', 'user_id'], + ['tag.id', 'tag.user_id'], + name='fk_chat_tag_tag', + ondelete='CASCADE', + ), + # No chat_id-only index: PK's leading column covers it. Name must + # match migration 17a6d37e23d2. + Index('chat_tag_user_tag_idx', 'user_id', 'tag_id'), + ) + + class ChatModel(BaseModel): model_config = ConfigDict(from_attributes=True) @@ -356,16 +394,66 @@ class ChatTable: db: Optional[AsyncSession] = None, ) -> list[ChatModel]: async with get_async_db_context(db) as db: - chats = [] + chats: list[Chat] = [] + new_chat_tag_rows: list[dict] = [] + # First display name per normalized tag_id wins - avoids + # composite-PK duplicates inside ensure_tags_exist. + display_name_by_tag_id: dict[str, str] = {} for form_data in chat_import_forms: - chat = self._chat_import_form_to_chat_model(user_id, form_data) - chats.append(Chat(**chat.model_dump())) + chat_model = self._chat_import_form_to_chat_model(user_id, form_data) + # Tags now live in chat_tag, not meta. Entries are validated + # per-item below; the list may contain non-strings from legacy data. + raw_tag_names: list = [] + if isinstance(chat_model.meta, dict) and 'tags' in chat_model.meta: + candidate = chat_model.meta.get('tags') + if isinstance(candidate, list): + raw_tag_names = candidate + chat_model.meta = {k: v for k, v in chat_model.meta.items() if k != 'tags'} + + chat_row = Chat(**chat_model.model_dump()) + chats.append(chat_row) + + seen_tag_ids_in_chat: set[str] = set() + for raw_tag_name in raw_tag_names: + if not isinstance(raw_tag_name, str): + continue + tag_id = normalize_tag_id(raw_tag_name) + # 'none' is the sentinel used by the search "tag:none" + # filter to mean "no tags" - skip it so it never becomes + # a real association. + if not tag_id or tag_id == RESERVED_TAG_ID_NONE or tag_id in seen_tag_ids_in_chat: + continue + seen_tag_ids_in_chat.add(tag_id) + display_name_by_tag_id.setdefault(tag_id, raw_tag_name) + new_chat_tag_rows.append( + {'chat_id': chat_row.id, 'tag_id': tag_id, 'user_id': user_id} + ) + + # One commit covers chats + tag rows + chat_tag rows. Flush the + # chat rows before raw tag/chat_tag inserts so the FKs resolve + # without depending on implicit autoflush ordering. db.add_all(chats) + await db.flush() + if display_name_by_tag_id: + await Tags.ensure_tags_exist( + list(display_name_by_tag_id.values()), user_id, db=db, commit=False + ) + # Flush so tag rows land before chat_tag's composite FK check. + await db.flush() + if new_chat_tag_rows: + # ON CONFLICT DO NOTHING matches the other write paths; the + # per-chat seen set already dedupes within a single import, + # but this stays consistent and idempotent. + await insert_all_on_conflict_nothing( + db, + ChatTag, + new_chat_tag_rows, + index_elements=['chat_id', 'tag_id', 'user_id'], + ) await db.commit() - # Dual-write messages to chat_message table try: for form_data, chat_obj in zip(chat_import_forms, chats): history = form_data.chat.get('history', {}) @@ -428,27 +516,79 @@ class ChatTable: return None async def update_chat_tags_by_id(self, id: str, tags: list[str], user) -> Optional[ChatModel]: + # Lock the chat row so concurrent callers see each other's chat_tag + # writes and the method has true replace semantics (instead of the + # set-union behavior a plain read + diff would produce under races). + # PG honors FOR UPDATE; SQLite is single-writer so the modifier is a + # no-op there. async with get_async_db_context() as db: - chat = await db.get(Chat, id) - if chat is None: + chat = await db.scalar( + select(Chat).where(Chat.id == id).with_for_update() + ) + if chat is None or chat.user_id != user.id: return None - old_tags = chat.meta.get('tags', []) - new_tags = [t for t in tags if t.replace(' ', '_').lower() != 'none'] - new_tag_ids = [t.replace(' ', '_').lower() for t in new_tags] + # First display name per normalized id wins; ['My Tag', 'my tag'] + # would otherwise hit a composite-PK error in ensure_tags_exist. + display_name_by_tag_id: dict[str, str] = {} + for raw_tag_name in tags: + if not isinstance(raw_tag_name, str): + continue + tag_id = normalize_tag_id(raw_tag_name) + if not tag_id or tag_id == RESERVED_TAG_ID_NONE: + continue + display_name_by_tag_id.setdefault(tag_id, raw_tag_name) + new_tag_ids_set = set(display_name_by_tag_id) - # Single meta update - chat.meta = {**chat.meta, 'tags': new_tag_ids} + previous_tag_id_rows = await db.execute( + select(ChatTag.tag_id).filter_by(chat_id=id, user_id=user.id) + ) + previous_tag_ids = {row[0] for row in previous_tag_id_rows.all()} + + to_add = new_tag_ids_set - previous_tag_ids + to_remove = previous_tag_ids - new_tag_ids_set + + # Nothing changed - release the FOR UPDATE lock without committing + # mid-batch work that doesn't exist. + if not to_add and not to_remove: + return ChatModel.model_validate(chat) + + # commit=False keeps the tag ensure + chat_tag diff atomic. + if to_add: + await Tags.ensure_tags_exist( + [display_name_by_tag_id[tag_id] for tag_id in to_add], + user.id, + db=db, + commit=False, + ) + # Bulk ON CONFLICT DO NOTHING so concurrent adds of the same + # (chat, tag, user) don't raise IntegrityError. + await insert_all_on_conflict_nothing( + db, + ChatTag, + [ + {'chat_id': id, 'tag_id': tag_id, 'user_id': user.id} + for tag_id in to_add + ], + index_elements=['chat_id', 'tag_id', 'user_id'], + ) + if to_remove: + await db.execute( + delete(ChatTag).where( + ChatTag.chat_id == id, + ChatTag.user_id == user.id, + ChatTag.tag_id.in_(to_remove), + ) + ) await db.commit() - await db.refresh(chat) - # Batch-create any missing tag rows - await Tags.ensure_tags_exist(new_tags, user.id, db=db) - - # Clean up orphaned old tags in one query - removed = set(old_tags) - set(new_tag_ids) - if removed: - await self.delete_orphan_tags_for_user(list(removed), user.id, db=db) + if to_remove: + # Best-effort - runs after commit, so a failure here would + # only leave orphan tag rows (cosmetic), not corrupt state. + try: + await self.delete_orphan_tags_for_user(list(to_remove), user.id, db=db) + except Exception: + log.exception('orphan tag cleanup failed for chat=%s', id) return ChatModel.model_validate(chat) @@ -555,7 +695,11 @@ class ChatTable: async def insert_shared_chat_by_chat_id( self, chat_id: str, db: Optional[AsyncSession] = None ) -> Optional[ChatModel]: - """Create a shared snapshot for a chat. Returns the original chat with share_id set.""" + """Create a shared snapshot for a chat. Returns the original chat with share_id set. + + chat_tag rows are intentionally NOT copied to the shared scope - + tags are personal organization metadata, not public snapshot state. + """ from open_webui.models.shared_chats import SharedChats async with get_async_db_context(db) as db: @@ -580,7 +724,10 @@ class ChatTable: async def update_shared_chat_by_chat_id( self, chat_id: str, db: Optional[AsyncSession] = None ) -> Optional[ChatModel]: - """Re-snapshot the shared chat with current chat data.""" + """Re-snapshot the shared chat with current chat data. + + Tags don't propagate to the shared scope - see insert_shared_chat_by_chat_id. + """ from open_webui.models.shared_chats import SharedChats try: @@ -1019,7 +1166,7 @@ class ChatTable: # search_text might contain 'tag:tag_name' format so we need to extract the tag_name tag_ids = [ - word.replace('tag:', '').replace(' ', '_').lower() for word in search_text_words if word.startswith('tag:') + normalize_tag_id(word.removeprefix('tag:')) for word in search_text_words if word.startswith('tag:') ] # Extract folder names @@ -1102,32 +1249,6 @@ class ChatTable: ) ) - # Check if there are any tags to filter - if 'none' in tag_ids: - stmt = stmt.filter( - text(""" - NOT EXISTS ( - SELECT 1 - FROM json_each(Chat.meta, '$.tags') AS tag - ) - """) - ) - elif tag_ids: - stmt = stmt.filter( - and_( - *[ - text(f""" - EXISTS ( - SELECT 1 - FROM json_each(Chat.meta, '$.tags') AS tag - WHERE tag.value = :tag_id_{tag_idx} - ) - """).params(**{f'tag_id_{tag_idx}': tag_id}) - for tag_idx, tag_id in enumerate(tag_ids) - ] - ) - ) - elif dialect_name == 'postgresql': # Safety filter: JSON field must not contain \u0000 stmt = stmt.filter(text("Chat.chat::text NOT LIKE '%\\\\u0000%'")) @@ -1153,33 +1274,26 @@ class ChatTable: ) ).params(title_key=f'%{search_text}%', content_key=search_text.lower()) - if 'none' in tag_ids: - stmt = stmt.filter( - text(""" - NOT EXISTS ( - SELECT 1 - FROM json_array_elements_text(Chat.meta->'tags') AS tag - ) - """) - ) - elif tag_ids: - stmt = stmt.filter( - and_( - *[ - text(f""" - EXISTS ( - SELECT 1 - FROM json_array_elements_text(Chat.meta->'tags') AS tag - WHERE tag = :tag_id_{tag_idx} - ) - """).params(**{f'tag_id_{tag_idx}': tag_id}) - for tag_idx, tag_id in enumerate(tag_ids) - ] - ) - ) else: raise NotImplementedError(f'Unsupported dialect: {dialect_name}') + # 'tag:none' = no associations; 'tag:X tag:Y' = has both. + # ChatTag.user_id filter is defense-in-depth + index alignment. + if RESERVED_TAG_ID_NONE in tag_ids: + any_tag_subquery = select(ChatTag.chat_id).where( + ChatTag.chat_id == Chat.id, + ChatTag.user_id == user_id, + ) + stmt = stmt.filter(~exists(any_tag_subquery)) + elif tag_ids: + for required_tag_id in tag_ids: + required_tag_subquery = select(ChatTag.chat_id).where( + ChatTag.chat_id == Chat.id, + ChatTag.user_id == user_id, + ChatTag.tag_id == required_tag_id, + ) + stmt = stmt.filter(exists(required_tag_subquery)) + # Perform pagination at the SQL level stmt = stmt.offset(skip).limit(limit) result = await db.execute(stmt) @@ -1258,14 +1372,42 @@ class ChatTable: except Exception: return None + # Returns the tag_ids associated with (id, user_id) in chat_tag. Empty + # list if the chat doesn't exist OR the user has no chat_tag rows for it. + async def get_chat_tag_ids_by_id_and_user_id( + self, id: str, user_id: str, db: Optional[AsyncSession] = None + ) -> list[str]: + async with get_async_db_context(db) as db: + rows = await db.execute(select(ChatTag.tag_id).filter_by(chat_id=id, user_id=user_id)) + return [row[0] for row in rows.all()] + + async def get_chat_tag_ids_by_chat_ids_and_user_id( + self, chat_ids: list[str], user_id: str, db: Optional[AsyncSession] = None + ) -> dict[str, list[str]]: + if not chat_ids: + return {} + # Every chat_id is present in the result; callers can look up without .get(). + tag_ids_by_chat_id: dict[str, list[str]] = {chat_id: [] for chat_id in chat_ids} + # Chunk the IN predicate to respect each dialect's bind-param ceiling. + async with get_async_db_context(db) as db: + batch_size = sql_param_batch(db.get_bind().dialect.name) + for start in range(0, len(chat_ids), batch_size): + batch_ids = chat_ids[start:start + batch_size] + rows = await db.execute( + select(ChatTag.chat_id, ChatTag.tag_id).where( + ChatTag.chat_id.in_(batch_ids), + ChatTag.user_id == user_id, + ) + ) + for chat_id, tag_id in rows.all(): + tag_ids_by_chat_id[chat_id].append(tag_id) + return tag_ids_by_chat_id + async def get_chat_tags_by_id_and_user_id( self, id: str, user_id: str, db: Optional[AsyncSession] = None ) -> list[TagModel]: async with get_async_db_context(db) as db: - stmt = select(Chat.meta).where(Chat.id == id) - result = await db.execute(stmt) - meta = result.scalar_one_or_none() - tag_ids = (meta or {}).get('tags', []) + tag_ids = await self.get_chat_tag_ids_by_id_and_user_id(id, user_id, db=db) return await Tags.get_tags_by_ids_and_user_id(tag_ids, user_id, db=db) async def get_chat_list_by_user_id_and_tag_name( @@ -1277,34 +1419,23 @@ class ChatTable: db: Optional[AsyncSession] = None, ) -> list[ChatTitleIdResponse]: async with get_async_db_context(db) as db: - stmt = select(Chat.id, Chat.title, Chat.updated_at, Chat.created_at, Chat.last_read_at).filter_by( - user_id=user_id + tag_id = normalize_tag_id(tag_name) + + # ChatTag.user_id filter is defense-in-depth + index alignment + # with chat_tag_user_tag_idx. + chat_list_query = ( + select(Chat.id, Chat.title, Chat.updated_at, Chat.created_at, Chat.last_read_at) + .join(ChatTag, ChatTag.chat_id == Chat.id) + .where(Chat.user_id == user_id, ChatTag.user_id == user_id, ChatTag.tag_id == tag_id) + .order_by(Chat.updated_at.desc(), Chat.id) ) - tag_id = tag_name.replace(' ', '_').lower() - - bind = await db.connection() - dialect_name = bind.dialect.name - log.info(f'DB dialect name: {dialect_name}') - if dialect_name == 'sqlite': - stmt = stmt.filter( - text(f"EXISTS (SELECT 1 FROM json_each(Chat.meta, '$.tags') WHERE json_each.value = :tag_id)") - ).params(tag_id=tag_id) - elif dialect_name == 'postgresql': - stmt = stmt.filter( - text("EXISTS (SELECT 1 FROM json_array_elements_text(Chat.meta->'tags') elem WHERE elem = :tag_id)") - ).params(tag_id=tag_id) - else: - raise NotImplementedError(f'Unsupported dialect: {dialect_name}') - - stmt = stmt.order_by(Chat.updated_at.desc(), Chat.id) if skip: - stmt = stmt.offset(skip) + chat_list_query = chat_list_query.offset(skip) if limit: - stmt = stmt.limit(limit) + chat_list_query = chat_list_query.limit(limit) - result = await db.execute(stmt) - all_chats = result.all() + all_chats = (await db.execute(chat_list_query)).all() return [ ChatTitleIdResponse.model_validate( { @@ -1321,44 +1452,57 @@ class ChatTable: async def add_chat_tag_by_id_and_user_id_and_tag_name( self, id: str, user_id: str, tag_name: str, db: Optional[AsyncSession] = None ) -> Optional[ChatModel]: - tag_id = tag_name.replace(' ', '_').lower() - await Tags.ensure_tags_exist([tag_name], user_id, db=db) + tag_id = normalize_tag_id(tag_name) + # 'none' is the search sentinel; '' would bind a garbage association. + if not tag_id or tag_id == RESERVED_TAG_ID_NONE: + return None try: async with get_async_db_context(db) as db: - chat = await db.get(Chat, id) - if tag_id not in chat.meta.get('tags', []): - chat.meta = { - **chat.meta, - 'tags': list(set(chat.meta.get('tags', []) + [tag_id])), - } + # FOR UPDATE on the chat row to match update_chat_tags_by_id's + # lock, so a concurrent "set tags" + "add tag" pair on the same + # chat serializes rather than racing. + chat = await db.scalar( + select(Chat).where(Chat.id == id).with_for_update() + ) + # Ownership check enforces the chat_tag invariant. + if chat is None or chat.user_id != user_id: + return None + + await Tags.ensure_tags_exist([tag_name], user_id, db=db, commit=False) + + # ON CONFLICT DO NOTHING avoids a TOCTOU race on concurrent adds. + await insert_on_conflict_nothing( + db, + ChatTag, + {'chat_id': id, 'tag_id': tag_id, 'user_id': user_id}, + index_elements=['chat_id', 'tag_id', 'user_id'], + ) + await db.commit() - await db.refresh(chat) return ChatModel.model_validate(chat) except Exception: + log.exception('add_chat_tag failed for chat=%s tag=%s', id, tag_name) return None + # UI-facing count: excludes archived chats, so a tag with only archived + # references shows count=0 but isn't treated as orphan by + # delete_orphan_tags_for_user (which intentionally counts archived too). async def count_chats_by_tag_name_and_user_id( self, tag_name: str, user_id: str, db: Optional[AsyncSession] = None ) -> int: async with get_async_db_context(db) as db: - stmt = select(func.count(Chat.id)).filter_by(user_id=user_id, archived=False) - tag_id = tag_name.replace(' ', '_').lower() - - bind = await db.connection() - dialect_name = bind.dialect.name - if dialect_name == 'sqlite': - stmt = stmt.filter( - text("EXISTS (SELECT 1 FROM json_each(Chat.meta, '$.tags') WHERE json_each.value = :tag_id)") - ).params(tag_id=tag_id) - elif dialect_name == 'postgresql': - stmt = stmt.filter( - text("EXISTS (SELECT 1 FROM json_array_elements_text(Chat.meta->'tags') elem WHERE elem = :tag_id)") - ).params(tag_id=tag_id) - else: - raise NotImplementedError(f'Unsupported dialect: {dialect_name}') - - result = await db.execute(stmt) - return result.scalar() + tag_id = normalize_tag_id(tag_name) + chat_count_query = ( + select(func.count(Chat.id)) + .join(ChatTag, ChatTag.chat_id == Chat.id) + .where( + Chat.user_id == user_id, + ChatTag.user_id == user_id, + ChatTag.tag_id == tag_id, + Chat.archived.is_(False), + ) + ) + return (await db.execute(chat_count_query)).scalar() async def delete_orphan_tags_for_user( self, @@ -1367,23 +1511,37 @@ class ChatTable: threshold: int = 0, db: Optional[AsyncSession] = None, ) -> None: - """Delete tag rows from *tag_ids* that appear in at most *threshold* - non-archived chats for *user_id*. One query to find orphans, one to - delete them. + """Delete tag rows from *tag_ids* whose chat reference count for + *user_id* is at most *threshold*. Counts across archived and + non-archived chats; see the in-body comment for why. - Use threshold=0 after a tag is already removed from a chat's meta. - Use threshold=1 when the chat itself is about to be deleted (the - referencing chat still exists at query time). + threshold=0: call after a tag has already been removed from a chat. + threshold=1: call before the referencing chat is deleted. """ if not tag_ids: return async with get_async_db_context(db) as db: - orphans = [] - for tag_id in tag_ids: - count = await self.count_chats_by_tag_name_and_user_id(tag_id, user_id, db=db) - if count <= threshold: - orphans.append(tag_id) - await Tags.delete_tags_by_ids_and_user_id(orphans, user_id, db=db) + # Counts across archived + non-archived: scoping to non-archived + # would combine with the chat_tag FK CASCADE to destroy archived + # chats' associations the next time any chat drops the tag. + # Chunk the IN for consistency with other IN predicates in this PR. + reference_count_by_tag_id: dict[str, int] = {} + batch_size = sql_param_batch(db.get_bind().dialect.name, cols_per_row=1) + for start in range(0, len(tag_ids), batch_size): + batch = tag_ids[start:start + batch_size] + reference_count_query = ( + select(ChatTag.tag_id, func.count()) + .where(ChatTag.user_id == user_id, ChatTag.tag_id.in_(batch)) + .group_by(ChatTag.tag_id) + ) + for tag_id, count in (await db.execute(reference_count_query)).all(): + reference_count_by_tag_id[tag_id] = count + orphan_tag_ids = [ + tag_id + for tag_id in tag_ids + if reference_count_by_tag_id.get(tag_id, 0) <= threshold + ] + await Tags.delete_tags_by_ids_and_user_id(orphan_tag_ids, user_id, db=db) async def count_chats_by_folder_id_and_user_id( self, folder_id: str, user_id: str, db: Optional[AsyncSession] = None @@ -1395,44 +1553,42 @@ class ChatTable: log.info(f"Count of chats for folder '{folder_id}': {count}") return count + # Callers that care about cleaning up now-unreferenced tag rows invoke + # delete_orphan_tags_for_user themselves (e.g. the DELETE /tags/all route). async def delete_tag_by_id_and_user_id_and_tag_name( self, id: str, user_id: str, tag_name: str, db: Optional[AsyncSession] = None ) -> bool: try: async with get_async_db_context(db) as db: - chat = await db.get(Chat, id) - tags = chat.meta.get('tags', []) - tag_id = tag_name.replace(' ', '_').lower() - - tags = [tag for tag in tags if tag != tag_id] - chat.meta = { - **chat.meta, - 'tags': list(set(tags)), - } + tag_id = normalize_tag_id(tag_name) + await db.execute( + delete(ChatTag).filter_by(chat_id=id, tag_id=tag_id, user_id=user_id) + ) await db.commit() return True except Exception: + log.exception('delete_tag failed for chat=%s tag=%s', id, tag_name) return False async def delete_all_tags_by_id_and_user_id(self, id: str, user_id: str, db: Optional[AsyncSession] = None) -> bool: try: async with get_async_db_context(db) as db: - chat = await db.get(Chat, id) - chat.meta = { - **chat.meta, - 'tags': [], - } + await db.execute(delete(ChatTag).filter_by(chat_id=id, user_id=user_id)) await db.commit() - return True except Exception: + log.exception('delete_all_tags failed for chat=%s', id) return False + # NOTE: ChatMessage / ChatTag are deleted explicitly - SQLite only + # enforces FK cascades when `PRAGMA foreign_keys = ON` is set, which + # isn't guaranteed at runtime. Keep sibling delete_* methods in sync. async def delete_chat_by_id(self, id: str, db: Optional[AsyncSession] = None) -> bool: try: async with get_async_db_context(db) as db: await db.execute(update(AutomationRun).filter_by(chat_id=id).values(chat_id=None)) await db.execute(delete(ChatMessage).filter_by(chat_id=id)) + await db.execute(delete(ChatTag).filter_by(chat_id=id)) await db.execute(delete(Chat).filter_by(id=id)) await db.commit() @@ -1445,6 +1601,7 @@ class ChatTable: async with get_async_db_context(db) as db: await db.execute(update(AutomationRun).filter_by(chat_id=id).values(chat_id=None)) await db.execute(delete(ChatMessage).filter_by(chat_id=id)) + await db.execute(delete(ChatTag).filter_by(chat_id=id)) await db.execute(delete(Chat).filter_by(id=id, user_id=user_id)) await db.commit() @@ -1466,6 +1623,8 @@ class ChatTable: await db.execute( delete(ChatMessage).filter(ChatMessage.chat_id.in_(select(Chat.id).filter_by(user_id=user_id))) ) + # Relies on chat_tag.user_id == chat.user_id invariant. + await db.execute(delete(ChatTag).filter_by(user_id=user_id)) await db.execute(delete(Chat).filter_by(user_id=user_id)) await db.commit() @@ -1483,6 +1642,7 @@ class ChatTable: update(AutomationRun).filter(AutomationRun.chat_id.in_(chat_ids_stmt)).values(chat_id=None) ) await db.execute(delete(ChatMessage).filter(ChatMessage.chat_id.in_(chat_ids_stmt))) + await db.execute(delete(ChatTag).filter(ChatTag.chat_id.in_(chat_ids_stmt))) await db.execute(delete(Chat).filter_by(user_id=user_id, folder_id=folder_id)) await db.commit() diff --git a/backend/open_webui/models/tags.py b/backend/open_webui/models/tags.py index ee2baefc01..9f5b293e1a 100644 --- a/backend/open_webui/models/tags.py +++ b/backend/open_webui/models/tags.py @@ -5,7 +5,7 @@ from typing import Optional from sqlalchemy import select, delete from sqlalchemy.ext.asyncio import AsyncSession -from open_webui.internal.db import Base, JSONField, get_async_db_context +from open_webui.internal.db import Base, JSONField, get_async_db_context, insert_all_on_conflict_nothing from pydantic import BaseModel, ConfigDict @@ -19,6 +19,18 @@ log = logging.getLogger(__name__) # To name a thing is to claim it. The creator has # already named everything stored in this table. #################### +# Reserved tag_id used as a sentinel by the tag:none search filter +# ("chat has no tags"). Writers must reject it so it can never become a +# real association. +RESERVED_TAG_ID_NONE = 'none' + + +def normalize_tag_id(raw: str) -> str: + """Canonical tag_id form. This is the PK for tag and chat_tag, so every + call site that derives an id from a user-supplied name must use this.""" + return raw.replace(' ', '_').lower() + + class Tag(Base): __tablename__ = 'tag' id = Column(String) @@ -31,9 +43,6 @@ class Tag(Base): Index('user_id_idx', 'user_id'), ) - # Unique constraint ensuring (id, user_id) is unique, not just the `id` column - __table_args__ = (PrimaryKeyConstraint('id', 'user_id', name='pk_id_user_id'),) - class TagModel(BaseModel): id: str @@ -56,7 +65,7 @@ class TagChatIdForm(BaseModel): class TagTable: async def insert_new_tag(self, name: str, user_id: str, db: Optional[AsyncSession] = None) -> Optional[TagModel]: async with get_async_db_context(db) as db: - id = name.replace(' ', '_').lower() + id = normalize_tag_id(name) tag = TagModel(**{'id': id, 'user_id': user_id, 'name': name}) try: result = Tag(**tag.model_dump()) @@ -75,7 +84,7 @@ class TagTable: self, name: str, user_id: str, db: Optional[AsyncSession] = None ) -> Optional[TagModel]: try: - id = name.replace(' ', '_').lower() + id = normalize_tag_id(name) async with get_async_db_context(db) as db: result = await db.execute(select(Tag).filter_by(id=id, user_id=user_id)) tag = result.scalars().first() @@ -98,7 +107,7 @@ class TagTable: async def delete_tag_by_name_and_user_id(self, name: str, user_id: str, db: Optional[AsyncSession] = None) -> bool: try: async with get_async_db_context(db) as db: - id = name.replace(' ', '_').lower() + id = normalize_tag_id(name) result = await db.execute(delete(Tag).filter_by(id=id, user_id=user_id)) log.debug(f'res: {result.rowcount}') await db.commit() @@ -122,19 +131,37 @@ class TagTable: log.error(f'delete_tags_by_ids: {e}') return False - async def ensure_tags_exist(self, names: list[str], user_id: str, db: Optional[AsyncSession] = None) -> None: - """Create tag rows for any *names* that don't already exist for *user_id*.""" + async def ensure_tags_exist( + self, + names: list[str], + user_id: str, + db: Optional[AsyncSession] = None, + commit: bool = True, + ) -> None: + """Create tag rows for any *names* that don't already exist for *user_id*. + + Pass ``commit=False`` when the caller owns a larger transaction that + must remain atomic (e.g. a dual-write that also touches chat_tag); + the caller is then responsible for committing the session. + """ + if not commit and db is None: + raise ValueError('ensure_tags_exist(commit=False) requires an explicit db session') if not names: return - ids = [n.replace(' ', '_').lower() for n in names] + # Dedupe on normalized id, first display name wins. ON CONFLICT DO + # NOTHING handles the concurrent-insert race so we don't need the + # old SELECT-then-add check. + values_by_tag_id: dict[str, dict] = {} + for name in names: + tag_id = normalize_tag_id(name) + values_by_tag_id.setdefault( + tag_id, {'id': tag_id, 'name': name, 'user_id': user_id} + ) async with get_async_db_context(db) as db: - result = await db.execute(select(Tag.id).filter(Tag.id.in_(ids), Tag.user_id == user_id)) - existing = {row[0] for row in result.all()} - new_tags = [ - Tag(id=tag_id, name=name, user_id=user_id) for tag_id, name in zip(ids, names) if tag_id not in existing - ] - if new_tags: - db.add_all(new_tags) + await insert_all_on_conflict_nothing( + db, Tag, list(values_by_tag_id.values()), index_elements=['id', 'user_id'] + ) + if commit: await db.commit() diff --git a/backend/open_webui/routers/analytics.py b/backend/open_webui/routers/analytics.py index fd045f79e7..916a6e90fd 100644 --- a/backend/open_webui/routers/analytics.py +++ b/backend/open_webui/routers/analytics.py @@ -5,13 +5,16 @@ import logging from fastapi import APIRouter, Depends, Query from pydantic import BaseModel +from collections import Counter + from open_webui.models.chat_messages import ChatMessages, ChatMessageModel -from open_webui.models.chats import Chats +from open_webui.models.chats import Chats, ChatTag from open_webui.models.groups import Groups from open_webui.models.users import Users from open_webui.models.feedbacks import Feedbacks from open_webui.utils.auth import get_admin_user -from open_webui.internal.db import get_async_session +from open_webui.internal.db import get_async_session, sql_param_batch +from sqlalchemy import select, func from sqlalchemy.ext.asyncio import AsyncSession log = logging.getLogger(__name__) @@ -428,15 +431,19 @@ async def get_model_overview( ) current += timedelta(days=1) - # Get chat tags - tag_counts: dict[str, int] = defaultdict(int) - for chat_id in chat_ids: - chat = await Chats.get_chat_by_id(chat_id, db=db) - if chat and chat.meta: - for tag in chat.meta.get('tags', []): - tag_counts[tag] += 1 - - # Sort by count and take top 10 - tags = [TagEntry(tag=tag, count=count) for tag, count in sorted(tag_counts.items(), key=lambda x: -x[1])[:10]] + # Chunk the IN clause to respect each dialect's bind-param ceiling, + # then aggregate in Python. + tag_counts: Counter[str] = Counter() + batch_size = sql_param_batch(db.get_bind().dialect.name) + for start in range(0, len(chat_ids), batch_size): + batch = chat_ids[start:start + batch_size] + rows = (await db.execute( + select(ChatTag.tag_id, func.count()) + .where(ChatTag.chat_id.in_(batch)) + .group_by(ChatTag.tag_id) + )).all() + for tag_id, count in rows: + tag_counts[tag_id] += count + tags = [TagEntry(tag=t, count=c) for t, c in tag_counts.most_common(10)] return ModelOverviewResponse(history=history, tags=tags) diff --git a/backend/open_webui/routers/chats.py b/backend/open_webui/routers/chats.py index 02f9a74662..0b82c0778f 100644 --- a/backend/open_webui/routers/chats.py +++ b/backend/open_webui/routers/chats.py @@ -25,7 +25,7 @@ from open_webui.models.chats import ( ) from open_webui.models.shared_chats import SharedChats, SharedChatResponse from open_webui.models.access_grants import AccessGrants -from open_webui.models.tags import TagModel, Tags +from open_webui.models.tags import TagModel, Tags, normalize_tag_id, RESERVED_TAG_ID_NONE from open_webui.models.folders import Folders from open_webui.internal.db import get_async_session @@ -105,6 +105,10 @@ async def get_session_user_chat_usage_stats( chats = result.items total = result.total + tag_ids_by_chat_id = await Chats.get_chat_tag_ids_by_chat_ids_and_user_id( + [chat.id for chat in chats], user.id, db=db + ) + chat_stats = [] for chat in chats: messages_map = chat.chat.get('history', {}).get('messages', {}) @@ -178,7 +182,7 @@ async def get_session_user_chat_usage_stats( 'average_response_time': average_response_time, 'average_user_message_content_length': average_user_message_content_length, 'average_assistant_message_content_length': average_assistant_message_content_length, - 'tags': chat.meta.get('tags', []), + 'tags': tag_ids_by_chat_id[chat.id], 'last_message_at': message_list[-1].get('timestamp', None), 'updated_at': chat.updated_at, 'created_at': chat.created_at, @@ -209,7 +213,7 @@ class ChatStatsExportList(BaseModel): page: int -def _process_chat_for_export(chat) -> Optional[ChatStatsExport]: +def _process_chat_for_export(chat, tag_ids: list[str]) -> Optional[ChatStatsExport]: try: def get_message_content_length(message): @@ -325,7 +329,7 @@ def _process_chat_for_export(chat) -> Optional[ChatStatsExport]: user_id=chat.user_id, created_at=chat.created_at, updated_at=chat.updated_at, - tags=chat.meta.get('tags', []), + tags=tag_ids, stats=stats, chat=chat_body, ) @@ -345,9 +349,12 @@ async def calculate_chat_stats(user_id, skip=0, limit=10, filter=None): filter=filter, ) + tag_ids_by_chat_id = await Chats.get_chat_tag_ids_by_chat_ids_and_user_id( + [c.id for c in result.items], user_id + ) chat_stats_export_list = [] for chat in result.items: - chat_stat = _process_chat_for_export(chat) + chat_stat = _process_chat_for_export(chat, tag_ids_by_chat_id.get(chat.id, [])) if chat_stat: chat_stats_export_list.append(chat_stat) @@ -380,9 +387,12 @@ async def generate_chat_stats_jsonl_generator(user_id, filter): if not result.items: break + tag_ids_by_chat_id = await Chats.get_chat_tag_ids_by_chat_ids_and_user_id( + [c.id for c in result.items], user_id + ) for chat in result.items: try: - chat_stat = _process_chat_for_export(chat) + chat_stat = _process_chat_for_export(chat, tag_ids_by_chat_id.get(chat.id, [])) if chat_stat: yield chat_stat.model_dump_json() + '\n' except Exception as e: @@ -471,8 +481,9 @@ async def export_single_chat_stats( detail=ERROR_MESSAGES.ACCESS_PROHIBITED, ) + tag_ids = await Chats.get_chat_tag_ids_by_id_and_user_id(chat.id, chat.user_id, db=db) # Process the chat for export (pure computation, no DB) - chat_stats = _process_chat_for_export(chat) + chat_stats = _process_chat_for_export(chat, tag_ids) if not chat_stats: raise HTTPException( @@ -1067,7 +1078,9 @@ async def delete_chat_by_id( status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND, ) - await Chats.delete_orphan_tags_for_user(chat.meta.get('tags', []), user.id, threshold=1, db=db) + # Orphan cleanup is scoped to the chat's owner, not the admin. + tag_ids = await Chats.get_chat_tag_ids_by_id_and_user_id(id, chat.user_id, db=db) + await Chats.delete_orphan_tags_for_user(tag_ids, chat.user_id, threshold=1, db=db) result = await Chats.delete_chat_by_id(id, db=db) @@ -1085,7 +1098,8 @@ async def delete_chat_by_id( status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND, ) - await Chats.delete_orphan_tags_for_user(chat.meta.get('tags', []), user.id, threshold=1, db=db) + tag_ids = await Chats.get_chat_tag_ids_by_id_and_user_id(id, user.id, db=db) + await Chats.delete_orphan_tags_for_user(tag_ids, user.id, threshold=1, db=db) result = await Chats.delete_chat_by_id_and_user_id(id, user.id, db=db) return result @@ -1251,16 +1265,8 @@ async def clone_shared_chat_by_id( async def archive_chat_by_id(id: str, user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)): chat = await Chats.get_chat_by_id_and_user_id(id, user.id, db=db) if chat: + # chat_tag rows persist; list/count queries already filter archived=False. chat = await Chats.toggle_chat_archive_by_id(id, db=db) - - tag_ids = chat.meta.get('tags', []) - if chat.archived: - # Archived chats are excluded from count — clean up orphans - await Chats.delete_orphan_tags_for_user(tag_ids, user.id, db=db) - else: - # Unarchived — ensure tag rows exist - await Tags.ensure_tags_exist(tag_ids, user.id, db=db) - return ChatResponse(**chat.model_dump()) else: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()) @@ -1460,8 +1466,7 @@ async def update_chat_folder_id_by_id( async def get_chat_tags_by_id(id: str, user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)): chat = await Chats.get_chat_by_id_and_user_id(id, user.id, db=db) if chat: - tags = chat.meta.get('tags', []) - return await Tags.get_tags_by_ids_and_user_id(tags, user.id, db=db) + return await Chats.get_chat_tags_by_id_and_user_id(id, user.id, db=db) else: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND) @@ -1480,21 +1485,19 @@ async def add_tag_by_id_and_tag_name( ): chat = await Chats.get_chat_by_id_and_user_id(id, user.id, db=db) if chat: - tags = chat.meta.get('tags', []) - tag_id = form_data.name.replace(' ', '_').lower() - - if tag_id == 'none': + if normalize_tag_id(form_data.name) == RESERVED_TAG_ID_NONE: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT("Tag name cannot be 'None'"), ) - if tag_id not in tags: - await Chats.add_chat_tag_by_id_and_user_id_and_tag_name(id, user.id, form_data.name, db=db) - - chat = await Chats.get_chat_by_id_and_user_id(id, user.id, db=db) - tags = chat.meta.get('tags', []) - return await Tags.get_tags_by_ids_and_user_id(tags, user.id, db=db) + # None => chat disappeared (or ownership broke) between the check + # above and the FOR UPDATE lock; surface as 404 rather than a + # misleading empty tag list. + result = await Chats.add_chat_tag_by_id_and_user_id_and_tag_name(id, user.id, form_data.name, db=db) + if result is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND) + return await Chats.get_chat_tags_by_id_and_user_id(id, user.id, db=db) else: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT()) @@ -1515,12 +1518,13 @@ async def delete_tag_by_id_and_tag_name( if chat: await Chats.delete_tag_by_id_and_user_id_and_tag_name(id, user.id, form_data.name, db=db) - if await Chats.count_chats_by_tag_name_and_user_id(form_data.name, user.id, db=db) == 0: - await Tags.delete_tag_by_name_and_user_id(form_data.name, user.id, db=db) + # Orphan cleanup counts archived chats too (see delete_orphan_tags_for_user), + # so a tag referenced only by archived chats is preserved for unarchive. + await Chats.delete_orphan_tags_for_user( + [normalize_tag_id(form_data.name)], user.id, db=db + ) - chat = await Chats.get_chat_by_id_and_user_id(id, user.id, db=db) - tags = chat.meta.get('tags', []) - return await Tags.get_tags_by_ids_and_user_id(tags, user.id, db=db) + return await Chats.get_chat_tags_by_id_and_user_id(id, user.id, db=db) else: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND) @@ -1536,9 +1540,10 @@ async def delete_all_tags_by_id( ): chat = await Chats.get_chat_by_id_and_user_id(id, user.id, db=db) if chat: - old_tags = chat.meta.get('tags', []) + # Snapshot before clearing so orphan cleanup knows which tags to re-count. + old_tag_ids = await Chats.get_chat_tag_ids_by_id_and_user_id(id, user.id, db=db) await Chats.delete_all_tags_by_id_and_user_id(id, user.id, db=db) - await Chats.delete_orphan_tags_for_user(old_tags, user.id, db=db) + await Chats.delete_orphan_tags_for_user(old_tag_ids, user.id, db=db) return True else: