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

feat(sqlalchemy): Replace peewee with sqlalchemy

Browse Source
This commit is contained in:
Jonathan Rohde
2024-06-27 07:48:08 +02:00
parent 8dac2a2140
commit df09d0830a
47 changed files with 2580 additions and 1003 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/webui/routers/chats.py
+88 -58
View File
@@ -1,6 +1,8 @@
from fastapi import Depends, Request, HTTPException, status
from datetime import datetime, timedelta
from typing import List, Union, Optional
from apps.webui.internal.db import get_db
from utils.utils import get_current_user, get_admin_user
from fastapi import APIRouter
from pydantic import BaseModel
@@ -43,9 +45,9 @@ router = APIRouter()
@router.get("/", response_model=List[ChatTitleIdResponse])
@router.get("/list", response_model=List[ChatTitleIdResponse])
async def get_session_user_chat_list(
user=Depends(get_current_user), skip: int = 0, limit: int = 50
user=Depends(get_current_user), skip: int = 0, limit: int = 50, db=Depends(get_db)
):
return Chats.get_chat_list_by_user_id(user.id, skip, limit)
return Chats.get_chat_list_by_user_id(db, user.id, skip, limit)
############################
@@ -54,7 +56,9 @@ async def get_session_user_chat_list(
@router.delete("/", response_model=bool)
async def delete_all_user_chats(request: Request, user=Depends(get_current_user)):
async def delete_all_user_chats(
request: Request, user=Depends(get_current_user), db=Depends(get_db)
):
if (
user.role == "user"
@@ -65,7 +69,7 @@ async def delete_all_user_chats(request: Request, user=Depends(get_current_user)
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
)
result = Chats.delete_chats_by_user_id(user.id)
result = Chats.delete_chats_by_user_id(db, user.id)
return result
@@ -76,10 +80,14 @@ async def delete_all_user_chats(request: Request, user=Depends(get_current_user)
@router.get("/list/user/{user_id}", response_model=List[ChatTitleIdResponse])
async def get_user_chat_list_by_user_id(
user_id: str, user=Depends(get_admin_user), skip: int = 0, limit: int = 50
user_id: str,
user=Depends(get_admin_user),
skip: int = 0,
limit: int = 50,
db=Depends(get_db),
):
return Chats.get_chat_list_by_user_id(
user_id, include_archived=True, skip=skip, limit=limit
db, user_id, include_archived=True, skip=skip, limit=limit
)
@@ -89,9 +97,11 @@ async def get_user_chat_list_by_user_id(
@router.post("/new", response_model=Optional[ChatResponse])
async def create_new_chat(form_data: ChatForm, user=Depends(get_current_user)):
async def create_new_chat(
form_data: ChatForm, user=Depends(get_current_user), db=Depends(get_db)
):
try:
chat = Chats.insert_new_chat(user.id, form_data)
chat = Chats.insert_new_chat(db, user.id, form_data)
return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
except Exception as e:
log.exception(e)
@@ -106,10 +116,10 @@ async def create_new_chat(form_data: ChatForm, user=Depends(get_current_user)):
@router.get("/all", response_model=List[ChatResponse])
async def get_user_chats(user=Depends(get_current_user)):
async def get_user_chats(user=Depends(get_current_user), db=Depends(get_db)):
return [
ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
for chat in Chats.get_chats_by_user_id(user.id)
for chat in Chats.get_chats_by_user_id(db, user.id)
]
@@ -119,10 +129,10 @@ async def get_user_chats(user=Depends(get_current_user)):
@router.get("/all/archived", response_model=List[ChatResponse])
async def get_user_chats(user=Depends(get_current_user)):
async def get_user_archived_chats(user=Depends(get_current_user), db=Depends(get_db)):
return [
ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
for chat in Chats.get_archived_chats_by_user_id(user.id)
for chat in Chats.get_archived_chats_by_user_id(db, user.id)
]
@@ -132,7 +142,7 @@ async def get_user_chats(user=Depends(get_current_user)):
@router.get("/all/db", response_model=List[ChatResponse])
async def get_all_user_chats_in_db(user=Depends(get_admin_user)):
async def get_all_user_chats_in_db(user=Depends(get_admin_user), db=Depends(get_db)):
if not ENABLE_ADMIN_EXPORT:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
@@ -140,7 +150,7 @@ async def get_all_user_chats_in_db(user=Depends(get_admin_user)):
)
return [
ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
for chat in Chats.get_chats()
for chat in Chats.get_chats(db)
]
@@ -151,9 +161,9 @@ async def get_all_user_chats_in_db(user=Depends(get_admin_user)):
@router.get("/archived", response_model=List[ChatTitleIdResponse])
async def get_archived_session_user_chat_list(
user=Depends(get_current_user), skip: int = 0, limit: int = 50
user=Depends(get_current_user), skip: int = 0, limit: int = 50, db=Depends(get_db)
):
return Chats.get_archived_chat_list_by_user_id(user.id, skip, limit)
return Chats.get_archived_chat_list_by_user_id(db, user.id, skip, limit)
############################
@@ -162,8 +172,8 @@ async def get_archived_session_user_chat_list(
@router.post("/archive/all", response_model=bool)
async def archive_all_chats(user=Depends(get_current_user)):
return Chats.archive_all_chats_by_user_id(user.id)
async def archive_all_chats(user=Depends(get_current_user), db=Depends(get_db)):
return Chats.archive_all_chats_by_user_id(db, user.id)
############################
@@ -172,16 +182,18 @@ async def archive_all_chats(user=Depends(get_current_user)):
@router.get("/share/{share_id}", response_model=Optional[ChatResponse])
async def get_shared_chat_by_id(share_id: str, user=Depends(get_current_user)):
async def get_shared_chat_by_id(
share_id: str, user=Depends(get_current_user), db=Depends(get_db)
):
if user.role == "pending":
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.NOT_FOUND
)
if user.role == "user":
chat = Chats.get_chat_by_share_id(share_id)
chat = Chats.get_chat_by_share_id(db, share_id)
elif user.role == "admin":
chat = Chats.get_chat_by_id(share_id)
chat = Chats.get_chat_by_id(db, share_id)
if chat:
return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
@@ -204,21 +216,23 @@ class TagNameForm(BaseModel):
@router.post("/tags", response_model=List[ChatTitleIdResponse])
async def get_user_chat_list_by_tag_name(
form_data: TagNameForm, user=Depends(get_current_user)
form_data: TagNameForm, user=Depends(get_current_user), db=Depends(get_db)
):
print(form_data)
chat_ids = [
chat_id_tag.chat_id
for chat_id_tag in Tags.get_chat_ids_by_tag_name_and_user_id(
form_data.name, user.id
db, form_data.name, user.id
)
]
chats = Chats.get_chat_list_by_chat_ids(chat_ids, form_data.skip, form_data.limit)
chats = Chats.get_chat_list_by_chat_ids(
db, chat_ids, form_data.skip, form_data.limit
)
if len(chats) == 0:
Tags.delete_tag_by_tag_name_and_user_id(form_data.name, user.id)
Tags.delete_tag_by_tag_name_and_user_id(db, form_data.name, user.id)
return chats
@@ -229,9 +243,9 @@ async def get_user_chat_list_by_tag_name(
@router.get("/tags/all", response_model=List[TagModel])
async def get_all_tags(user=Depends(get_current_user)):
async def get_all_tags(user=Depends(get_current_user), db=Depends(get_db)):
try:
tags = Tags.get_tags_by_user_id(user.id)
tags = Tags.get_tags_by_user_id(db, user.id)
return tags
except Exception as e:
log.exception(e)
@@ -246,8 +260,8 @@ async def get_all_tags(user=Depends(get_current_user)):
@router.get("/{id}", response_model=Optional[ChatResponse])
async def get_chat_by_id(id: str, user=Depends(get_current_user)):
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
async def get_chat_by_id(id: str, user=Depends(get_current_user), db=Depends(get_db)):
chat = Chats.get_chat_by_id_and_user_id(db, id, user.id)
if chat:
return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
@@ -264,13 +278,13 @@ async def get_chat_by_id(id: str, user=Depends(get_current_user)):
@router.post("/{id}", response_model=Optional[ChatResponse])
async def update_chat_by_id(
id: str, form_data: ChatForm, user=Depends(get_current_user)
id: str, form_data: ChatForm, user=Depends(get_current_user), db=Depends(get_db)
):
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
chat = Chats.get_chat_by_id_and_user_id(db, id, user.id)
if chat:
updated_chat = {**json.loads(chat.chat), **form_data.chat}
chat = Chats.update_chat_by_id(id, updated_chat)
chat = Chats.update_chat_by_id(db, id, updated_chat)
return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
else:
raise HTTPException(
@@ -285,10 +299,12 @@ async def update_chat_by_id(
@router.delete("/{id}", response_model=bool)
async def delete_chat_by_id(request: Request, id: str, user=Depends(get_current_user)):
async def delete_chat_by_id(
request: Request, id: str, user=Depends(get_current_user), db=Depends(get_db)
):
if user.role == "admin":
result = Chats.delete_chat_by_id(id)
result = Chats.delete_chat_by_id(db, id)
return result
else:
if not request.app.state.config.USER_PERMISSIONS["chat"]["deletion"]:
@@ -297,7 +313,7 @@ async def delete_chat_by_id(request: Request, id: str, user=Depends(get_current_
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
)
result = Chats.delete_chat_by_id_and_user_id(id, user.id)
result = Chats.delete_chat_by_id_and_user_id(db, id, user.id)
return result
@@ -307,8 +323,8 @@ async def delete_chat_by_id(request: Request, id: str, user=Depends(get_current_
@router.get("/{id}/clone", response_model=Optional[ChatResponse])
async def clone_chat_by_id(id: str, user=Depends(get_current_user)):
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
async def clone_chat_by_id(id: str, user=Depends(get_current_user), db=Depends(get_db)):
chat = Chats.get_chat_by_id_and_user_id(db, id, user.id)
if chat:
chat_body = json.loads(chat.chat)
@@ -319,7 +335,7 @@ async def clone_chat_by_id(id: str, user=Depends(get_current_user)):
"title": f"Clone of {chat.title}",
}
chat = Chats.insert_new_chat(user.id, ChatForm(**{"chat": updated_chat}))
chat = Chats.insert_new_chat(db, user.id, ChatForm(**{"chat": updated_chat}))
return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
else:
raise HTTPException(
@@ -333,10 +349,12 @@ async def clone_chat_by_id(id: str, user=Depends(get_current_user)):
@router.get("/{id}/archive", response_model=Optional[ChatResponse])
async def archive_chat_by_id(id: str, user=Depends(get_current_user)):
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
async def archive_chat_by_id(
id: str, user=Depends(get_current_user), db=Depends(get_db)
):
chat = Chats.get_chat_by_id_and_user_id(db, id, user.id)
if chat:
chat = Chats.toggle_chat_archive_by_id(id)
chat = Chats.toggle_chat_archive_by_id(db, id)
return ChatResponse(**{**chat.model_dump(), "chat": json.loads(chat.chat)})
else:
raise HTTPException(
@@ -350,16 +368,16 @@ async def archive_chat_by_id(id: str, user=Depends(get_current_user)):
@router.post("/{id}/share", response_model=Optional[ChatResponse])
async def share_chat_by_id(id: str, user=Depends(get_current_user)):
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
async def share_chat_by_id(id: str, user=Depends(get_current_user), db=Depends(get_db)):
chat = Chats.get_chat_by_id_and_user_id(db, id, user.id)
if chat:
if chat.share_id:
shared_chat = Chats.update_shared_chat_by_chat_id(chat.id)
shared_chat = Chats.update_shared_chat_by_chat_id(db, chat.id)
return ChatResponse(
**{**shared_chat.model_dump(), "chat": json.loads(shared_chat.chat)}
)
shared_chat = Chats.insert_shared_chat_by_chat_id(chat.id)
shared_chat = Chats.insert_shared_chat_by_chat_id(db, chat.id)
if not shared_chat:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
@@ -382,14 +400,16 @@ async def share_chat_by_id(id: str, user=Depends(get_current_user)):
@router.delete("/{id}/share", response_model=Optional[bool])
async def delete_shared_chat_by_id(id: str, user=Depends(get_current_user)):
chat = Chats.get_chat_by_id_and_user_id(id, user.id)
async def delete_shared_chat_by_id(
id: str, user=Depends(get_current_user), db=Depends(get_db)
):
chat = Chats.get_chat_by_id_and_user_id(db, id, user.id)
if chat:
if not chat.share_id:
return False
result = Chats.delete_shared_chat_by_chat_id(id)
update_result = Chats.update_chat_share_id_by_id(id, None)
result = Chats.delete_shared_chat_by_chat_id(db, id)
update_result = Chats.update_chat_share_id_by_id(db, id, None)
return result and update_result != None
else:
@@ -405,8 +425,10 @@ async def delete_shared_chat_by_id(id: str, user=Depends(get_current_user)):
@router.get("/{id}/tags", response_model=List[TagModel])
async def get_chat_tags_by_id(id: str, user=Depends(get_current_user)):
tags = Tags.get_tags_by_chat_id_and_user_id(id, user.id)
async def get_chat_tags_by_id(
id: str, user=Depends(get_current_user), db=Depends(get_db)
):
tags = Tags.get_tags_by_chat_id_and_user_id(db, id, user.id)
if tags != None:
return tags
@@ -423,12 +445,15 @@ async def get_chat_tags_by_id(id: str, user=Depends(get_current_user)):
@router.post("/{id}/tags", response_model=Optional[ChatIdTagModel])
async def add_chat_tag_by_id(
id: str, form_data: ChatIdTagForm, user=Depends(get_current_user)
id: str,
form_data: ChatIdTagForm,
user=Depends(get_current_user),
db=Depends(get_db),
):
tags = Tags.get_tags_by_chat_id_and_user_id(id, user.id)
tags = Tags.get_tags_by_chat_id_and_user_id(db, id, user.id)
if form_data.tag_name not in tags:
tag = Tags.add_tag_to_chat(user.id, form_data)
tag = Tags.add_tag_to_chat(db, user.id, form_data)
if tag:
return tag
@@ -450,10 +475,13 @@ async def add_chat_tag_by_id(
@router.delete("/{id}/tags", response_model=Optional[bool])
async def delete_chat_tag_by_id(
id: str, form_data: ChatIdTagForm, user=Depends(get_current_user)
id: str,
form_data: ChatIdTagForm,
user=Depends(get_current_user),
db=Depends(get_db),
):
result = Tags.delete_tag_by_tag_name_and_chat_id_and_user_id(
form_data.tag_name, id, user.id
db, form_data.tag_name, id, user.id
)
if result:
@@ -470,8 +498,10 @@ async def delete_chat_tag_by_id(
@router.delete("/{id}/tags/all", response_model=Optional[bool])
async def delete_all_chat_tags_by_id(id: str, user=Depends(get_current_user)):
result = Tags.delete_tags_by_chat_id_and_user_id(id, user.id)
async def delete_all_chat_tags_by_id(
id: str, user=Depends(get_current_user), db=Depends(get_db)
):
result = Tags.delete_tags_by_chat_id_and_user_id(db, id, user.id)
if result:
return result