trauma/api/message/views.py

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
from trauma.api.message.schemas import (AllMessageWrapper,
                                        AllMessageResponse,
                                        CreateMessageRequest,
                                        CreateMessageResponse)
from trauma.api.message.utils import transform_messages_to_openai
from trauma.core.wrappers import TraumaResponseWrapper


@message_router.get('/{chatId}/all')
async def get_all_chat_messages(
        chatId: str
) -> AllMessageWrapper:
    messages, _ = await get_all_chat_messages_obj(chatId)
    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,
) -> TraumaResponseWrapper[CreateMessageResponse]:
    messages, chat = await get_all_chat_messages_obj(chatId)
    message_history = transform_messages_to_openai(messages)
    response = await search_entities(message_data.text, message_history, chat)
    return TraumaResponseWrapper(data=response)