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

refac/enh: db session sharing

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2025-12-28 22:00:44 +04:00
parent d4de26bd05
commit 2041ab483e
20 changed files with 600 additions and 562 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/models/chats.py
+96 -90
View File
@@ -4,7 +4,8 @@ import time
import uuid
from typing import Optional
from open_webui.internal.db import Base, get_db
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, JSONField, get_db, get_db_context
from open_webui.models.tags import TagModel, Tag, Tags
from open_webui.models.folders import Folders
from open_webui.utils.misc import sanitize_data_for_db, sanitize_text_for_db
@@ -280,8 +281,8 @@ class ChatTable:
return changed
def insert_new_chat(self, user_id: str, form_data: ChatForm) -> Optional[ChatModel]:
with get_db() as db:
def insert_new_chat(self, user_id: str, form_data: ChatForm, db: Optional[Session] = None) -> Optional[ChatModel]:
with get_db_context(db) as db:
id = str(uuid.uuid4())
chat = ChatModel(
**{
@@ -331,9 +332,9 @@ class ChatTable:
return chat
def import_chats(
self, user_id: str, chat_import_forms: list[ChatImportForm]
self, user_id: str, chat_import_forms: list[ChatImportForm], db: Optional[Session] = None
) -> list[ChatModel]:
with get_db() as db:
with get_db_context(db) as db:
chats = []
for form_data in chat_import_forms:
@@ -344,9 +345,9 @@ class ChatTable:
db.commit()
return [ChatModel.model_validate(chat) for chat in chats]
def update_chat_by_id(self, id: str, chat: dict) -> Optional[ChatModel]:
def update_chat_by_id(self, id: str, chat: dict, db: Optional[Session] = None) -> Optional[ChatModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
chat_item = db.get(Chat, id)
chat_item.chat = self._clean_null_bytes(chat)
chat_item.title = (
@@ -483,13 +484,13 @@ class ChatTable:
self.update_chat_by_id(id, chat)
return message_files
def insert_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
with get_db() as db:
def insert_shared_chat_by_chat_id(self, chat_id: str, db: Optional[Session] = None) -> Optional[ChatModel]:
with get_db_context(db) as db:
# Get the existing chat to share
chat = db.get(Chat, chat_id)
# Check if the chat is already shared
if chat.share_id:
return self.get_chat_by_id_and_user_id(chat.share_id, "shared")
return self.get_chat_by_id_and_user_id(chat.share_id, "shared", db=db)
# Create a new chat with the same data, but with a new ID
shared_chat = ChatModel(
**{
@@ -518,16 +519,16 @@ class ChatTable:
db.commit()
return shared_chat if (shared_result and result) else None
def update_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
def update_shared_chat_by_chat_id(self, chat_id: str, db: Optional[Session] = None) -> Optional[ChatModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
chat = db.get(Chat, chat_id)
shared_chat = (
db.query(Chat).filter_by(user_id=f"shared-{chat_id}").first()
)
if shared_chat is None:
return self.insert_shared_chat_by_chat_id(chat_id)
return self.insert_shared_chat_by_chat_id(chat_id, db=db)
shared_chat.title = chat.title
shared_chat.chat = chat.chat
@@ -542,9 +543,9 @@ class ChatTable:
except Exception:
return None
def delete_shared_chat_by_chat_id(self, chat_id: str) -> bool:
def delete_shared_chat_by_chat_id(self, chat_id: str, db: Optional[Session] = None) -> bool:
try:
with get_db() as db:
with get_db_context(db) as db:
db.query(Chat).filter_by(user_id=f"shared-{chat_id}").delete()
db.commit()
@@ -552,9 +553,9 @@ class ChatTable:
except Exception:
return False
def unarchive_all_chats_by_user_id(self, user_id: str) -> bool:
def unarchive_all_chats_by_user_id(self, user_id: str, db: Optional[Session] = None) -> bool:
try:
with get_db() as db:
with get_db_context(db) as db:
db.query(Chat).filter_by(user_id=user_id).update({"archived": False})
db.commit()
return True
@@ -562,10 +563,10 @@ class ChatTable:
return False
def update_chat_share_id_by_id(
self, id: str, share_id: Optional[str]
self, id: str, share_id: Optional[str], db: Optional[Session] = None
) -> Optional[ChatModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
chat = db.get(Chat, id)
chat.share_id = share_id
db.commit()
@@ -574,9 +575,9 @@ class ChatTable:
except Exception:
return None
def toggle_chat_pinned_by_id(self, id: str) -> Optional[ChatModel]:
def toggle_chat_pinned_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ChatModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
chat = db.get(Chat, id)
chat.pinned = not chat.pinned
chat.updated_at = int(time.time())
@@ -586,9 +587,9 @@ class ChatTable:
except Exception:
return None
def toggle_chat_archive_by_id(self, id: str) -> Optional[ChatModel]:
def toggle_chat_archive_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ChatModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
chat = db.get(Chat, id)
chat.archived = not chat.archived
chat.folder_id = None
@@ -599,9 +600,9 @@ class ChatTable:
except Exception:
return None
def archive_all_chats_by_user_id(self, user_id: str) -> bool:
def archive_all_chats_by_user_id(self, user_id: str, db: Optional[Session] = None) -> bool:
try:
with get_db() as db:
with get_db_context(db) as db:
db.query(Chat).filter_by(user_id=user_id).update({"archived": True})
db.commit()
return True
@@ -614,9 +615,10 @@ class ChatTable:
filter: Optional[dict] = None,
skip: int = 0,
limit: int = 50,
db: Optional[Session] = None,
) -> list[ChatModel]:
with get_db() as db:
with get_db_context(db) as db:
query = db.query(Chat).filter_by(user_id=user_id, archived=True)
if filter:
@@ -655,8 +657,9 @@ class ChatTable:
filter: Optional[dict] = None,
skip: int = 0,
limit: int = 50,
db: Optional[Session] = None,
) -> list[ChatModel]:
with get_db() as db:
with get_db_context(db) as db:
query = db.query(Chat).filter_by(user_id=user_id)
if not include_archived:
query = query.filter_by(archived=False)
@@ -695,8 +698,9 @@ class ChatTable:
include_pinned: bool = False,
skip: Optional[int] = None,
limit: Optional[int] = None,
db: Optional[Session] = None,
) -> list[ChatTitleIdResponse]:
with get_db() as db:
with get_db_context(db) as db:
query = db.query(Chat).filter_by(user_id=user_id)
if not include_folders:
@@ -733,9 +737,9 @@ class ChatTable:
]
def get_chat_list_by_chat_ids(
self, chat_ids: list[str], skip: int = 0, limit: int = 50
self, chat_ids: list[str], skip: int = 0, limit: int = 50, db: Optional[Session] = None
) -> list[ChatModel]:
with get_db() as db:
with get_db_context(db) as db:
all_chats = (
db.query(Chat)
.filter(Chat.id.in_(chat_ids))
@@ -745,9 +749,9 @@ class ChatTable:
)
return [ChatModel.model_validate(chat) for chat in all_chats]
def get_chat_by_id(self, id: str) -> Optional[ChatModel]:
def get_chat_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ChatModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
chat_item = db.get(Chat, id)
if chat_item is None:
return None
@@ -760,30 +764,30 @@ class ChatTable:
except Exception:
return None
def get_chat_by_share_id(self, id: str) -> Optional[ChatModel]:
def get_chat_by_share_id(self, id: str, db: Optional[Session] = None) -> Optional[ChatModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
# it is possible that the shared link was deleted. hence,
# we check if the chat is still shared by checking if a chat with the share_id exists
chat = db.query(Chat).filter_by(share_id=id).first()
if chat:
return self.get_chat_by_id(id)
return self.get_chat_by_id(id, db=db)
else:
return None
except Exception:
return None
def get_chat_by_id_and_user_id(self, id: str, user_id: str) -> Optional[ChatModel]:
def get_chat_by_id_and_user_id(self, id: str, user_id: str, db: Optional[Session] = None) -> Optional[ChatModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
chat = db.query(Chat).filter_by(id=id, user_id=user_id).first()
return ChatModel.model_validate(chat)
except Exception:
return None
def get_chats(self, skip: int = 0, limit: int = 50) -> list[ChatModel]:
with get_db() as db:
def get_chats(self, skip: int = 0, limit: int = 50, db: Optional[Session] = None) -> list[ChatModel]:
with get_db_context(db) as db:
all_chats = (
db.query(Chat)
# .limit(limit).offset(skip)
@@ -797,8 +801,9 @@ class ChatTable:
filter: Optional[dict] = None,
skip: Optional[int] = None,
limit: Optional[int] = None,
db: Optional[Session] = None,
) -> ChatListResponse:
with get_db() as db:
with get_db_context(db) as db:
query = db.query(Chat).filter_by(user_id=user_id)
if filter:
@@ -838,8 +843,8 @@ class ChatTable:
}
)
def get_pinned_chats_by_user_id(self, user_id: str) -> list[ChatModel]:
with get_db() as db:
def get_pinned_chats_by_user_id(self, user_id: str, db: Optional[Session] = None) -> list[ChatModel]:
with get_db_context(db) as db:
all_chats = (
db.query(Chat)
.filter_by(user_id=user_id, pinned=True, archived=False)
@@ -847,8 +852,8 @@ class ChatTable:
)
return [ChatModel.model_validate(chat) for chat in all_chats]
def get_archived_chats_by_user_id(self, user_id: str) -> list[ChatModel]:
with get_db() as db:
def get_archived_chats_by_user_id(self, user_id: str, db: Optional[Session] = None) -> list[ChatModel]:
with get_db_context(db) as db:
all_chats = (
db.query(Chat)
.filter_by(user_id=user_id, archived=True)
@@ -863,6 +868,7 @@ class ChatTable:
include_archived: bool = False,
skip: int = 0,
limit: int = 60,
db: Optional[Session] = None,
) -> list[ChatModel]:
"""
Filters chats based on a search query using Python, allowing pagination using skip and limit.
@@ -871,7 +877,7 @@ class ChatTable:
if not search_text:
return self.get_chat_list_by_user_id(
user_id, include_archived, filter={}, skip=skip, limit=limit
user_id, include_archived, filter={}, skip=skip, limit=limit, db=db
)
search_text_words = search_text.split(" ")
@@ -926,7 +932,7 @@ class ChatTable:
search_text = " ".join(search_text_words)
with get_db() as db:
with get_db_context(db) as db:
query = db.query(Chat).filter(Chat.user_id == user_id)
if is_archived is not None:
@@ -1067,9 +1073,9 @@ class ChatTable:
return [ChatModel.model_validate(chat) for chat in all_chats]
def get_chats_by_folder_id_and_user_id(
self, folder_id: str, user_id: str, skip: int = 0, limit: int = 60
self, folder_id: str, user_id: str, skip: int = 0, limit: int = 60, db: Optional[Session] = None
) -> list[ChatModel]:
with get_db() as db:
with get_db_context(db) as db:
query = db.query(Chat).filter_by(folder_id=folder_id, user_id=user_id)
query = query.filter(or_(Chat.pinned == False, Chat.pinned == None))
query = query.filter_by(archived=False)
@@ -1085,9 +1091,9 @@ class ChatTable:
return [ChatModel.model_validate(chat) for chat in all_chats]
def get_chats_by_folder_ids_and_user_id(
self, folder_ids: list[str], user_id: str
self, folder_ids: list[str], user_id: str, db: Optional[Session] = None
) -> list[ChatModel]:
with get_db() as db:
with get_db_context(db) as db:
query = db.query(Chat).filter(
Chat.folder_id.in_(folder_ids), Chat.user_id == user_id
)
@@ -1100,10 +1106,10 @@ class ChatTable:
return [ChatModel.model_validate(chat) for chat in all_chats]
def update_chat_folder_id_by_id_and_user_id(
self, id: str, user_id: str, folder_id: str
self, id: str, user_id: str, folder_id: str, db: Optional[Session] = None
) -> Optional[ChatModel]:
try:
with get_db() as db:
with get_db_context(db) as db:
chat = db.get(Chat, id)
chat.folder_id = folder_id
chat.updated_at = int(time.time())
@@ -1114,16 +1120,16 @@ class ChatTable:
except Exception:
return None
def get_chat_tags_by_id_and_user_id(self, id: str, user_id: str) -> list[TagModel]:
with get_db() as db:
def get_chat_tags_by_id_and_user_id(self, id: str, user_id: str, db: Optional[Session] = None) -> list[TagModel]:
with get_db_context(db) as db:
chat = db.get(Chat, id)
tags = chat.meta.get("tags", [])
return [Tags.get_tag_by_name_and_user_id(tag, user_id) for tag in tags]
def get_chat_list_by_user_id_and_tag_name(
self, user_id: str, tag_name: str, skip: int = 0, limit: int = 50
self, user_id: str, tag_name: str, skip: int = 0, limit: int = 50, db: Optional[Session] = None
) -> list[ChatModel]:
with get_db() as db:
with get_db_context(db) as db:
query = db.query(Chat).filter_by(user_id=user_id)
tag_id = tag_name.replace(" ", "_").lower()
@@ -1152,13 +1158,13 @@ class ChatTable:
return [ChatModel.model_validate(chat) for chat in all_chats]
def add_chat_tag_by_id_and_user_id_and_tag_name(
self, id: str, user_id: str, tag_name: str
self, id: str, user_id: str, tag_name: str, db: Optional[Session] = None
) -> Optional[ChatModel]:
tag = Tags.get_tag_by_name_and_user_id(tag_name, user_id)
if tag is None:
tag = Tags.insert_new_tag(tag_name, user_id)
try:
with get_db() as db:
with get_db_context(db) as db:
chat = db.get(Chat, id)
tag_id = tag.id
@@ -1174,8 +1180,8 @@ class ChatTable:
except Exception:
return None
def count_chats_by_tag_name_and_user_id(self, tag_name: str, user_id: str) -> int:
with get_db() as db: # Assuming `get_db()` returns a session object
def count_chats_by_tag_name_and_user_id(self, tag_name: str, user_id: str, db: Optional[Session] = None) -> int:
with get_db_context(db) as db: # Assuming `get_db()` returns a session object
query = db.query(Chat).filter_by(user_id=user_id, archived=False)
# Normalize the tag_name for consistency
@@ -1210,8 +1216,8 @@ class ChatTable:
return count
def count_chats_by_folder_id_and_user_id(self, folder_id: str, user_id: str) -> int:
with get_db() as db:
def count_chats_by_folder_id_and_user_id(self, folder_id: str, user_id: str, db: Optional[Session] = None) -> int:
with get_db_context(db) as db:
query = db.query(Chat).filter_by(user_id=user_id)
query = query.filter_by(folder_id=folder_id)
@@ -1221,10 +1227,10 @@ class ChatTable:
return count
def delete_tag_by_id_and_user_id_and_tag_name(
self, id: str, user_id: str, tag_name: str
self, id: str, user_id: str, tag_name: str, db: Optional[Session] = None
) -> bool:
try:
with get_db() as db:
with get_db_context(db) as db:
chat = db.get(Chat, id)
tags = chat.meta.get("tags", [])
tag_id = tag_name.replace(" ", "_").lower()
@@ -1239,9 +1245,9 @@ class ChatTable:
except Exception:
return False
def delete_all_tags_by_id_and_user_id(self, id: str, user_id: str) -> bool:
def delete_all_tags_by_id_and_user_id(self, id: str, user_id: str, db: Optional[Session] = None) -> bool:
try:
with get_db() as db:
with get_db_context(db) as db:
chat = db.get(Chat, id)
chat.meta = {
**chat.meta,
@@ -1253,30 +1259,30 @@ class ChatTable:
except Exception:
return False
def delete_chat_by_id(self, id: str) -> bool:
def delete_chat_by_id(self, id: str, db: Optional[Session] = None) -> bool:
try:
with get_db() as db:
with get_db_context(db) as db:
db.query(Chat).filter_by(id=id).delete()
db.commit()
return True and self.delete_shared_chat_by_chat_id(id)
return True and self.delete_shared_chat_by_chat_id(id, db=db)
except Exception:
return False
def delete_chat_by_id_and_user_id(self, id: str, user_id: str) -> bool:
def delete_chat_by_id_and_user_id(self, id: str, user_id: str, db: Optional[Session] = None) -> bool:
try:
with get_db() as db:
with get_db_context(db) as db:
db.query(Chat).filter_by(id=id, user_id=user_id).delete()
db.commit()
return True and self.delete_shared_chat_by_chat_id(id)
return True and self.delete_shared_chat_by_chat_id(id, db=db)
except Exception:
return False
def delete_chats_by_user_id(self, user_id: str) -> bool:
def delete_chats_by_user_id(self, user_id: str, db: Optional[Session] = None) -> bool:
try:
with get_db() as db:
self.delete_shared_chats_by_user_id(user_id)
with get_db_context(db) as db:
self.delete_shared_chats_by_user_id(user_id, db=db)
db.query(Chat).filter_by(user_id=user_id).delete()
db.commit()
@@ -1286,10 +1292,10 @@ class ChatTable:
return False
def delete_chats_by_user_id_and_folder_id(
self, user_id: str, folder_id: str
self, user_id: str, folder_id: str, db: Optional[Session] = None
) -> bool:
try:
with get_db() as db:
with get_db_context(db) as db:
db.query(Chat).filter_by(user_id=user_id, folder_id=folder_id).delete()
db.commit()
@@ -1298,10 +1304,10 @@ class ChatTable:
return False
def move_chats_by_user_id_and_folder_id(
self, user_id: str, folder_id: str, new_folder_id: Optional[str]
self, user_id: str, folder_id: str, new_folder_id: Optional[str], db: Optional[Session] = None
) -> bool:
try:
with get_db() as db:
with get_db_context(db) as db:
db.query(Chat).filter_by(user_id=user_id, folder_id=folder_id).update(
{"folder_id": new_folder_id}
)
@@ -1311,9 +1317,9 @@ class ChatTable:
except Exception:
return False
def delete_shared_chats_by_user_id(self, user_id: str) -> bool:
def delete_shared_chats_by_user_id(self, user_id: str, db: Optional[Session] = None) -> bool:
try:
with get_db() as db:
with get_db_context(db) as db:
chats_by_user = db.query(Chat).filter_by(user_id=user_id).all()
shared_chat_ids = [f"shared-{chat.id}" for chat in chats_by_user]
@@ -1325,7 +1331,7 @@ class ChatTable:
return False
def insert_chat_files(
self, chat_id: str, message_id: str, file_ids: list[str], user_id: str
self, chat_id: str, message_id: str, file_ids: list[str], user_id: str, db: Optional[Session] = None
) -> Optional[list[ChatFileModel]]:
if not file_ids:
return None
@@ -1333,7 +1339,7 @@ class ChatTable:
chat_message_file_ids = [
item.id
for item in self.get_chat_files_by_chat_id_and_message_id(
chat_id, message_id
chat_id, message_id, db=db
)
]
# Remove duplicates and existing file_ids
@@ -1350,7 +1356,7 @@ class ChatTable:
return None
try:
with get_db() as db:
with get_db_context(db) as db:
now = int(time.time())
chat_files = [
@@ -1378,9 +1384,9 @@ class ChatTable:
return None
def get_chat_files_by_chat_id_and_message_id(
self, chat_id: str, message_id: str
self, chat_id: str, message_id: str, db: Optional[Session] = None
) -> list[ChatFileModel]:
with get_db() as db:
with get_db_context(db) as db:
all_chat_files = (
db.query(ChatFile)
.filter_by(chat_id=chat_id, message_id=message_id)
@@ -1391,17 +1397,17 @@ class ChatTable:
ChatFileModel.model_validate(chat_file) for chat_file in all_chat_files
]
def delete_chat_file(self, chat_id: str, file_id: str) -> bool:
def delete_chat_file(self, chat_id: str, file_id: str, db: Optional[Session] = None) -> bool:
try:
with get_db() as db:
with get_db_context(db) as db:
db.query(ChatFile).filter_by(chat_id=chat_id, file_id=file_id).delete()
db.commit()
return True
except Exception:
return False
def get_shared_chats_by_file_id(self, file_id: str) -> list[ChatModel]:
with get_db() as db:
def get_shared_chats_by_file_id(self, file_id: str, db: Optional[Session] = None) -> list[ChatModel]:
with get_db_context(db) as db:
# Join Chat and ChatFile tables to get shared chats associated with the file_id
all_chats = (
db.query(Chat)