OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-03 11:16:48 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy J. Baek 2024-07-03 23:32:39 -07:00
parent 1b65df3acc
commit 864646094e
11 changed files with 789 additions and 616 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
backend/apps/webui/models/auths.py
View File

@ -7,7 +7,7 @@ from sqlalchemy import String, Column, Boolean, Text
from apps.webui.models.users import UserModel, Users from apps.webui.models.users import UserModel, Users
from utils.utils import verify_password from utils.utils import verify_password
from apps.webui.internal.db import Base, Session from apps.webui.internal.db import Base, get_db
from config import SRC_LOG_LEVELS from config import SRC_LOG_LEVELS
@ -110,14 +110,14 @@ class AuthsTable:
**{"id": id, "email": email, "password": password, "active": True} **{"id": id, "email": email, "password": password, "active": True}
) )
result = Auth(**auth.model_dump()) result = Auth(**auth.model_dump())
Session.add(result) db.add(result)
user = Users.insert_new_user( user = Users.insert_new_user(
id, name, email, profile_image_url, role, oauth_sub id, name, email, profile_image_url, role, oauth_sub
) )
Session.commit() db.commit()
Session.refresh(result) db.refresh(result)
if result and user: if result and user:
return user return user
@ -127,7 +127,7 @@ class AuthsTable:
def authenticate_user(self, email: str, password: str) -> Optional[UserModel]: def authenticate_user(self, email: str, password: str) -> Optional[UserModel]:
log.info(f"authenticate_user: {email}") log.info(f"authenticate_user: {email}")
try: try:
auth = Session.query(Auth).filter_by(email=email, active=True).first() auth = db.query(Auth).filter_by(email=email, active=True).first()
if auth: if auth:
if verify_password(password, auth.password): if verify_password(password, auth.password):
user = Users.get_user_by_id(auth.id) user = Users.get_user_by_id(auth.id)
@ -154,7 +154,7 @@ class AuthsTable:
def authenticate_user_by_trusted_header(self, email: str) -> Optional[UserModel]: def authenticate_user_by_trusted_header(self, email: str) -> Optional[UserModel]:
log.info(f"authenticate_user_by_trusted_header: {email}") log.info(f"authenticate_user_by_trusted_header: {email}")
try: try:
auth = Session.query(Auth).filter(email=email, active=True).first() auth = db.query(Auth).filter(email=email, active=True).first()
if auth: if auth:
user = Users.get_user_by_id(auth.id) user = Users.get_user_by_id(auth.id)
return user return user
@ -163,16 +163,14 @@ class AuthsTable:
def update_user_password_by_id(self, id: str, new_password: str) -> bool: def update_user_password_by_id(self, id: str, new_password: str) -> bool:
try: try:
result = ( result = db.query(Auth).filter_by(id=id).update({"password": new_password})
Session.query(Auth).filter_by(id=id).update({"password": new_password})
)
return True if result == 1 else False return True if result == 1 else False
except: except:
return False return False
def update_email_by_id(self, id: str, email: str) -> bool: def update_email_by_id(self, id: str, email: str) -> bool:
try: try:
result = Session.query(Auth).filter_by(id=id).update({"email": email}) result = db.query(Auth).filter_by(id=id).update({"email": email})
return True if result == 1 else False return True if result == 1 else False
except: except:
return False return False
@ -183,7 +181,7 @@ class AuthsTable:
result = Users.delete_user_by_id(id) result = Users.delete_user_by_id(id)
if result: if result:
Session.query(Auth).filter_by(id=id).delete() db.query(Auth).filter_by(id=id).delete()
return True return True
else: else:

322
backend/apps/webui/models/chats.py
View File

