Spaces:
Running
Running
File size: 2,835 Bytes
50553ea c6cc0f2 50553ea 3ec35ef 50553ea c6cc0f2 50553ea c6cc0f2 50553ea c6cc0f2 50553ea c6cc0f2 50553ea c6cc0f2 50553ea c6cc0f2 50553ea c6cc0f2 50553ea c6cc0f2 50553ea 9150f8e 50553ea 9150f8e 50553ea 9150f8e 50553ea 9150f8e 3ec35ef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import asyncio
from fastapi import HTTPException
from trauma.api.account.dto import AccountType
from trauma.api.account.model import AccountModel
from trauma.api.chat.model import ChatModel
from trauma.api.chat.schemas import CreateChatRequest, ChatTitleRequest
from trauma.api.message.dto import Author
from trauma.api.message.model import MessageModel
from trauma.core.config import settings
async def get_chat_obj(chat_id: str, account: AccountModel) -> ChatModel:
chat = await settings.DB_CLIENT.chats.find_one({"id": chat_id})
if not chat:
raise HTTPException(status_code=404, detail="Chat not found")
chat = ChatModel.from_mongo(chat)
if chat.account.id != account.id and account.accountType != AccountType.Admin:
raise HTTPException(status_code=403, detail="Not authorized")
return chat
async def create_chat_obj(chat_request: CreateChatRequest, account: AccountModel) -> ChatModel:
chat = ChatModel(model=chat_request.model, account=account)
await settings.DB_CLIENT.chats.insert_one(chat.to_mongo())
return chat
async def delete_chat_obj(chat_id: str, account: AccountModel) -> None:
chat = await settings.DB_CLIENT.chats.find_one({"id": chat_id})
if not chat:
raise HTTPException(status_code=404, detail="Chat not found")
chat = ChatModel.from_mongo(chat)
if chat.account.id != account.id and account.accountType != AccountType.Admin:
raise HTTPException(status_code=403, detail="Not authorized")
await settings.DB_CLIENT.chats.delete_one({"id": chat_id})
async def update_chat_obj_title(chat_id: str, chat_request: ChatTitleRequest, account: AccountModel) -> ChatModel:
chat = await settings.DB_CLIENT.chats.find_one({"id": chat_id})
if not chat:
raise HTTPException(status_code=404, detail="Chat not found")
chat = ChatModel.from_mongo(chat)
if chat.account.id != account.id and account.accountType != AccountType.Admin:
raise HTTPException(status_code=403, detail="Not authorized")
chat.title = chat_request.title
await settings.DB_CLIENT.chats.update_one({"id": chat_id}, {"$set": chat.to_mongo()})
return chat
async def get_all_chats_obj(page_size: int, page_index: int) -> tuple[list[ChatModel], int]:
skip = page_size * page_index
objects, total_count = await asyncio.gather(
settings.DB_CLIENT.chats
.find({})
.sort("_id", -1)
.skip(skip)
.limit(page_size)
.to_list(length=page_size),
settings.DB_CLIENT.chats.count_documents({}),
)
return objects, total_count
async def save_intro_message(chat_id: str) -> None:
message = settings.INTRO_MESSAGE
message = MessageModel(author=Author.Assistant, text=message, chatId=chat_id)
await settings.DB_CLIENT.messages.insert_one(message.to_mongo())
|