OllamaInterface/OllamaChatMessageBox.py

# -*- coding: utf-8 -*-

# Copyright (c) 2024 Detlev Offenbach <detlev@die-offenbachs.de>
#

"""
Module implementing a message box widget showing the role and content of a message.
"""

import os

from PyQt6.QtCore import QSize, Qt
from PyQt6.QtWidgets import QHBoxLayout, QLabel, QWidget

from eric7.EricGui import EricPixmapCache
from eric7.EricWidgets.EricApplication import ericApp

from .AutoResizeTextBrowser import AutoResizeTextBrowser


class OllamaChatMessageBox(QWidget):
    """
    Class implementing a message box widget showing the role and content of a message.
    """

    def __init__(self, role, message, parent=None):
        """
        Constructor

        @param role role of the message sender (one of 'user' or 'assistant')
        @type str
        @param message message text
        @type str
        @param parent reference to the parent widget (defaults to None)
        @type QWidget (optional)
        """
        super().__init__(parent)

        self.__roleLabel = QLabel(self)
        self.__roleLabel.setFixedSize(22, 22)
        pixmapName = "{0}-{1}".format(
            "user" if role == "user" else "ollama22",
            "dark" if ericApp().usesDarkPalette() else "light",
        )
        self.__roleLabel.setPixmap(
            EricPixmapCache.getPixmap(
                os.path.join("OllamaInterface", "icons", pixmapName),
                QSize(22, 22),
            )
        )
        self.__roleLabel.setAlignment(
            Qt.AlignmentFlag.AlignHCenter | Qt.AlignmentFlag.AlignTop
        )

        self.__messageBrowser = AutoResizeTextBrowser(self)

        self.__layout = QHBoxLayout(self)
        self.__layout.setAlignment(Qt.AlignmentFlag.AlignTop)
        self.__layout.setContentsMargins(0, 0, 0, 0)
        self.__layout.addWidget(self.__roleLabel, Qt.AlignmentFlag.AlignTop)
        self.__layout.addWidget(self.__messageBrowser)
        self.setLayout(self.__layout)

        self.__message = ""
        self.appendMessage(message)

    def appendMessage(self, msg):
        """
        Public method to append the given message text to the current content.

        @param msg message to be appended
        @type str
        """
        if msg:
            self.__message += msg
            self.__messageBrowser.setMarkdown(self.__message)

    def getMessage(self):
        """
        Public method to get the message content.

        @return message content
        @rtype str
        """
        return self.__message