@ -7,7 +7,7 @@ import time
from sqlalchemy import Column, String, BigInteger, Boolean, Text from sqlalchemy import Column, String, BigInteger, Boolean, Text
from apps.webui.internal.db import Base, Session from apps.webui.internal.db import Base, get_db
#################### ####################
@ -79,87 +79,99 @@ class ChatTitleIdResponse(BaseModel):
class ChatTable: class ChatTable:
def insert_new_chat(self, user_id: str, form_data: ChatForm) -> Optional[ChatModel]: def insert_new_chat(self, user_id: str, form_data: ChatForm) -> Optional[ChatModel]:
id = str(uuid.uuid4()) with get_db() as db:
chat = ChatModel(
**{
"id": id,
"user_id": user_id,
"title": (
form_data.chat["title"] if "title" in form_data.chat else "New Chat"
),
"chat": json.dumps(form_data.chat),
"created_at": int(time.time()),
"updated_at": int(time.time()),
}
)
result = Chat(**chat.model_dump()) id = str(uuid.uuid4())
Session.add(result) chat = ChatModel(
Session.commit() **{
Session.refresh(result) "id": id,
return ChatModel.model_validate(result) if result else None "user_id": user_id,
"title": (
form_data.chat["title"]
if "title" in form_data.chat
else "New Chat"
),
"chat": json.dumps(form_data.chat),
"created_at": int(time.time()),
"updated_at": int(time.time()),
}
)
result = Chat(**chat.model_dump())
db.add(result)
db.commit()
db.refresh(result)
return ChatModel.model_validate(result) if result else None
def update_chat_by_id(self, id: str, chat: dict) -> Optional[ChatModel]: def update_chat_by_id(self, id: str, chat: dict) -> Optional[ChatModel]:
try: try:
chat_obj = Session.get(Chat, id) with get_db() as db:
chat_obj.chat = json.dumps(chat)
chat_obj.title = chat["title"] if "title" in chat else "New Chat"
chat_obj.updated_at = int(time.time())
Session.commit()
Session.refresh(chat_obj)
return ChatModel.model_validate(chat_obj) chat_obj = db.get(Chat, id)
chat_obj.chat = json.dumps(chat)
chat_obj.title = chat["title"] if "title" in chat else "New Chat"
chat_obj.updated_at = int(time.time())
db.commit()
db.refresh(chat_obj)
return ChatModel.model_validate(chat_obj)
except Exception as e: except Exception as e:
return None return None
def insert_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]: def insert_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
# Get the existing chat to share with get_db() as db:
chat = Session.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")
# Create a new chat with the same data, but with a new ID
shared_chat = ChatModel(
**{
"id": str(uuid.uuid4()),
"user_id": f"shared-{chat_id}",
"title": chat.title,
"chat": chat.chat,
"created_at": chat.created_at,
"updated_at": int(time.time()),
}
)
shared_result = Chat(**shared_chat.model_dump())
Session.add(shared_result)
Session.commit()
Session.refresh(shared_result)
# Update the original chat with the share_id
result = (
Session.query(Chat)
.filter_by(id=chat_id)
.update({"share_id": shared_chat.id})
)
return shared_chat if (shared_result and result) else None # 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")
# Create a new chat with the same data, but with a new ID
shared_chat = ChatModel(
**{
"id": str(uuid.uuid4()),
"user_id": f"shared-{chat_id}",
"title": chat.title,
"chat": chat.chat,
"created_at": chat.created_at,
"updated_at": int(time.time()),
}
)
shared_result = Chat(**shared_chat.model_dump())
db.add(shared_result)
db.commit()
db.refresh(shared_result)
# Update the original chat with the share_id
result = (
db.query(Chat)
.filter_by(id=chat_id)
.update({"share_id": shared_chat.id})
)
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) -> Optional[ChatModel]:
try: try:
print("update_shared_chat_by_id") with get_db() as db:
chat = Session.get(Chat, chat_id)
print(chat)
chat.title = chat.title
chat.chat = chat.chat
Session.commit()
Session.refresh(chat)
return self.get_chat_by_id(chat.share_id) print("update_shared_chat_by_id")
chat = db.get(Chat, chat_id)
print(chat)
chat.title = chat.title
chat.chat = chat.chat
db.commit()
db.refresh(chat)
return self.get_chat_by_id(chat.share_id)
except: except:
return None 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) -> bool:
try: try:
Session.query(Chat).filter_by(user_id=f"shared-{chat_id}").delete() with get_db() as db:
return True
db.query(Chat).filter_by(user_id=f"shared-{chat_id}").delete()
return True
except: except:
return False return False
@ -167,42 +179,50 @@ class ChatTable:
self, id: str, share_id: Optional[str] self, id: str, share_id: Optional[str]
) -> Optional[ChatModel]: ) -> Optional[ChatModel]:
try: try:
chat = Session.get(Chat, id) with get_db() as db:
chat.share_id = share_id
Session.commit() chat = db.get(Chat, id)
Session.refresh(chat) chat.share_id = share_id
return ChatModel.model_validate(chat) db.commit()
db.refresh(chat)
return ChatModel.model_validate(chat)
except: except:
return None return None
def toggle_chat_archive_by_id(self, id: str) -> Optional[ChatModel]: def toggle_chat_archive_by_id(self, id: str) -> Optional[ChatModel]:
try: try:
chat = Session.get(Chat, id) with get_db() as db:
chat.archived = not chat.archived
Session.commit() chat = db.get(Chat, id)
Session.refresh(chat) chat.archived = not chat.archived
return ChatModel.model_validate(chat) db.commit()
db.refresh(chat)
return ChatModel.model_validate(chat)
except: except:
return None 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) -> bool:
try: try:
Session.query(Chat).filter_by(user_id=user_id).update({"archived": True}) with get_db() as db:
return True
db.query(Chat).filter_by(user_id=user_id).update({"archived": True})
return True
except: except:
return False return False
def get_archived_chat_list_by_user_id( def get_archived_chat_list_by_user_id(
self, user_id: str, skip: int = 0, limit: int = 50 self, user_id: str, skip: int = 0, limit: int = 50
) -> List[ChatModel]: ) -> List[ChatModel]:
all_chats = ( with get_db() as db:
Session.query(Chat)
.filter_by(user_id=user_id, archived=True) all_chats = (
.order_by(Chat.updated_at.desc()) db.query(Chat)
# .limit(limit).offset(skip) .filter_by(user_id=user_id, archived=True)
.all() .order_by(Chat.updated_at.desc())
) # .limit(limit).offset(skip)
return [ChatModel.model_validate(chat) for chat in all_chats] .all()
)
return [ChatModel.model_validate(chat) for chat in all_chats]
def get_chat_list_by_user_id( def get_chat_list_by_user_id(
self, self,
@ -211,110 +231,136 @@ class ChatTable:
skip: int = 0, skip: int = 0,
limit: int = 50, limit: int = 50,
) -> List[ChatModel]: ) -> List[ChatModel]:
query = Session.query(Chat).filter_by(user_id=user_id) with get_db() as db:
if not include_archived: query = db.query(Chat).filter_by(user_id=user_id)
query = query.filter_by(archived=False) if not include_archived:
all_chats = ( query = query.filter_by(archived=False)
query.order_by(Chat.updated_at.desc()) all_chats = (
# .limit(limit).offset(skip) query.order_by(Chat.updated_at.desc())
.all() # .limit(limit).offset(skip)
) .all()
return [ChatModel.model_validate(chat) for chat in all_chats] )
return [ChatModel.model_validate(chat) for chat in all_chats]
def get_chat_list_by_chat_ids( 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
) -> List[ChatModel]: ) -> List[ChatModel]:
all_chats = (
Session.query(Chat) with get_db() as db:
.filter(Chat.id.in_(chat_ids))
.filter_by(archived=False) all_chats = (
.order_by(Chat.updated_at.desc()) db.query(Chat)
.all() .filter(Chat.id.in_(chat_ids))
) .filter_by(archived=False)
return [ChatModel.model_validate(chat) for chat in all_chats] .order_by(Chat.updated_at.desc())
.all()
)
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) -> Optional[ChatModel]:
try: try:
chat = Session.get(Chat, id) with get_db() as db:
return ChatModel.model_validate(chat)
chat = db.get(Chat, id)
return ChatModel.model_validate(chat)
except: except:
return None return None
def get_chat_by_share_id(self, id: str) -> Optional[ChatModel]: def get_chat_by_share_id(self, id: str) -> Optional[ChatModel]:
try: try:
chat = Session.query(Chat).filter_by(share_id=id).first() with get_db() as db:
if chat: chat = db.query(Chat).filter_by(share_id=id).first()
return self.get_chat_by_id(id)
else: if chat:
return None return self.get_chat_by_id(id)
else:
return None
except Exception as e: except Exception as e:
return None 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) -> Optional[ChatModel]:
try: try:
chat = Session.query(Chat).filter_by(id=id, user_id=user_id).first() with get_db() as db:
return ChatModel.model_validate(chat)
chat = db.query(Chat).filter_by(id=id, user_id=user_id).first()
return ChatModel.model_validate(chat)
except: except:
return None return None
def get_chats(self, skip: int = 0, limit: int = 50) -> List[ChatModel]: def get_chats(self, skip: int = 0, limit: int = 50) -> List[ChatModel]:
all_chats = ( with get_db() as db:
Session.query(Chat)
# .limit(limit).offset(skip) all_chats = (
.order_by(Chat.updated_at.desc()) db.query(Chat)
) # .limit(limit).offset(skip)
return [ChatModel.model_validate(chat) for chat in all_chats] .order_by(Chat.updated_at.desc())
)
return [ChatModel.model_validate(chat) for chat in all_chats]
def get_chats_by_user_id(self, user_id: str) -> List[ChatModel]: def get_chats_by_user_id(self, user_id: str) -> List[ChatModel]:
all_chats = ( with get_db() as db:
Session.query(Chat)
.filter_by(user_id=user_id) all_chats = (
.order_by(Chat.updated_at.desc()) db.query(Chat)
) .filter_by(user_id=user_id)
return [ChatModel.model_validate(chat) for chat in all_chats] .order_by(Chat.updated_at.desc())
)
return [ChatModel.model_validate(chat) for chat in all_chats]
def get_archived_chats_by_user_id(self, user_id: str) -> List[ChatModel]: def get_archived_chats_by_user_id(self, user_id: str) -> List[ChatModel]:
all_chats = ( with get_db() as db:
Session.query(Chat)
.filter_by(user_id=user_id, archived=True) all_chats = (
.order_by(Chat.updated_at.desc()) db.query(Chat)
) .filter_by(user_id=user_id, archived=True)
return [ChatModel.model_validate(chat) for chat in all_chats] .order_by(Chat.updated_at.desc())
)
return [ChatModel.model_validate(chat) for chat in all_chats]
def delete_chat_by_id(self, id: str) -> bool: def delete_chat_by_id(self, id: str) -> bool:
try: try:
Session.query(Chat).filter_by(id=id).delete() with get_db() as db:
return True and self.delete_shared_chat_by_chat_id(id) db.query(Chat).filter_by(id=id).delete()
return True and self.delete_shared_chat_by_chat_id(id)
except: except:
return False 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) -> bool:
try: try:
Session.query(Chat).filter_by(id=id, user_id=user_id).delete() with get_db() as db:
return True and self.delete_shared_chat_by_chat_id(id) db.query(Chat).filter_by(id=id, user_id=user_id).delete()
return True and self.delete_shared_chat_by_chat_id(id)
except: except:
return False return False
def delete_chats_by_user_id(self, user_id: str) -> bool: def delete_chats_by_user_id(self, user_id: str) -> bool:
try: try:
self.delete_shared_chats_by_user_id(user_id)
Session.query(Chat).filter_by(user_id=user_id).delete() with get_db() as db:
return True
self.delete_shared_chats_by_user_id(user_id)
db.query(Chat).filter_by(user_id=user_id).delete()
return True
except: except:
return False 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) -> bool:
try: try:
chats_by_user = Session.query(Chat).filter_by(user_id=user_id).all()
shared_chat_ids = [f"shared-{chat.id}" for chat in chats_by_user]
Session.query(Chat).filter(Chat.user_id.in_(shared_chat_ids)).delete() with get_db() as db:
return True 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]
db.query(Chat).filter(Chat.user_id.in_(shared_chat_ids)).delete()
return True
except: except:
return False return False

