Repository: LAION-AI/Open-Assistant
File: inference/server/oasst_inference_server/user_chat_repository.py (Python)

Code before:
        )
        self.session.add(message)

        self.maybe_commit()
        logger.debug(f"Added prompter message {len(content)=} to chat {chat_id}")
        self.session.refresh(message)
        return message

    def initiate_assistant_message(self, parent_id: str, work_parameters: inference.WorkParameters) -> models.DbMessage:
        logger.info(f"Adding stub assistant message to {parent_id=}")
        parent = self.get_prompter_message_by_id(parent_id)
        message = models.DbMessage(
            role="assistant",
            chat_id=parent.chat_id,
            chat=parent.chat,
            parent_id=parent_id,
            state=inference.MessageState.pending,
            work_parameters=work_parameters,
        )
        self.session.add(message)
        self.maybe_commit()
        logger.debug(f"Initiated assistant message of {parent_id=}")
        self.session.refresh(message)
        return message

    def update_score(self, message_id: str, score: int):
        logger.info(f"Updating message score to {message_id=}: {score=}")
        message = self.get_assistant_message_by_id(message_id)
        message.score = score
        self.session.commit()
        return message

    def add_report(self, message_id: str, reason: str, report_type: inference.ReportType) -> models.DbReport:
        logger.info(f"Adding report to {message_id=}: {reason=}")
        message = self.get_assistant_message_by_id(message_id)
        report = models.DbReport(message_id=message.id, reason=reason, report_type=report_type)
        self.session.add(report)
        self.maybe_commit()
        self.session.refresh(report)
        return report


Review comment (suggestion):
If we are returning message can we add the type hint here as `models.DbMessage`

Code after:
        )
        self.session.add(message)

        self.maybe_commit()
        logger.debug(f"Added prompter message {len(content)=} to chat {chat_id}")
        self.session.refresh(message)
        return message

    def initiate_assistant_message(self, parent_id: str, work_parameters: inference.WorkParameters) -> models.DbMessage:
        logger.info(f"Adding stub assistant message to {parent_id=}")
        parent = self.get_prompter_message_by_id(parent_id)
        message = models.DbMessage(
            role="assistant",
            chat_id=parent.chat_id,
            chat=parent.chat,
            parent_id=parent_id,
            state=inference.MessageState.pending,
            work_parameters=work_parameters,
        )
        self.session.add(message)
        self.maybe_commit()
        logger.debug(f"Initiated assistant message of {parent_id=}")
        self.session.refresh(message)
        return message

    def update_score(self, message_id: str, score: int) -> models.DbMessage:
        if score < -1 or score > 1:
            raise fastapi.HTTPException(status_code=400, detail="Invalid score")

        logger.info(f"Updating message score to {message_id=}: {score=}")
        message = self.get_assistant_message_by_id(message_id)
        message.score = score
        self.session.commit()
        return message

    def add_report(self, message_id: str, reason: str, report_type: inference.ReportType) -> models.DbReport:
        logger.info(f"Adding report to {message_id=}: {reason=}")
        message = self.get_assistant_message_by_id(message_id)
        report = models.DbReport(message_id=message.id, reason=reason, report_type=report_type)
        self.session.add(report)
        self.maybe_commit()
        self.session.refresh(report)
        return report