trauma/api/message/views.py

from fastapi.params import Depends
from starlette.responses import StreamingResponse

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.openai_request import response_generator
from trauma.api.message.db_requests import get_all_chat_messages_obj, create_message_obj
from trauma.api.message.schemas import (AllMessageWrapper,
                                        AllMessageResponse,
                                        CreateMessageRequest)
from trauma.core.security import PermissionDependency


@message_router.get('/{chatId}/all')
async def get_all_chat_messages(
        chatId: str, account: AccountModel = Depends(PermissionDependency(is_public=True))
) -> 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(is_public=True))
) -> StreamingResponse:
    user_message, chat = await create_message_obj(chatId, message_data, account)
    return StreamingResponse(response_generator(chat, user_message), media_type='text/event-stream')