96
backend/apps/webui/models/documents.py
View File

@ -5,7 +5,7 @@ import logging
from sqlalchemy import String, Column, BigInteger, Text from sqlalchemy import String, Column, BigInteger, Text
from apps.webui.internal.db import Base, Session from apps.webui.internal.db import Base, get_db
import json import json
@ -74,51 +74,59 @@ class DocumentsTable:
def insert_new_doc( def insert_new_doc(
self, user_id: str, form_data: DocumentForm self, user_id: str, form_data: DocumentForm
) -> Optional[DocumentModel]: ) -> Optional[DocumentModel]:
document = DocumentModel( with get_db() as db:
**{
**form_data.model_dump(),
"user_id": user_id,
"timestamp": int(time.time()),
}
)
try: document = DocumentModel(
result = Document(**document.model_dump()) **{
Session.add(result) **form_data.model_dump(),
Session.commit() "user_id": user_id,
Session.refresh(result) "timestamp": int(time.time()),
if result: }
return DocumentModel.model_validate(result) )
else:
try:
result = Document(**document.model_dump())
db.add(result)
db.commit()
db.refresh(result)
if result:
return DocumentModel.model_validate(result)
else:
return None
except:
return None return None
except:
return None
def get_doc_by_name(self, name: str) -> Optional[DocumentModel]: def get_doc_by_name(self, name: str) -> Optional[DocumentModel]:
try: try:
document = Session.query(Document).filter_by(name=name).first() with get_db() as db:
return DocumentModel.model_validate(document) if document else None
document = db.query(Document).filter_by(name=name).first()
return DocumentModel.model_validate(document) if document else None
except: except:
return None return None
def get_docs(self) -> List[DocumentModel]: def get_docs(self) -> List[DocumentModel]:
return [ with get_db() as db:
DocumentModel.model_validate(doc) for doc in Session.query(Document).all()
] return [
DocumentModel.model_validate(doc) for doc in db.query(Document).all()
]
def update_doc_by_name( def update_doc_by_name(
self, name: str, form_data: DocumentUpdateForm self, name: str, form_data: DocumentUpdateForm
) -> Optional[DocumentModel]: ) -> Optional[DocumentModel]:
try: try:
Session.query(Document).filter_by(name=name).update( with get_db() as db:
{
"title": form_data.title, db.query(Document).filter_by(name=name).update(
"name": form_data.name, {
"timestamp": int(time.time()), "title": form_data.title,
} "name": form_data.name,
) "timestamp": int(time.time()),
Session.commit() }
return self.get_doc_by_name(form_data.name) )
db.commit()
return self.get_doc_by_name(form_data.name)
except Exception as e: except Exception as e:
log.exception(e) log.exception(e)
return None return None
@ -131,22 +139,26 @@ class DocumentsTable:
doc_content = json.loads(doc.content if doc.content else "{}") doc_content = json.loads(doc.content if doc.content else "{}")
doc_content = {**doc_content, **updated} doc_content = {**doc_content, **updated}
Session.query(Document).filter_by(name=name).update( with get_db() as db:
{
"content": json.dumps(doc_content), db.query(Document).filter_by(name=name).update(
"timestamp": int(time.time()), {
} "content": json.dumps(doc_content),
) "timestamp": int(time.time()),
Session.commit() }
return self.get_doc_by_name(name) )
db.commit()
return self.get_doc_by_name(name)
except Exception as e: except Exception as e:
log.exception(e) log.exception(e)
return None return None
def delete_doc_by_name(self, name: str) -> bool: def delete_doc_by_name(self, name: str) -> bool:
try: try:
Session.query(Document).filter_by(name=name).delete() with get_db() as db:
return True
db.query(Document).filter_by(name=name).delete()
return True
except: except:
return False return False

82
backend/apps/webui/models/files.py
View File

@ -5,7 +5,7 @@ import logging
from sqlalchemy import Column, String, BigInteger, Text from sqlalchemy import Column, String, BigInteger, Text
from apps.webui.internal.db import JSONField, Base, Session from apps.webui.internal.db import JSONField, Base, get_db
import json import json
@ -61,50 +61,62 @@ class FileForm(BaseModel):
class FilesTable: class FilesTable:
def insert_new_file(self, user_id: str, form_data: FileForm) -> Optional[FileModel]: def insert_new_file(self, user_id: str, form_data: FileForm) -> Optional[FileModel]:
file = FileModel( with get_db() as db:
**{
**form_data.model_dump(),
"user_id": user_id,
"created_at": int(time.time()),
}
)
try: file = FileModel(
result = File(**file.model_dump()) **{
Session.add(result) **form_data.model_dump(),
Session.commit() "user_id": user_id,
Session.refresh(result) "created_at": int(time.time()),
if result: }
return FileModel.model_validate(result) )
else:
try:
result = File(**file.model_dump())
db.add(result)
db.commit()
db.refresh(result)
if result:
return FileModel.model_validate(result)
else:
return None
except Exception as e:
print(f"Error creating tool: {e}")
return None return None
except Exception as e:
print(f"Error creating tool: {e}")
return None
def get_file_by_id(self, id: str) -> Optional[FileModel]: def get_file_by_id(self, id: str) -> Optional[FileModel]:
try: with get_db() as db:
file = Session.get(File, id)
return FileModel.model_validate(file) try:
except: file = db.get(File, id)
return None return FileModel.model_validate(file)
except:
return None
def get_files(self) -> List[FileModel]: def get_files(self) -> List[FileModel]:
return [FileModel.model_validate(file) for file in Session.query(File).all()] with get_db() as db:
return [FileModel.model_validate(file) for file in db.query(File).all()]
def delete_file_by_id(self, id: str) -> bool: def delete_file_by_id(self, id: str) -> bool:
try:
Session.query(File).filter_by(id=id).delete() with get_db() as db:
return True
except: try:
return False db.query(File).filter_by(id=id).delete()
return True
except:
return False
def delete_all_files(self) -> bool: def delete_all_files(self) -> bool:
try:
Session.query(File).delete() with get_db() as db:
return True
except: try:
return False db.query(File).delete()
return True
except:
return False
Files = FilesTable() Files = FilesTable()

