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