bhetherman/open-webui
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-05-21 14:01:57 +04:00
parent cac4c6da2e
commit 260ead64da
19 changed files with 772 additions and 771 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/models/chats.py
+180 -178
View File
@@ -1,10 +1,11 @@
"""Chat models, forms, and database operations."""
from __future__ import annotations
import json
import logging
import time
import uuid
from typing import Optional
from open_webui.internal.db import Base, JSONField, get_async_db_context
from open_webui.models.automations import AutomationRun
@@ -35,12 +36,6 @@ from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql import exists
from sqlalchemy.sql.expression import bindparam
####################
# Chat DB Schema
# Let no word spoken in this house be lost, and when the
# record is read again, let it still serve the one who spoke.
####################
log = logging.getLogger(__name__)
@@ -49,14 +44,14 @@ class Chat(Base):
id = Column(String, primary_key=True, unique=True)
user_id = Column(String)
title = Column(Text)
title = Column(Text) # user-visible conversation title
chat = Column(JSON)
created_at = Column(BigInteger)
updated_at = Column(BigInteger)
share_id = Column(Text, unique=True, nullable=True)
archived = Column(Boolean, default=False)
share_id = Column(Text, unique=True, nullable=True) # public share link token
archived = Column(Boolean, default=False) # hidden from main chat list
pinned = Column(Boolean, default=False, nullable=True)
meta = Column(JSON, server_default='{}')
@@ -302,7 +297,7 @@ class ChatTable:
async def insert_new_chat(
self, id: str, user_id: str, form_data: ChatForm, db: AsyncSession | None = None
) -> ChatModel | None:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
chat = ChatModel(
**{
'id': id,
@@ -318,9 +313,9 @@ class ChatTable:
)
chat_item = Chat(**chat.model_dump())
db.add(chat_item)
await db.commit()
await db.refresh(chat_item)
session.add(chat_item)
await session.commit()
await session.refresh(chat_item)
# Dual-write initial messages to chat_message table
try:
@@ -362,7 +357,7 @@ class ChatTable:
chat_import_forms: list[ChatImportForm],
db: AsyncSession | None = None,
) -> list[ChatModel]:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
chats = []
for form_data in chat_import_forms:
@@ -370,7 +365,7 @@ class ChatTable:
chats.append(Chat(**chat.model_dump()))
db.add_all(chats)
await db.commit()
await session.commit()
# Dual-write messages to chat_message table
for form_data, chat_obj in zip(chat_import_forms, chats):
@@ -390,10 +385,13 @@ class ChatTable:
return [ChatModel.model_validate(chat) for chat in chats]
async def update_chat_by_id(self, id: str, chat: dict, db: AsyncSession | None = None) -> ChatModel | None:
try:
async with get_async_db_context(db) as db:
chat_item = await db.get(Chat, id)
async def update_chat_by_id(
self, id: str, chat: dict, db: AsyncSession | None = None,
) -> ChatModel | None:
"""Persist updated chat content, sanitizing null bytes."""
try: # load the chat record for in-place mutation
async with get_async_db_context(db) as session:
chat_item = await session.get(Chat, id)
if chat_item is None:
return None
@@ -402,19 +400,19 @@ class ChatTable:
chat_item.updated_at = int(time.time())
await db.commit()
await session.commit()
return ChatModel.model_validate(chat_item)
except Exception:
return None
return
async def update_chat_last_read_at_by_id(self, id: str, user_id: str, db: AsyncSession | None = None) -> bool:
try:
async with get_async_db_context(db) as db:
chat = await db.get(Chat, id)
async with get_async_db_context(db) as session:
chat = await session.get(Chat, id)
if chat and chat.user_id == user_id:
chat.last_read_at = int(time.time())
await db.commit()
await session.commit()
return True
return False
except Exception:
@@ -423,22 +421,22 @@ class ChatTable:
async def update_chat_title_by_id(self, id: str, title: str) -> ChatModel | None:
try:
async with get_async_db_context() as db:
chat_item = await db.get(Chat, id)
chat_item = await session.get(Chat, id)
if chat_item is None:
return None
clean_title = self._clean_null_bytes(title)
chat_item.title = clean_title
chat_item.chat = {**(chat_item.chat or {}), 'title': clean_title}
chat_item.updated_at = int(time.time())
await db.commit()
await db.refresh(chat_item)
await session.commit()
await session.refresh(chat_item)
return ChatModel.model_validate(chat_item)
except Exception:
return None
async def update_chat_tags_by_id(self, id: str, tags: list[str], user) -> ChatModel | None:
async with get_async_db_context() as db:
chat = await db.get(Chat, id)
chat = await session.get(Chat, id)
if chat is None:
return None
@@ -448,22 +446,22 @@ class ChatTable:
# Single meta update
chat.meta = {**chat.meta, 'tags': new_tag_ids}
await db.commit()
await db.refresh(chat)
await session.commit()
await session.refresh(chat)
# Batch-create any missing tag rows
await Tags.ensure_tags_exist(new_tags, user.id, db=db)
await Tags.ensure_tags_exist(new_tags, user.id, db=session)
# 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)
await self.delete_orphan_tags_for_user(list(removed), user.id, db=session)
return ChatModel.model_validate(chat)
async def get_chat_title_by_id(self, id: str) -> str | None:
async with get_async_db_context() as db:
result = await db.execute(select(Chat.title).filter_by(id=id))
result = await session.execute(select(Chat.title).filter_by(id=id))
row = result.first()
if row is None:
return None
@@ -638,7 +636,7 @@ class ChatTable:
async def add_message_files_by_id_and_message_id(self, id: str, message_id: str, files: list[dict]) -> list[dict]:
async with get_async_db_context() as db:
chat = await self.get_chat_by_id(id, db=db)
chat = await self.get_chat_by_id(id, db=session)
if chat is None:
return None
@@ -653,61 +651,62 @@ class ChatTable:
history['messages'][message_id]['files'] = message_files
chat['history'] = history
await self.update_chat_by_id(id, chat, db=db)
await self.update_chat_by_id(id, chat, db=session)
return message_files
async def insert_shared_chat_by_chat_id(self, chat_id: str, db: AsyncSession | None = None) -> ChatModel | None:
"""Create a shared snapshot for a chat. Returns the original chat with share_id set."""
from open_webui.models.shared_chats import SharedChats
async with get_async_db_context(db) as db:
chat = await db.get(Chat, chat_id)
async with get_async_db_context(db) as session:
chat = await session.get(Chat, chat_id)
if not chat:
return None
# If already shared, just update the existing snapshot
if chat.share_id:
return await self.update_shared_chat_by_chat_id(chat_id, db=db)
return await self.update_shared_chat_by_chat_id(chat_id, db=session)
shared = await SharedChats.create(chat_id, chat.user_id, db=db)
shared = await SharedChats.create(chat_id, chat.user_id, db=session)
if not shared:
return None
# Set share_id on the original chat
chat.share_id = shared.id
await db.commit()
await db.refresh(chat)
return ChatModel.model_validate(chat)
await session.commit()
await session.refresh(chat)
return ChatModel.model_validate(chat) # return the updated original
async def update_shared_chat_by_chat_id(self, chat_id: str, db: AsyncSession | None = None) -> ChatModel | None:
"""Re-snapshot the shared chat with current chat data."""
from open_webui.models.shared_chats import SharedChats
async def update_shared_chat_by_chat_id(
self, chat_id: str, db: AsyncSession | None = None,
) -> ChatModel | None:
"""Re-snapshot the shared copy with the latest chat content."""
from open_webui.models.shared_chats import SharedChats # deferred — circular
try:
async with get_async_db_context(db) as db:
chat = await db.get(Chat, chat_id)
async with get_async_db_context(db) as session:
chat = await session.get(Chat, chat_id)
if not chat or not chat.share_id:
return await self.insert_shared_chat_by_chat_id(chat_id, db=db)
await SharedChats.update(chat.share_id, db=db)
return await self.insert_shared_chat_by_chat_id(chat_id, db=session)
await SharedChats.update(chat.share_id, db=session)
return ChatModel.model_validate(chat)
except Exception:
return None
return
async def delete_shared_chat_by_chat_id(self, chat_id: str, db: AsyncSession | None = None) -> bool:
"""Delete shared snapshot for a chat."""
from open_webui.models.shared_chats import SharedChats
try:
return await SharedChats.delete_by_chat_id(chat_id, db=db)
return await SharedChats.delete_by_chat_id(chat_id, db=session)
except Exception:
return False
async def unarchive_all_chats_by_user_id(self, user_id: str, db: AsyncSession | None = None) -> bool:
try:
async with get_async_db_context(db) as db:
await db.execute(update(Chat).filter_by(user_id=user_id).values(archived=False))
await db.commit()
async with get_async_db_context(db) as session:
await session.execute(update(Chat).filter_by(user_id=user_id).values(archived=False))
await session.commit()
return True
except Exception:
return False
@@ -716,45 +715,45 @@ class ChatTable:
self, id: str, share_id: str | None, db: AsyncSession | None = None
) -> ChatModel | None:
try:
async with get_async_db_context(db) as db:
chat = await db.get(Chat, id)
async with get_async_db_context(db) as session:
chat = await session.get(Chat, id)
chat.share_id = share_id
await db.commit()
await db.refresh(chat)
await session.commit()
await session.refresh(chat)
return ChatModel.model_validate(chat)
except Exception:
return None
async def toggle_chat_pinned_by_id(self, id: str, db: AsyncSession | None = None) -> ChatModel | None:
try:
async with get_async_db_context(db) as db:
chat = await db.get(Chat, id)
async with get_async_db_context(db) as session:
chat = await session.get(Chat, id)
chat.pinned = not chat.pinned
chat.updated_at = int(time.time())
await db.commit()
await db.refresh(chat)
await session.commit()
await session.refresh(chat)
return ChatModel.model_validate(chat)
except Exception:
return None
async def toggle_chat_archive_by_id(self, id: str, db: AsyncSession | None = None) -> ChatModel | None:
try:
async with get_async_db_context(db) as db:
chat = await db.get(Chat, id)
async with get_async_db_context(db) as session:
chat = await session.get(Chat, id)
chat.archived = not chat.archived
chat.folder_id = None
chat.updated_at = int(time.time())
await db.commit()
await db.refresh(chat)
await session.commit()
await session.refresh(chat)
return ChatModel.model_validate(chat)
except Exception:
return None
async def archive_all_chats_by_user_id(self, user_id: str, db: AsyncSession | None = None) -> bool:
try:
async with get_async_db_context(db) as db:
await db.execute(update(Chat).filter_by(user_id=user_id).values(archived=True))
await db.commit()
async with get_async_db_context(db) as session:
await session.execute(update(Chat).filter_by(user_id=user_id).values(archived=True))
await session.commit()
return True
except Exception:
return False
@@ -767,7 +766,7 @@ class ChatTable:
limit: int = 50,
db: AsyncSession | None = None,
) -> list[ChatTitleIdResponse]:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
stmt = select(Chat.id, Chat.title, Chat.updated_at, Chat.created_at).filter_by(
user_id=user_id, archived=True
)
@@ -798,7 +797,7 @@ class ChatTable:
if limit:
stmt = stmt.limit(limit)
result = await db.execute(stmt)
result = await session.execute(stmt)
all_chats = result.all()
return [
ChatTitleIdResponse.model_validate(
@@ -823,7 +822,7 @@ class ChatTable:
"""Delegate to SharedChats for listing shared chats by user."""
from open_webui.models.shared_chats import SharedChats
return await SharedChats.get_by_user_id(user_id, filter=filter, skip=skip, limit=limit, db=db)
return await SharedChats.get_by_user_id(user_id, filter=filter, skip=skip, limit=limit, db=session)
async def get_chat_list_by_user_id(
self,
@@ -834,7 +833,7 @@ class ChatTable:
limit: int = 50,
db: AsyncSession | None = None,
) -> list[ChatTitleIdResponse]:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
stmt = select(Chat.id, Chat.title, Chat.updated_at, Chat.created_at, Chat.last_read_at).filter_by(
user_id=user_id
)
@@ -864,7 +863,7 @@ class ChatTable:
if limit:
stmt = stmt.limit(limit)
result = await db.execute(stmt)
result = await session.execute(stmt)
all_chats = result.all()
return [
ChatTitleIdResponse.model_validate(
@@ -889,7 +888,7 @@ class ChatTable:
limit: int | None = None,
db: AsyncSession | None = None,
) -> list[ChatTitleIdResponse]:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
stmt = select(Chat.id, Chat.title, Chat.updated_at, Chat.created_at, Chat.last_read_at).filter_by(
user_id=user_id
)
@@ -910,7 +909,7 @@ class ChatTable:
if limit:
stmt = stmt.limit(limit)
result = await db.execute(stmt)
result = await session.execute(stmt)
all_chats = result.all()
return [
@@ -933,23 +932,26 @@ class ChatTable:
limit: int = 50,
db: AsyncSession | None = None,
) -> list[ChatModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
async with get_async_db_context(db) as session:
result = await session.execute(
select(Chat).filter(Chat.id.in_(chat_ids)).filter_by(archived=False).order_by(Chat.updated_at.desc())
)
all_chats = result.scalars().all()
return [ChatModel.model_validate(chat) for chat in all_chats]
async def get_chat_by_id(self, id: str, db: AsyncSession | None = None) -> ChatModel | None:
async def get_chat_by_id(
self, id: str, db: AsyncSession | None = None,
) -> ChatModel | None:
"""Fetch a chat by PK, auto-sanitizing null bytes on read."""
try:
async with get_async_db_context(db) as db:
chat_item = await db.get(Chat, id)
async with get_async_db_context(db) as session:
chat_item = await session.get(Chat, id)
if chat_item is None:
return None
if self._sanitize_chat_row(chat_item):
await db.commit()
await db.refresh(chat_item)
await session.commit()
await session.refresh(chat_item)
return ChatModel.model_validate(chat_item)
except Exception:
@@ -960,7 +962,7 @@ class ChatTable:
from open_webui.models.shared_chats import SharedChats
try:
shared = await SharedChats.get_by_id(id, db=db)
shared = await SharedChats.get_by_id(id, db=session)
if shared:
# Return a ChatModel-compatible view of the snapshot
return ChatModel(
@@ -980,8 +982,8 @@ class ChatTable:
self, id: str, user_id: str, db: AsyncSession | None = None
) -> ChatModel | None:
try:
async with get_async_db_context(db) as db:
result = await db.execute(select(Chat).filter_by(id=id, user_id=user_id))
async with get_async_db_context(db) as session:
result = await session.execute(select(Chat).filter_by(id=id, user_id=user_id))
chat = result.scalars().first()
return ChatModel.model_validate(chat) if chat else None
except Exception:
@@ -993,8 +995,8 @@ class ChatTable:
the full Chat row (which includes the potentially large JSON blob).
"""
try:
async with get_async_db_context(db) as db:
result = await db.execute(select(exists().where(and_(Chat.id == id, Chat.user_id == user_id))))
async with get_async_db_context(db) as session:
result = await session.execute(select(exists().where(and_(Chat.id == id, Chat.user_id == user_id))))
return result.scalar()
except Exception:
return False
@@ -1005,16 +1007,16 @@ class ChatTable:
JSON blob. Returns None if chat doesn't exist or doesn't belong to user.
"""
try:
async with get_async_db_context(db) as db:
result = await db.execute(select(Chat.folder_id).filter_by(id=id, user_id=user_id))
async with get_async_db_context(db) as session:
result = await session.execute(select(Chat.folder_id).filter_by(id=id, user_id=user_id))
row = result.first()
return row[0] if row else None
except Exception:
return None
async def get_chats(self, skip: int = 0, limit: int = 50, db: AsyncSession | None = None) -> list[ChatModel]:
async with get_async_db_context(db) as db:
result = await db.execute(select(Chat).order_by(Chat.updated_at.desc()))
async with get_async_db_context(db) as session:
result = await session.execute(select(Chat).order_by(Chat.updated_at.desc()))
all_chats = result.scalars().all()
return [ChatModel.model_validate(chat) for chat in all_chats]
@@ -1026,7 +1028,7 @@ class ChatTable:
limit: int | None = None,
db: AsyncSession | None = None,
) -> ChatListResponse:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
stmt = select(Chat).filter_by(user_id=user_id)
if filter:
@@ -1048,7 +1050,7 @@ class ChatTable:
else:
stmt = stmt.order_by(Chat.updated_at.desc(), Chat.id)
count_result = await db.execute(select(func.count()).select_from(stmt.subquery()))
count_result = await session.execute(select(func.count()).select_from(stmt.subquery()))
total = count_result.scalar()
if skip is not None:
@@ -1056,7 +1058,7 @@ class ChatTable:
if limit is not None:
stmt = stmt.limit(limit)
result = await db.execute(stmt)
result = await session.execute(stmt)
all_chats = result.scalars().all()
return ChatListResponse(
@@ -1069,8 +1071,8 @@ class ChatTable:
async def get_pinned_chats_by_user_id(
self, user_id: str, db: AsyncSession | None = None
) -> list[ChatTitleIdResponse]:
async with get_async_db_context(db) as db:
result = await db.execute(
async with get_async_db_context(db) as session:
result = await session.execute(
select(Chat.id, Chat.title, Chat.updated_at, Chat.created_at, Chat.last_read_at)
.filter_by(user_id=user_id, pinned=True, archived=False)
.order_by(Chat.updated_at.desc())
@@ -1090,8 +1092,8 @@ class ChatTable:
]
async def get_archived_chats_by_user_id(self, user_id: str, db: AsyncSession | None = None) -> list[ChatModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
async with get_async_db_context(db) as session:
result = await session.execute(
select(Chat).filter_by(user_id=user_id, archived=True).order_by(Chat.updated_at.desc())
)
return [ChatModel.model_validate(chat) for chat in result.scalars().all()]
@@ -1161,7 +1163,7 @@ class ChatTable:
search_text = ' '.join(search_text_words)
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
stmt = select(Chat).filter(Chat.user_id == user_id)
if is_archived is not None:
@@ -1184,7 +1186,7 @@ class ChatTable:
stmt = stmt.order_by(Chat.updated_at.desc(), Chat.id)
# Check if the database dialect is either 'sqlite' or 'postgresql'
bind = await db.connection()
bind = await session.connection()
dialect_name = bind.dialect.name
if dialect_name == 'sqlite':
# SQLite case: using JSON1 extension for JSON searching
@@ -1282,7 +1284,7 @@ class ChatTable:
# Perform pagination at the SQL level
stmt = stmt.offset(skip).limit(limit)
result = await db.execute(stmt)
result = await session.execute(stmt)
all_chats = result.scalars().all()
log.info(f'The number of chats: {len(all_chats)}')
@@ -1298,7 +1300,7 @@ class ChatTable:
limit: int = 60,
db: AsyncSession | None = None,
) -> list[ChatTitleIdResponse]:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
stmt = (
select(Chat.id, Chat.title, Chat.updated_at, Chat.created_at, Chat.last_read_at)
.filter_by(folder_id=folder_id, user_id=user_id)
@@ -1312,7 +1314,7 @@ class ChatTable:
if limit:
stmt = stmt.limit(limit)
result = await db.execute(stmt)
result = await session.execute(stmt)
all_chats = result.all()
return [
ChatTitleIdResponse.model_validate(
@@ -1330,7 +1332,7 @@ class ChatTable:
async def get_chats_by_folder_ids_and_user_id(
self, folder_ids: list[str], user_id: str, db: AsyncSession | None = None
) -> list[ChatModel]:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
stmt = (
select(Chat)
.filter(Chat.folder_id.in_(folder_ids), Chat.user_id == user_id)
@@ -1339,7 +1341,7 @@ class ChatTable:
.order_by(Chat.updated_at.desc())
)
result = await db.execute(stmt)
result = await session.execute(stmt)
all_chats = result.scalars().all()
return [ChatModel.model_validate(chat) for chat in all_chats]
@@ -1347,13 +1349,13 @@ class ChatTable:
self, id: str, user_id: str, folder_id: str, db: AsyncSession | None = None
) -> ChatModel | None:
try:
async with get_async_db_context(db) as db:
chat = await db.get(Chat, id)
async with get_async_db_context(db) as session:
chat = await session.get(Chat, id)
chat.folder_id = folder_id
chat.updated_at = int(time.time())
chat.pinned = False
await db.commit()
await db.refresh(chat)
await session.commit()
await session.refresh(chat)
return ChatModel.model_validate(chat)
except Exception:
return None
@@ -1361,12 +1363,12 @@ class ChatTable:
async def get_chat_tags_by_id_and_user_id(
self, id: str, user_id: str, db: AsyncSession | None = None
) -> list[TagModel]:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
stmt = select(Chat.meta).where(Chat.id == id)
result = await db.execute(stmt)
result = await session.execute(stmt)
meta = result.scalar_one_or_none()
tag_ids = (meta or {}).get('tags', [])
return await Tags.get_tags_by_ids_and_user_id(tag_ids, user_id, db=db)
return await Tags.get_tags_by_ids_and_user_id(tag_ids, user_id, db=session)
async def get_chat_list_by_user_id_and_tag_name(
self,
@@ -1376,13 +1378,13 @@ class ChatTable:
limit: int = 50,
db: AsyncSession | None = None,
) -> list[ChatTitleIdResponse]:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
stmt = select(Chat.id, Chat.title, Chat.updated_at, Chat.created_at, Chat.last_read_at).filter_by(
user_id=user_id
)
tag_id = tag_name.replace(' ', '_').lower()
bind = await db.connection()
bind = await session.connection()
dialect_name = bind.dialect.name
log.info(f'DB dialect name: {dialect_name}')
if dialect_name == 'sqlite':
@@ -1403,7 +1405,7 @@ class ChatTable:
if limit:
stmt = stmt.limit(limit)
result = await db.execute(stmt)
result = await session.execute(stmt)
all_chats = result.all()
return [
ChatTitleIdResponse.model_validate(
@@ -1422,17 +1424,17 @@ class ChatTable:
self, id: str, user_id: str, tag_name: str, db: AsyncSession | None = None
) -> ChatModel | None:
tag_id = tag_name.replace(' ', '_').lower()
await Tags.ensure_tags_exist([tag_name], user_id, db=db)
await Tags.ensure_tags_exist([tag_name], user_id, db=session)
try:
async with get_async_db_context(db) as db:
chat = await db.get(Chat, id)
async with get_async_db_context(db) as session:
chat = await session.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])),
}
await db.commit()
await db.refresh(chat)
await session.commit()
await session.refresh(chat)
return ChatModel.model_validate(chat)
except Exception:
return None
@@ -1440,11 +1442,11 @@ class ChatTable:
async def count_chats_by_tag_name_and_user_id(
self, tag_name: str, user_id: str, db: AsyncSession | None = None
) -> int:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
stmt = select(func.count(Chat.id)).filter_by(user_id=user_id, archived=False)
tag_id = tag_name.replace(' ', '_').lower()
bind = await db.connection()
bind = await session.connection()
dialect_name = bind.dialect.name
if dialect_name == 'sqlite':
stmt = stmt.filter(
@@ -1457,7 +1459,7 @@ class ChatTable:
else:
raise NotImplementedError(f'Unsupported dialect: {dialect_name}')
result = await db.execute(stmt)
result = await session.execute(stmt)
return result.scalar()
async def delete_orphan_tags_for_user(
@@ -1477,19 +1479,19 @@ class ChatTable:
"""
if not tag_ids:
return
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
orphans = []
for tag_id in tag_ids:
count = await self.count_chats_by_tag_name_and_user_id(tag_id, user_id, db=db)
count = await self.count_chats_by_tag_name_and_user_id(tag_id, user_id, db=session)
if count <= threshold:
orphans.append(tag_id)
await Tags.delete_tags_by_ids_and_user_id(orphans, user_id, db=db)
await Tags.delete_tags_by_ids_and_user_id(orphans, user_id, db=session)
async def count_chats_by_folder_id_and_user_id(
self, folder_id: str, user_id: str, db: AsyncSession | None = None
) -> int:
async with get_async_db_context(db) as db:
result = await db.execute(select(func.count(Chat.id)).filter_by(user_id=user_id, folder_id=folder_id))
async with get_async_db_context(db) as session:
result = await session.execute(select(func.count(Chat.id)).filter_by(user_id=user_id, folder_id=folder_id))
count = result.scalar()
log.info(f"Count of chats for folder '{folder_id}': {count}")
@@ -1499,8 +1501,8 @@ class ChatTable:
self, id: str, user_id: str, tag_name: str, db: AsyncSession | None = None
) -> bool:
try:
async with get_async_db_context(db) as db:
chat = await db.get(Chat, id)
async with get_async_db_context(db) as session:
chat = await session.get(Chat, id)
tags = chat.meta.get('tags', [])
tag_id = tag_name.replace(' ', '_').lower()
@@ -1509,51 +1511,51 @@ class ChatTable:
**chat.meta,
'tags': list(set(tags)),
}
await db.commit()
await session.commit()
return True
except Exception:
return False
async def delete_chat_by_id(self, id: str, db: AsyncSession | None = 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(Chat).filter_by(id=id))
await db.commit()
async with get_async_db_context(db) as session:
await session.execute(update(AutomationRun).filter_by(chat_id=id).values(chat_id=None))
await session.execute(delete(ChatMessage).filter_by(chat_id=id))
await session.execute(delete(Chat).filter_by(id=id))
await session.commit()
return True and await self.delete_shared_chat_by_chat_id(id, db=db)
return True and await self.delete_shared_chat_by_chat_id(id, db=session)
except Exception:
return False
async def delete_chat_by_id_and_user_id(self, id: str, user_id: str, db: AsyncSession | None = 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(Chat).filter_by(id=id, user_id=user_id))
await db.commit()
async with get_async_db_context(db) as session:
await session.execute(update(AutomationRun).filter_by(chat_id=id).values(chat_id=None))
await session.execute(delete(ChatMessage).filter_by(chat_id=id))
await session.execute(delete(Chat).filter_by(id=id, user_id=user_id))
await session.commit()
return True and await self.delete_shared_chat_by_chat_id(id, db=db)
return True and await self.delete_shared_chat_by_chat_id(id, db=session)
except Exception:
return False
async def delete_chats_by_user_id(self, user_id: str, db: AsyncSession | None = None) -> bool:
try:
async with get_async_db_context(db) as db:
await self.delete_shared_chats_by_user_id(user_id, db=db)
async with get_async_db_context(db) as session:
await self.delete_shared_chats_by_user_id(user_id, db=session)
chat_id_subquery = select(Chat.id).filter_by(user_id=user_id).scalar_subquery()
await db.execute(
await session.execute(
update(AutomationRun)
.filter(AutomationRun.chat_id.in_(select(Chat.id).filter_by(user_id=user_id)))
.values(chat_id=None)
)
await db.execute(
await session.execute(
delete(ChatMessage).filter(ChatMessage.chat_id.in_(select(Chat.id).filter_by(user_id=user_id)))
)
await db.execute(delete(Chat).filter_by(user_id=user_id))
await db.commit()
await session.execute(delete(Chat).filter_by(user_id=user_id))
await session.commit()
return True
except Exception:
@@ -1563,14 +1565,14 @@ class ChatTable:
self, user_id: str, folder_id: str, db: AsyncSession | None = None
) -> bool:
try:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
chat_ids_stmt = select(Chat.id).filter_by(user_id=user_id, folder_id=folder_id)
await db.execute(
await session.execute(
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(Chat).filter_by(user_id=user_id, folder_id=folder_id))
await db.commit()
await session.execute(delete(ChatMessage).filter(ChatMessage.chat_id.in_(chat_ids_stmt)))
await session.execute(delete(Chat).filter_by(user_id=user_id, folder_id=folder_id))
await session.commit()
return True
except Exception:
@@ -1584,11 +1586,11 @@ class ChatTable:
db: AsyncSession | None = None,
) -> bool:
try:
async with get_async_db_context(db) as db:
await db.execute(
async with get_async_db_context(db) as session:
await session.execute(
update(Chat).filter_by(user_id=user_id, folder_id=folder_id).values(folder_id=new_folder_id)
)
await db.commit()
await session.commit()
return True
except Exception:
@@ -1600,13 +1602,13 @@ class ChatTable:
from open_webui.models.shared_chats import SharedChats
try:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
# Delete shared_chat rows for this user's chats
await db.execute(delete(SharedChatTable).filter_by(user_id=user_id))
await session.execute(delete(SharedChatTable).filter_by(user_id=user_id))
# Clear share_id on all of this user's chats
await db.execute(update(Chat).filter_by(user_id=user_id).values(share_id=None))
await db.commit()
await session.execute(update(Chat).filter_by(user_id=user_id).values(share_id=None))
await session.commit()
return True
except Exception:
@@ -1624,7 +1626,7 @@ class ChatTable:
return None
chat_message_file_ids = {
item.id for item in await self.get_chat_files_by_chat_id_and_message_id(chat_id, message_id, db=db)
item.id for item in await self.get_chat_files_by_chat_id_and_message_id(chat_id, message_id, db=session)
}
# Remove duplicates and existing file_ids
file_ids = list({file_id for file_id in file_ids if file_id and file_id not in chat_message_file_ids})
@@ -1632,7 +1634,7 @@ class ChatTable:
return None
try:
async with get_async_db_context(db) as db:
async with get_async_db_context(db) as session:
now = int(time.time())
chat_files = [
@@ -1651,7 +1653,7 @@ class ChatTable:
results = [ChatFile(**chat_file.model_dump()) for chat_file in chat_files]
db.add_all(results)
await db.commit()
await session.commit()
return chat_files
except Exception:
@@ -1660,8 +1662,8 @@ class ChatTable:
async def get_chat_files_by_chat_id_and_message_id(
self, chat_id: str, message_id: str, db: AsyncSession | None = None
) -> list[ChatFileModel]:
async with get_async_db_context(db) as db:
result = await db.execute(
async with get_async_db_context(db) as session:
result = await session.execute(
select(ChatFile).filter_by(chat_id=chat_id, message_id=message_id).order_by(ChatFile.created_at.asc())
)
all_chat_files = result.scalars().all()
@@ -1669,17 +1671,17 @@ class ChatTable:
async def delete_chat_file(self, chat_id: str, file_id: str, db: AsyncSession | None = None) -> bool:
try:
async with get_async_db_context(db) as db:
await db.execute(delete(ChatFile).filter_by(chat_id=chat_id, file_id=file_id))
await db.commit()
async with get_async_db_context(db) as session:
await session.execute(delete(ChatFile).filter_by(chat_id=chat_id, file_id=file_id))
await session.commit()
return True
except Exception:
return False
async def get_shared_chat_ids_by_file_id(self, file_id: str, db: AsyncSession | None = None) -> list[str]:
"""Return IDs of chats that contain this file and have an active share link."""
async with get_async_db_context(db) as db:
result = await db.execute(
async with get_async_db_context(db) as session:
result = await session.execute(
select(Chat.id)
.join(ChatFile, Chat.id == ChatFile.chat_id)
.filter(ChatFile.file_id == file_id, Chat.share_id.isnot(None))
@@ -1690,12 +1692,12 @@ class ChatTable:
"""Update the tasks list on a chat."""
try:
async with get_async_db_context() as db:
chat = await db.get(Chat, id)
chat = await session.get(Chat, id)
if chat is None:
return None
chat.tasks = tasks
await db.commit()
await db.refresh(chat)
await session.commit()
await session.refresh(chat)
return ChatModel.model_validate(chat)
except Exception:
return None
@@ -1703,7 +1705,7 @@ class ChatTable:
async def get_chat_tasks_by_id(self, id: str) -> list[dict]:
"""Read the tasks list from a chat (lightweight column query)."""
async with get_async_db_context() as db:
result = await db.execute(select(Chat.tasks).filter_by(id=id))
result = await session.execute(select(Chat.tasks).filter_by(id=id))
row = result.first()
if row is None or row[0] is None:
return []