184
backend/apps/webui/models/functions.py
View File

@ -5,7 +5,7 @@ import logging
from sqlalchemy import Column, String, Text, BigInteger, Boolean from sqlalchemy import Column, String, Text, BigInteger, Boolean
from apps.webui.internal.db import JSONField, Base, Session from apps.webui.internal.db import JSONField, Base, get_db
from apps.webui.models.users import Users from apps.webui.models.users import Users
import json import json
@ -91,6 +91,7 @@ class FunctionsTable:
def insert_new_function( def insert_new_function(
self, user_id: str, type: str, form_data: FunctionForm self, user_id: str, type: str, form_data: FunctionForm
) -> Optional[FunctionModel]: ) -> Optional[FunctionModel]:
function = FunctionModel( function = FunctionModel(
**{ **{
**form_data.model_dump(), **form_data.model_dump(),
@ -102,85 +103,99 @@ class FunctionsTable:
) )
try: try:
result = Function(**function.model_dump()) with get_db() as db:
Session.add(result) result = Function(**function.model_dump())
Session.commit() db.add(result)
Session.refresh(result) db.commit()
if result: db.refresh(result)
return FunctionModel.model_validate(result) if result:
else: return FunctionModel.model_validate(result)
return None else:
return None
except Exception as e: except Exception as e:
print(f"Error creating tool: {e}") print(f"Error creating tool: {e}")
return None return None
def get_function_by_id(self, id: str) -> Optional[FunctionModel]: def get_function_by_id(self, id: str) -> Optional[FunctionModel]:
try: try:
function = Session.get(Function, id) with get_db() as db:
return FunctionModel.model_validate(function)
function = db.get(Function, id)
return FunctionModel.model_validate(function)
except: except:
return None return None
def get_functions(self, active_only=False) -> List[FunctionModel]: def get_functions(self, active_only=False) -> List[FunctionModel]:
if active_only: with get_db() as db:
return [
FunctionModel.model_validate(function) if active_only:
for function in Session.query(Function).filter_by(is_active=True).all() return [
] FunctionModel.model_validate(function)
else: for function in db.query(Function).filter_by(is_active=True).all()
return [ ]
FunctionModel.model_validate(function) else:
for function in Session.query(Function).all() return [
] FunctionModel.model_validate(function)
for function in db.query(Function).all()
]
def get_functions_by_type( def get_functions_by_type(
self, type: str, active_only=False self, type: str, active_only=False
) -> List[FunctionModel]: ) -> List[FunctionModel]:
if active_only: with get_db() as db:
return [
FunctionModel.model_validate(function) if active_only:
for function in Session.query(Function) return [
.filter_by(type=type, is_active=True) FunctionModel.model_validate(function)
.all() for function in db.query(Function)
] .filter_by(type=type, is_active=True)
else: .all()
return [ ]
FunctionModel.model_validate(function) else:
for function in Session.query(Function).filter_by(type=type).all() return [
] FunctionModel.model_validate(function)
for function in db.query(Function).filter_by(type=type).all()
]
def get_global_filter_functions(self) -> List[FunctionModel]: def get_global_filter_functions(self) -> List[FunctionModel]:
return [ with get_db() as db:
FunctionModel.model_validate(function)
for function in Session.query(Function) return [
.filter_by(type="filter", is_active=True, is_global=True) FunctionModel.model_validate(function)
.all() for function in db.query(Function)
] .filter_by(type="filter", is_active=True, is_global=True)
.all()
]
def get_function_valves_by_id(self, id: str) -> Optional[dict]: def get_function_valves_by_id(self, id: str) -> Optional[dict]:
try: with get_db() as db:
function = Session.get(Function, id)
return function.valves if function.valves else {} try:
except Exception as e: function = db.get(Function, id)
print(f"An error occurred: {e}") return function.valves if function.valves else {}
return None except Exception as e:
print(f"An error occurred: {e}")
return None
def update_function_valves_by_id( def update_function_valves_by_id(
self, id: str, valves: dict self, id: str, valves: dict
) -> Optional[FunctionValves]: ) -> Optional[FunctionValves]:
try: with get_db() as db:
function = Session.get(Function, id)
function.valves = valves try:
function.updated_at = int(time.time()) function = db.get(Function, id)
Session.commit() function.valves = valves
Session.refresh(function) function.updated_at = int(time.time())
return self.get_function_by_id(id) db.commit()
except: db.refresh(function)
return None return self.get_function_by_id(id)
except:
return None
def get_user_valves_by_id_and_user_id( def get_user_valves_by_id_and_user_id(
self, id: str, user_id: str self, id: str, user_id: str
) -> Optional[dict]: ) -> Optional[dict]:
try: try:
user = Users.get_user_by_id(user_id) user = Users.get_user_by_id(user_id)
user_settings = user.settings.model_dump() user_settings = user.settings.model_dump()
@ -199,6 +214,7 @@ class FunctionsTable:
def update_user_valves_by_id_and_user_id( def update_user_valves_by_id_and_user_id(
self, id: str, user_id: str, valves: dict self, id: str, user_id: str, valves: dict
) -> Optional[dict]: ) -> Optional[dict]:
try: try:
user = Users.get_user_by_id(user_id) user = Users.get_user_by_id(user_id)
user_settings = user.settings.model_dump() user_settings = user.settings.model_dump()
@ -220,37 +236,43 @@ class FunctionsTable:
return None return None
def update_function_by_id(self, id: str, updated: dict) -> Optional[FunctionModel]: def update_function_by_id(self, id: str, updated: dict) -> Optional[FunctionModel]:
try: with get_db() as db:
Session.query(Function).filter_by(id=id).update(
{ try:
**updated, db.query(Function).filter_by(id=id).update(
"updated_at": int(time.time()), {
} **updated,
) "updated_at": int(time.time()),
Session.commit() }
return self.get_function_by_id(id) )
except: db.commit()
return None return self.get_function_by_id(id)
except:
return None
def deactivate_all_functions(self) -> Optional[bool]: def deactivate_all_functions(self) -> Optional[bool]:
try: with get_db() as db:
Session.query(Function).update(
{ try:
"is_active": False, db.query(Function).update(
"updated_at": int(time.time()), {
} "is_active": False,
) "updated_at": int(time.time()),
Session.commit() }
return True )
except: db.commit()
return None return True
except:
return None
def delete_function_by_id(self, id: str) -> bool: def delete_function_by_id(self, id: str) -> bool:
try: with get_db() as db:
Session.query(Function).filter_by(id=id).delete()
return True try:
except: db.query(Function).filter_by(id=id).delete()
return False return True
except:
return False
Functions = FunctionsTable() Functions = FunctionsTable()

130
backend/apps/webui/models/memories.py
View File

@ -3,7 +3,7 @@ from typing import List, Union, Optional
from sqlalchemy import Column, String, BigInteger, Text from sqlalchemy import Column, String, BigInteger, Text
from apps.webui.internal.db import Base, Session from apps.webui.internal.db import Base, get_db
import time import time
import uuid import uuid
@ -45,82 +45,98 @@ class MemoriesTable:
user_id: str, user_id: str,
content: str, content: str,
) -> Optional[MemoryModel]: ) -> Optional[MemoryModel]:
id = str(uuid.uuid4())
memory = MemoryModel( with get_db() as db:
**{ id = str(uuid.uuid4())
"id": id,
"user_id": user_id, memory = MemoryModel(
"content": content, **{
"created_at": int(time.time()), "id": id,
"updated_at": int(time.time()), "user_id": user_id,
} "content": content,
) "created_at": int(time.time()),
result = Memory(**memory.model_dump()) "updated_at": int(time.time()),
Session.add(result) }
Session.commit() )
Session.refresh(result) result = Memory(**memory.model_dump())
if result: db.add(result)
return MemoryModel.model_validate(result) db.commit()
else: db.refresh(result)
return None if result:
return MemoryModel.model_validate(result)
else:
return None
def update_memory_by_id( def update_memory_by_id(
self, self,
id: str, id: str,
content: str, content: str,
) -> Optional[MemoryModel]: ) -> Optional[MemoryModel]:
try: with get_db() as db:
Session.query(Memory).filter_by(id=id).update(
{"content": content, "updated_at": int(time.time())} try:
) db.query(Memory).filter_by(id=id).update(
Session.commit() {"content": content, "updated_at": int(time.time())}
return self.get_memory_by_id(id) )
except: db.commit()
return None return self.get_memory_by_id(id)
except:
return None
def get_memories(self) -> List[MemoryModel]: def get_memories(self) -> List[MemoryModel]:
try: with get_db() as db:
memories = Session.query(Memory).all()
return [MemoryModel.model_validate(memory) for memory in memories] try:
except: memories = db.query(Memory).all()
return None return [MemoryModel.model_validate(memory) for memory in memories]
except:
return None
def get_memories_by_user_id(self, user_id: str) -> List[MemoryModel]: def get_memories_by_user_id(self, user_id: str) -> List[MemoryModel]:
try: with get_db() as db:
memories = Session.query(Memory).filter_by(user_id=user_id).all()
return [MemoryModel.model_validate(memory) for memory in memories] try:
except: memories = db.query(Memory).filter_by(user_id=user_id).all()
return None return [MemoryModel.model_validate(memory) for memory in memories]
except:
return None
def get_memory_by_id(self, id: str) -> Optional[MemoryModel]: def get_memory_by_id(self, id: str) -> Optional[MemoryModel]:
try: with get_db() as db:
memory = Session.get(Memory, id)
return MemoryModel.model_validate(memory) try:
except: memory = db.get(Memory, id)
return None return MemoryModel.model_validate(memory)
except:
return None
def delete_memory_by_id(self, id: str) -> bool: def delete_memory_by_id(self, id: str) -> bool:
try: with get_db() as db:
Session.query(Memory).filter_by(id=id).delete()
return True
except: try:
return False db.query(Memory).filter_by(id=id).delete()
return True
except:
return False
def delete_memories_by_user_id(self, user_id: str) -> bool: def delete_memories_by_user_id(self, user_id: str) -> bool:
try: with get_db() as db:
Session.query(Memory).filter_by(user_id=user_id).delete()
return True try:
except: db.query(Memory).filter_by(user_id=user_id).delete()
return False return True
except:
return False
def delete_memory_by_id_and_user_id(self, id: str, user_id: str) -> bool: def delete_memory_by_id_and_user_id(self, id: str, user_id: str) -> bool:
try: with get_db() as db:
Session.query(Memory).filter_by(id=id, user_id=user_id).delete()
return True try:
except: db.query(Memory).filter_by(id=id, user_id=user_id).delete()
return False return True
except:
return False
Memories = MemoriesTable() Memories = MemoriesTable()

