OllamaInterface/OllamaChatWidget.py

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

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

"""
Module implementing a widget showing the chat with the 'ollama' server.
"""

from PyQt6.QtCore import Qt, QTimer, pyqtSlot
from PyQt6.QtWidgets import QVBoxLayout, QWidget

from .OllamaChatMessageBox import OllamaChatMessageBox
from .Ui_OllamaChatWidget import Ui_OllamaChatWidget


class OllamaChatWidget(QWidget, Ui_OllamaChatWidget):
    """
    Class implementing a widget showing the chat with the 'ollama' server.
    """

    def __init__(self, hid, title, model, parent=None):
        """
        Constructor

        @param hid ID of the chat history
        @type str
        @param title title of the chat
        @type str
        @param model model name used for the chat
        @type str
        @param parent reference to the parent widget (defaults to None)
        @type QWidget (optional)
        """
        super().__init__(parent)
        self.setupUi(self)

        self.__hid = hid

        self.headerLabel.setText(
            self.tr("<b>{0} - {1}</b>", "title, model name").format(title, model)
        )

        self.__messagesLayout = QVBoxLayout()
        self.__messagesLayout.setContentsMargins(4, 4, 4, 4)
        self.__messagesLayout.setAlignment(Qt.AlignmentFlag.AlignTop)
        self.chatMessagesWidget.setLayout(self.__messagesLayout)

    def addMessage(self, role, message):
        """
        Public method to add a new message.

        @param role role of the message sender (one of 'user' or 'assistant')
        @type str
        @param message message text
        @type str
        """
        msgWidget = OllamaChatMessageBox(role=role, message=message)
        self.__messagesLayout.addWidget(msgWidget)

        QTimer.singleShot(0, self.__scrollChatToBottom)

    def appendMessage(self, message):
        """
        Public method to append a given message to the bottom most message box.

        @param message message text to be appended
        @type str
        """
        msgBox = self.__messagesLayout.itemAt(
            self.__messagesLayout.count() - 1
        ).widget()
        msgBox.appendMessage(message)

        QTimer.singleShot(0, self.__scrollChatToBottom)

    @pyqtSlot()
    def __scrollChatToBottom(self):
        """
        Private slot to scroll the chat scroll area to the bottom.
        """
        scrollbar = self.chatMessagesScrollArea.verticalScrollBar()
        scrollbar.setValue(scrollbar.maximum())

    def getHistoryId(self):
        """
        Public method to get the history ID of this chat.

        @return DESCRIPTION
        @rtype TYPE
        """
        return self.__hid

    def getRecentMessage(self):
        """
        Public method to get the message of the last message box.

        @return message content
        @rtype str
        """
        msgBox = self.__messagesLayout.itemAt(
            self.__messagesLayout.count() - 1
        ).widget()
        return msgBox.getMessage()