trauma/api/message/views.py

from fastapi import Depends

from trauma.api.account.dto import AccountType
from trauma.api.account.model import AccountModel
from trauma.api.common.dto import Paging
from trauma.api.message import message_router
from trauma.api.message.ai.engine import search_entities
from trauma.api.message.db_requests import get_all_chat_messages_obj, update_message_feedback_obj
from trauma.api.message.dto import Feedback
from trauma.api.message.model import MessageModel
from trauma.api.message.schemas import (AllMessageWrapper,
                                        AllMessageResponse,
                                        CreateMessageRequest)
from trauma.core.security import PermissionDependency
from trauma.core.wrappers import TraumaResponseWrapper
from trauma.api.message.utils import decrypt_messages

@message_router.get('/{chatId}/all')
async def get_all_chat_messages(
        chatId: str,
        account: AccountModel = Depends(PermissionDependency([AccountType.Admin, AccountType.User])),
) -> AllMessageWrapper:
    messages, _ = await get_all_chat_messages_obj(chatId, account)
    response = AllMessageResponse(
        paging=Paging(pageSize=len(messages), pageIndex=0, totalCount=len(messages)),
        data=messages
    )
    return AllMessageWrapper(data=response)


@message_router.post('/{chatId}')
async def create_message(
        chatId: str,
        message_data: CreateMessageRequest,
        account: AccountModel = Depends(PermissionDependency([AccountType.Admin, AccountType.User]))
) -> TraumaResponseWrapper[MessageModel]:
    messages, chat = await get_all_chat_messages_obj(chatId, account)
    messages = decrypt_messages(messages)
    response = await search_entities(message_data.text, messages, chat)
    return TraumaResponseWrapper(data=response)


@message_router.patch('/{messageId}/feedback')
async def create_feedback(
        messageId: str,
        feedback_data: Feedback,
        _: AccountModel = Depends(PermissionDependency([AccountType.Admin, AccountType.User]))
) -> TraumaResponseWrapper[MessageModel]:
    message = await update_message_feedback_obj(messageId, feedback_data)
    return TraumaResponseWrapper(data=message)