53
backend/apps/webui/models/models.py
View File

@ -5,7 +5,7 @@ from typing import Optional
from pydantic import BaseModel, ConfigDict from pydantic import BaseModel, ConfigDict
from sqlalchemy import String, Column, BigInteger, Text from sqlalchemy import String, Column, BigInteger, Text
from apps.webui.internal.db import Base, JSONField, Session from apps.webui.internal.db import Base, JSONField, get_db
from typing import List, Union, Optional from typing import List, Union, Optional
from config import SRC_LOG_LEVELS from config import SRC_LOG_LEVELS
@ -126,39 +126,46 @@ class ModelsTable:
} }
) )
try: try:
result = Model(**model.model_dump())
Session.add(result)
Session.commit()
Session.refresh(result)
if result: with get_db() as db:
return ModelModel.model_validate(result)
else: result = Model(**model.model_dump())
return None db.add(result)
db.commit()
db.refresh(result)
if result:
return ModelModel.model_validate(result)
else:
return None
except Exception as e: except Exception as e:
print(e) print(e)
return None return None
def get_all_models(self) -> List[ModelModel]: def get_all_models(self) -> List[ModelModel]:
return [ with get_db() as db:
ModelModel.model_validate(model) for model in Session.query(Model).all()
] return [ModelModel.model_validate(model) for model in db.query(Model).all()]
def get_model_by_id(self, id: str) -> Optional[ModelModel]: def get_model_by_id(self, id: str) -> Optional[ModelModel]:
try: try:
model = Session.get(Model, id) with get_db() as db:
return ModelModel.model_validate(model)
model = db.get(Model, id)
return ModelModel.model_validate(model)
except: except:
return None return None
def update_model_by_id(self, id: str, model: ModelForm) -> Optional[ModelModel]: def update_model_by_id(self, id: str, model: ModelForm) -> Optional[ModelModel]:
try: try:
# update only the fields that are present in the model with get_db() as db:
model = Session.query(Model).get(id)
model.update(**model.model_dump()) # update only the fields that are present in the model
Session.commit() model = db.query(Model).get(id)
Session.refresh(model) model.update(**model.model_dump())
return ModelModel.model_validate(model) db.commit()
db.refresh(model)
return ModelModel.model_validate(model)
except Exception as e: except Exception as e:
print(e) print(e)
@ -166,8 +173,10 @@ class ModelsTable:
def delete_model_by_id(self, id: str) -> bool: def delete_model_by_id(self, id: str) -> bool:
try: try:
Session.query(Model).filter_by(id=id).delete() with get_db() as db:
return True
db.query(Model).filter_by(id=id).delete()
return True
except: except:
return False return False

54
backend/apps/webui/models/prompts.py
View File

