ocr/api/message/db_requests.py

import asyncio

from fastapi import HTTPException

from ocr.api.message.dto import Author
from ocr.api.message.models import MessageModel
from ocr.api.message.schemas import CreateMessageRequest
from ocr.core.config import settings


async def create_message_obj(
        report_id: str, message_data: CreateMessageRequest
) -> MessageModel:
    report = await settings.DB_CLIENT.reports.find_one({"id": report_id})
    if not report:
        raise HTTPException(status_code=404, detail="Report not found.")

    message = MessageModel(**message_data.model_dump(), reportId=report_id, author=Author.User)
    await settings.DB_CLIENT.messages.insert_one(message.to_mongo())
    return message


async def get_all_chat_messages_obj(report_id: str) -> list[MessageModel]:
    messages, report = await asyncio.gather(
        settings.DB_CLIENT.messages.find({"reportId": report_id}).to_list(length=None),
        settings.DB_CLIENT.reports.find_one({"id": report_id})
    )
    messages = [MessageModel.from_mongo(message) for message in messages]
    if not report:
        raise HTTPException(status_code=404, detail="Report not found")
    return messages


async def save_assistant_user_message(user_message: str, assistant_message: str, report_id: str) -> MessageModel:
    user_message = MessageModel(reportId=report_id, author=Author.User, text=user_message)
    assistant_message = MessageModel(reportId=report_id, author=Author.Assistant, text=assistant_message)

    await settings.DB_CLIENT.messages.insert_one(user_message.to_mongo())
    await settings.DB_CLIENT.messages.insert_one(assistant_message.to_mongo())

    return assistant_message