@ -4,7 +4,7 @@ import time
from sqlalchemy import String, Column, BigInteger, Text from sqlalchemy import String, Column, BigInteger, Text
from apps.webui.internal.db import Base, Session from apps.webui.internal.db import Base, get_db
import json import json
@ -60,46 +60,56 @@ class PromptsTable:
) )
try: try:
result = Prompt(**prompt.dict()) with get_db() as db:
Session.add(result)
Session.commit() result = Prompt(**prompt.dict())
Session.refresh(result) db.add(result)
if result: db.commit()
return PromptModel.model_validate(result) db.refresh(result)
else: if result:
return None return PromptModel.model_validate(result)
else:
return None
except Exception as e: except Exception as e:
return None return None
def get_prompt_by_command(self, command: str) -> Optional[PromptModel]: def get_prompt_by_command(self, command: str) -> Optional[PromptModel]:
try: try:
prompt = Session.query(Prompt).filter_by(command=command).first() with get_db() as db:
return PromptModel.model_validate(prompt)
prompt = db.query(Prompt).filter_by(command=command).first()
return PromptModel.model_validate(prompt)
except: except:
return None return None
def get_prompts(self) -> List[PromptModel]: def get_prompts(self) -> List[PromptModel]:
return [ with get_db() as db:
PromptModel.model_validate(prompt) for prompt in Session.query(Prompt).all()
] return [
PromptModel.model_validate(prompt) for prompt in db.query(Prompt).all()
]
def update_prompt_by_command( def update_prompt_by_command(
self, command: str, form_data: PromptForm self, command: str, form_data: PromptForm
) -> Optional[PromptModel]: ) -> Optional[PromptModel]:
try: try:
prompt = Session.query(Prompt).filter_by(command=command).first() with get_db() as db:
prompt.title = form_data.title
prompt.content = form_data.content prompt = db.query(Prompt).filter_by(command=command).first()
prompt.timestamp = int(time.time()) prompt.title = form_data.title
Session.commit() prompt.content = form_data.content
return PromptModel.model_validate(prompt) prompt.timestamp = int(time.time())
db.commit()
return PromptModel.model_validate(prompt)
except: except:
return None return None
def delete_prompt_by_command(self, command: str) -> bool: def delete_prompt_by_command(self, command: str) -> bool:
try: try:
Session.query(Prompt).filter_by(command=command).delete() with get_db() as db:
return True
db.query(Prompt).filter_by(command=command).delete()
return True
except: except:
return False return False

212
backend/apps/webui/models/tags.py
View File

@ -8,7 +8,7 @@ import logging
from sqlalchemy import String, Column, BigInteger, Text from sqlalchemy import String, Column, BigInteger, Text
from apps.webui.internal.db import Base, Session from apps.webui.internal.db import Base, get_db
from config import SRC_LOG_LEVELS from config import SRC_LOG_LEVELS
@ -79,26 +79,29 @@ class ChatTagsResponse(BaseModel):
class TagTable: class TagTable:
def insert_new_tag(self, name: str, user_id: str) -> Optional[TagModel]: def insert_new_tag(self, name: str, user_id: str) -> Optional[TagModel]:
id = str(uuid.uuid4()) with get_db() as db:
tag = TagModel(**{"id": id, "user_id": user_id, "name": name})
try: id = str(uuid.uuid4())
result = Tag(**tag.model_dump()) tag = TagModel(**{"id": id, "user_id": user_id, "name": name})
Session.add(result) try:
Session.commit() result = Tag(**tag.model_dump())
Session.refresh(result) db.add(result)
if result: db.commit()
return TagModel.model_validate(result) db.refresh(result)
else: if result:
return TagModel.model_validate(result)
else:
return None
except Exception as e:
return None return None
except Exception as e:
return None
def get_tag_by_name_and_user_id( def get_tag_by_name_and_user_id(
self, name: str, user_id: str self, name: str, user_id: str
) -> Optional[TagModel]: ) -> Optional[TagModel]:
try: try:
tag = Session.query(Tag).filter(name=name, user_id=user_id).first() with get_db() as db:
return TagModel.model_validate(tag) tag = db.query(Tag).filter(name=name, user_id=user_id).first()
return TagModel.model_validate(tag)
except Exception as e: except Exception as e:
return None return None
@ -120,98 +123,109 @@ class TagTable:
} }
) )
try: try:
result = ChatIdTag(**chatIdTag.model_dump()) with get_db() as db:
Session.add(result) result = ChatIdTag(**chatIdTag.model_dump())
Session.commit() db.add(result)
Session.refresh(result) db.commit()
if result: db.refresh(result)
return ChatIdTagModel.model_validate(result) if result:
else: return ChatIdTagModel.model_validate(result)
return None else:
return None
except: except:
return None return None
def get_tags_by_user_id(self, user_id: str) -> List[TagModel]: def get_tags_by_user_id(self, user_id: str) -> List[TagModel]:
tag_names = [ with get_db() as db:
chat_id_tag.tag_name tag_names = [
for chat_id_tag in ( chat_id_tag.tag_name
Session.query(ChatIdTag) for chat_id_tag in (
.filter_by(user_id=user_id) db.query(ChatIdTag)
.order_by(ChatIdTag.timestamp.desc()) .filter_by(user_id=user_id)
.all() .order_by(ChatIdTag.timestamp.desc())
) .all()
] )
]
return [ return [
TagModel.model_validate(tag) TagModel.model_validate(tag)
for tag in ( for tag in (
Session.query(Tag) db.query(Tag)
.filter_by(user_id=user_id) .filter_by(user_id=user_id)
.filter(Tag.name.in_(tag_names)) .filter(Tag.name.in_(tag_names))
.all() .all()
) )
] ]
def get_tags_by_chat_id_and_user_id( def get_tags_by_chat_id_and_user_id(
self, chat_id: str, user_id: str self, chat_id: str, user_id: str
) -> List[TagModel]: ) -> List[TagModel]:
tag_names = [ with get_db() as db:
chat_id_tag.tag_name
for chat_id_tag in (
Session.query(ChatIdTag)
.filter_by(user_id=user_id, chat_id=chat_id)
.order_by(ChatIdTag.timestamp.desc())
.all()
)
]
return [ tag_names = [
TagModel.model_validate(tag) chat_id_tag.tag_name
for tag in ( for chat_id_tag in (
Session.query(Tag) db.query(ChatIdTag)
.filter_by(user_id=user_id) .filter_by(user_id=user_id, chat_id=chat_id)
.filter(Tag.name.in_(tag_names)) .order_by(ChatIdTag.timestamp.desc())
.all() .all()
) )
] ]
return [
TagModel.model_validate(tag)
for tag in (
db.query(Tag)
.filter_by(user_id=user_id)
.filter(Tag.name.in_(tag_names))
.all()
)
]
def get_chat_ids_by_tag_name_and_user_id( def get_chat_ids_by_tag_name_and_user_id(
self, tag_name: str, user_id: str self, tag_name: str, user_id: str
) -> List[ChatIdTagModel]: ) -> List[ChatIdTagModel]:
return [ with get_db() as db:
ChatIdTagModel.model_validate(chat_id_tag)
for chat_id_tag in ( return [
Session.query(ChatIdTag) ChatIdTagModel.model_validate(chat_id_tag)
.filter_by(user_id=user_id, tag_name=tag_name) for chat_id_tag in (
.order_by(ChatIdTag.timestamp.desc()) db.query(ChatIdTag)
.all() .filter_by(user_id=user_id, tag_name=tag_name)
) .order_by(ChatIdTag.timestamp.desc())
] .all()
)
]
def count_chat_ids_by_tag_name_and_user_id( def count_chat_ids_by_tag_name_and_user_id(
self, tag_name: str, user_id: str self, tag_name: str, user_id: str
) -> int: ) -> int:
return ( with get_db() as db:
Session.query(ChatIdTag)
.filter_by(tag_name=tag_name, user_id=user_id) return (
.count() db.query(ChatIdTag)
) .filter_by(tag_name=tag_name, user_id=user_id)
.count()
)
def delete_tag_by_tag_name_and_user_id(self, tag_name: str, user_id: str) -> bool: def delete_tag_by_tag_name_and_user_id(self, tag_name: str, user_id: str) -> bool:
try: try:
res = ( with get_db() as db:
Session.query(ChatIdTag) res = (
.filter_by(tag_name=tag_name, user_id=user_id) db.query(ChatIdTag)
.delete() .filter_by(tag_name=tag_name, user_id=user_id)
) .delete()
log.debug(f"res: {res}") )
Session.commit() log.debug(f"res: {res}")
db.commit()
tag_count = self.count_chat_ids_by_tag_name_and_user_id(tag_name, user_id) tag_count = self.count_chat_ids_by_tag_name_and_user_id(
if tag_count == 0: tag_name, user_id
# Remove tag item from Tag col as well )
Session.query(Tag).filter_by(name=tag_name, user_id=user_id).delete() if tag_count == 0:
return True # Remove tag item from Tag col as well
db.query(Tag).filter_by(name=tag_name, user_id=user_id).delete()
return True
except Exception as e: except Exception as e:
log.error(f"delete_tag: {e}") log.error(f"delete_tag: {e}")
return False return False
@ -220,20 +234,24 @@ class TagTable:
self, tag_name: str, chat_id: str, user_id: str self, tag_name: str, chat_id: str, user_id: str
) -> bool: ) -> bool:
try: try:
res = ( with get_db() as db:
Session.query(ChatIdTag)
.filter_by(tag_name=tag_name, chat_id=chat_id, user_id=user_id)
.delete()
)
log.debug(f"res: {res}")
Session.commit()
tag_count = self.count_chat_ids_by_tag_name_and_user_id(tag_name, user_id) res = (
if tag_count == 0: db.query(ChatIdTag)
# Remove tag item from Tag col as well .filter_by(tag_name=tag_name, chat_id=chat_id, user_id=user_id)
Session.query(Tag).filter_by(name=tag_name, user_id=user_id).delete() .delete()
)
log.debug(f"res: {res}")
db.commit()
return True tag_count = self.count_chat_ids_by_tag_name_and_user_id(
tag_name, user_id
)
if tag_count == 0:
# Remove tag item from Tag col as well
db.query(Tag).filter_by(name=tag_name, user_id=user_id).delete()
return True
except Exception as e: except Exception as e:
log.error(f"delete_tag: {e}") log.error(f"delete_tag: {e}")
return False return False

89
backend/apps/webui/models/tools.py
View File

@ -4,7 +4,7 @@ import time
import logging import logging
from sqlalchemy import String, Column, BigInteger, Text from sqlalchemy import String, Column, BigInteger, Text
from apps.webui.internal.db import Base, JSONField, Session from apps.webui.internal.db import Base, JSONField, get_db
from apps.webui.models.users import Users from apps.webui.models.users import Users
import json import json
@ -83,54 +83,63 @@ class ToolsTable:
def insert_new_tool( def insert_new_tool(
self, user_id: str, form_data: ToolForm, specs: List[dict] self, user_id: str, form_data: ToolForm, specs: List[dict]
) -> Optional[ToolModel]: ) -> Optional[ToolModel]:
tool = ToolModel(
**{
**form_data.model_dump(),
"specs": specs,
"user_id": user_id,
"updated_at": int(time.time()),
"created_at": int(time.time()),
}
)
try: with get_db() as db:
result = Tool(**tool.model_dump())
Session.add(result) tool = ToolModel(
Session.commit() **{
Session.refresh(result) **form_data.model_dump(),
if result: "specs": specs,
return ToolModel.model_validate(result) "user_id": user_id,
else: "updated_at": int(time.time()),
"created_at": int(time.time()),
}
)
try:
result = Tool(**tool.model_dump())
db.add(result)
db.commit()
db.refresh(result)
if result:
return ToolModel.model_validate(result)
else:
return None
except Exception as e:
print(f"Error creating tool: {e}")
return None return None
except Exception as e:
print(f"Error creating tool: {e}")
return None
def get_tool_by_id(self, id: str) -> Optional[ToolModel]: def get_tool_by_id(self, id: str) -> Optional[ToolModel]:
try: try:
tool = Session.get(Tool, id) with get_db() as db:
return ToolModel.model_validate(tool)
tool = db.get(Tool, id)
return ToolModel.model_validate(tool)
except: except:
return None return None
def get_tools(self) -> List[ToolModel]: def get_tools(self) -> List[ToolModel]:
return [ToolModel.model_validate(tool) for tool in Session.query(Tool).all()] return [ToolModel.model_validate(tool) for tool in db.query(Tool).all()]
def get_tool_valves_by_id(self, id: str) -> Optional[dict]: def get_tool_valves_by_id(self, id: str) -> Optional[dict]:
try: try:
tool = Session.get(Tool, id) with get_db() as db:
return tool.valves if tool.valves else {}
tool = db.get(Tool, id)
return tool.valves if tool.valves else {}
except Exception as e: except Exception as e:
print(f"An error occurred: {e}") print(f"An error occurred: {e}")
return None return None
def update_tool_valves_by_id(self, id: str, valves: dict) -> Optional[ToolValves]: def update_tool_valves_by_id(self, id: str, valves: dict) -> Optional[ToolValves]:
try: try:
Session.query(Tool).filter_by(id=id).update( with get_db() as db:
{"valves": valves, "updated_at": int(time.time())}
) db.query(Tool).filter_by(id=id).update(
Session.commit() {"valves": valves, "updated_at": int(time.time())}
return self.get_tool_by_id(id) )
db.commit()
return self.get_tool_by_id(id)
except: except:
return None return None
@ -177,19 +186,21 @@ class ToolsTable:
def update_tool_by_id(self, id: str, updated: dict) -> Optional[ToolModel]: def update_tool_by_id(self, id: str, updated: dict) -> Optional[ToolModel]:
try: try:
tool = Session.get(Tool, id) with get_db() as db:
tool.update(**updated) tool = db.get(Tool, id)
tool.updated_at = int(time.time()) tool.update(**updated)
Session.commit() tool.updated_at = int(time.time())
Session.refresh(tool) db.commit()
return ToolModel.model_validate(tool) db.refresh(tool)
return ToolModel.model_validate(tool)
except: except:
return None return None
def delete_tool_by_id(self, id: str) -> bool: def delete_tool_by_id(self, id: str) -> bool:
try: try:
Session.query(Tool).filter_by(id=id).delete() with get_db() as db:
return True db.query(Tool).filter_by(id=id).delete()
return True
except: except:
return False return False

163
backend/apps/webui/models/users.py
View File

@ -6,7 +6,7 @@ from sqlalchemy import String, Column, BigInteger, Text
from utils.misc import get_gravatar_url from utils.misc import get_gravatar_url
from apps.webui.internal.db import Base, JSONField, Session from apps.webui.internal.db import Base, JSONField, Session, get_db
from apps.webui.models.chats import Chats from apps.webui.models.chats import Chats
#################### ####################
@ -88,81 +88,92 @@ class UsersTable:
role: str = "pending", role: str = "pending",
oauth_sub: Optional[str] = None, oauth_sub: Optional[str] = None,
) -> Optional[UserModel]: ) -> Optional[UserModel]:
user = UserModel( with get_db() as db:
**{ user = UserModel(
"id": id, **{
"name": name, "id": id,
"email": email, "name": name,
"role": role, "email": email,
"profile_image_url": profile_image_url, "role": role,
"last_active_at": int(time.time()), "profile_image_url": profile_image_url,
"created_at": int(time.time()), "last_active_at": int(time.time()),
"updated_at": int(time.time()), "created_at": int(time.time()),
"oauth_sub": oauth_sub, "updated_at": int(time.time()),
} "oauth_sub": oauth_sub,
) }
result = User(**user.model_dump()) )
Session.add(result) result = User(**user.model_dump())
Session.commit() db.add(result)
Session.refresh(result) db.commit()
if result: db.refresh(result)
return user if result:
else: return user
return None else:
return None
def get_user_by_id(self, id: str) -> Optional[UserModel]: def get_user_by_id(self, id: str) -> Optional[UserModel]:
try: try:
user = Session.query(User).filter_by(id=id).first() with get_db() as db:
return UserModel.model_validate(user) user = db.query(User).filter_by(id=id).first()
return UserModel.model_validate(user)
except Exception as e: except Exception as e:
return None return None
def get_user_by_api_key(self, api_key: str) -> Optional[UserModel]: def get_user_by_api_key(self, api_key: str) -> Optional[UserModel]:
try: try:
user = Session.query(User).filter_by(api_key=api_key).first() with get_db() as db:
return UserModel.model_validate(user)
user = db.query(User).filter_by(api_key=api_key).first()
return UserModel.model_validate(user)
except: except:
return None return None
def get_user_by_email(self, email: str) -> Optional[UserModel]: def get_user_by_email(self, email: str) -> Optional[UserModel]:
try: try:
user = Session.query(User).filter_by(email=email).first() with get_db() as db:
return UserModel.model_validate(user)
user = db.query(User).filter_by(email=email).first()
return UserModel.model_validate(user)
except: except:
return None return None
def get_user_by_oauth_sub(self, sub: str) -> Optional[UserModel]: def get_user_by_oauth_sub(self, sub: str) -> Optional[UserModel]:
try: try:
user = Session.query(User).filter_by(oauth_sub=sub).first() with get_db() as db:
return UserModel.model_validate(user)
user = db.query(User).filter_by(oauth_sub=sub).first()
return UserModel.model_validate(user)
except: except:
return None return None
def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]: def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]:
users = ( with get_db() as db:
Session.query(User) users = (
# .offset(skip).limit(limit) db.query(User)
.all() # .offset(skip).limit(limit)
) .all()
return [UserModel.model_validate(user) for user in users] )
return [UserModel.model_validate(user) for user in users]
def get_num_users(self) -> Optional[int]: def get_num_users(self) -> Optional[int]:
return Session.query(User).count() with get_db() as db:
return db.query(User).count()
def get_first_user(self) -> UserModel: def get_first_user(self) -> UserModel:
try: try:
user = Session.query(User).order_by(User.created_at).first() with get_db() as db:
return UserModel.model_validate(user) user = db.query(User).order_by(User.created_at).first()
return UserModel.model_validate(user)
except: except:
return None return None
def update_user_role_by_id(self, id: str, role: str) -> Optional[UserModel]: def update_user_role_by_id(self, id: str, role: str) -> Optional[UserModel]:
try: try:
Session.query(User).filter_by(id=id).update({"role": role}) with get_db() as db:
Session.commit() db.query(User).filter_by(id=id).update({"role": role})
db.commit()
user = Session.query(User).filter_by(id=id).first() user = db.query(User).filter_by(id=id).first()
return UserModel.model_validate(user) return UserModel.model_validate(user)
except: except:
return None return None
@ -170,25 +181,28 @@ class UsersTable:
self, id: str, profile_image_url: str self, id: str, profile_image_url: str
) -> Optional[UserModel]: ) -> Optional[UserModel]:
try: try:
Session.query(User).filter_by(id=id).update( with get_db() as db:
{"profile_image_url": profile_image_url} db.query(User).filter_by(id=id).update(
) {"profile_image_url": profile_image_url}
Session.commit() )
db.commit()
user = Session.query(User).filter_by(id=id).first() user = db.query(User).filter_by(id=id).first()
return UserModel.model_validate(user) return UserModel.model_validate(user)
except: except:
return None return None
def update_user_last_active_by_id(self, id: str) -> Optional[UserModel]: def update_user_last_active_by_id(self, id: str) -> Optional[UserModel]:
try: try:
Session.query(User).filter_by(id=id).update( with get_db() as db:
{"last_active_at": int(time.time())}
)
Session.commit()
user = Session.query(User).filter_by(id=id).first() db.query(User).filter_by(id=id).update(
return UserModel.model_validate(user) {"last_active_at": int(time.time())}
)
db.commit()
user = db.query(User).filter_by(id=id).first()
return UserModel.model_validate(user)
except: except:
return None return None
@ -196,21 +210,23 @@ class UsersTable:
self, id: str, oauth_sub: str self, id: str, oauth_sub: str
) -> Optional[UserModel]: ) -> Optional[UserModel]:
try: try:
Session.query(User).filter_by(id=id).update({"oauth_sub": oauth_sub}) with get_db() as db:
db.query(User).filter_by(id=id).update({"oauth_sub": oauth_sub})
user = Session.query(User).filter_by(id=id).first() user = db.query(User).filter_by(id=id).first()
return UserModel.model_validate(user) return UserModel.model_validate(user)
except: except:
return None return None
def update_user_by_id(self, id: str, updated: dict) -> Optional[UserModel]: def update_user_by_id(self, id: str, updated: dict) -> Optional[UserModel]:
try: try:
Session.query(User).filter_by(id=id).update(updated) with get_db() as db:
Session.commit() db.query(User).filter_by(id=id).update(updated)
db.commit()
user = Session.query(User).filter_by(id=id).first() user = db.query(User).filter_by(id=id).first()
return UserModel.model_validate(user) return UserModel.model_validate(user)
# return UserModel(**user.dict()) # return UserModel(**user.dict())
except Exception as e: except Exception as e:
return None return None
@ -220,9 +236,10 @@ class UsersTable:
result = Chats.delete_chats_by_user_id(id) result = Chats.delete_chats_by_user_id(id)
if result: if result:
# Delete User with get_db() as db:
Session.query(User).filter_by(id=id).delete() # Delete User
Session.commit() db.query(User).filter_by(id=id).delete()
db.commit()
return True return True
else: else:
@ -232,16 +249,18 @@ class UsersTable:
def update_user_api_key_by_id(self, id: str, api_key: str) -> str: def update_user_api_key_by_id(self, id: str, api_key: str) -> str:
try: try:
result = Session.query(User).filter_by(id=id).update({"api_key": api_key}) with get_db() as db:
Session.commit() result = db.query(User).filter_by(id=id).update({"api_key": api_key})
return True if result == 1 else False db.commit()
return True if result == 1 else False
except: except:
return False return False
def get_user_api_key_by_id(self, id: str) -> Optional[str]: def get_user_api_key_by_id(self, id: str) -> Optional[str]:
try: try:
user = Session.query(User).filter_by(id=id).first() with get_db() as db:
return user.api_key user = db.query(User).filter_by(id=id).first()
return user.api_key
except Exception as e: except Exception as e:
return None return None