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 QDialog, QDialogButtonBox, 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.__setHeader(title, model)

        self.__messagesLayout = QVBoxLayout()
        if parent is None:
            self.__messagesLayout.setContentsMargins(4, 4, 4, 4)
        else:
            self.__messagesLayout.setContentsMargins(0, 0, 0, 0)
        self.__messagesLayout.setAlignment(Qt.AlignmentFlag.AlignTop)
        self.chatMessagesWidget.setLayout(self.__messagesLayout)

    def __setHeader(self, title, model):
        """
        Private method to set the header label.

        @param title chat title
        @type str
        @param model model of the chat
        @type str
        """
        self.headerLabel.setText(
            self.tr("<b>{0} - {1}</b>", "title, model name").format(title, model)
        )

    def addMessage(self, role, message, scrollToBottom=True):
        """
        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
        @param scrollToBottom flag indicating to scroll the view to the bottom
            (defaults to True)
        @type bool (optional)
        """
        msgWidget = OllamaChatMessageBox(role=role, message=message)
        self.__messagesLayout.addWidget(msgWidget)

        if scrollToBottom:
            QTimer.singleShot(100, 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 history ID of the chat
        @rtype str
        """
        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()

    def clear(self):
        """
        Public method to clear the list of messages.
        """
        while not self.__messagesLayout.isEmpty():
            itm = self.__messagesLayout.takeAt(0)
            itm.widget().deleteLater()

        scrollbar = self.chatMessagesScrollArea.verticalScrollBar()
        scrollbar.setValue(0)

    @pyqtSlot(str, str, str)
    def chatParametersChanged(self, _hid, title, model):
        """
        Public slot to handle a change of the chat parameters.

        @param _hid history ID (unused)
        @type str
        @param title new chat title
        @type str
        @param model new chat model
        @type str
        """
        self.__setHeader(title, model)


class OllamaChatDialog(QDialog):
    """
    Class implementing a dialog 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.__layout = QVBoxLayout(self)
        self.__chatWidget = OllamaChatWidget(
            hid=hid, title=title, model=model, parent=self
        )
        self.__layout.addWidget(self.__chatWidget)
        self.__buttonBox = QDialogButtonBox(
            QDialogButtonBox.StandardButton.Close, Qt.Orientation.Horizontal, self
        )
        self.__buttonBox.rejected.connect(self.reject)
        self.__layout.addWidget(self.__buttonBox)

        self.setModal(False)
        self.setSizeGripEnabled(True)
        self.resize(600, 750)

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

        @return DESCRIPTION
        @rtype TYPE
        """
        return self.__chatWidget.getHistoryId()

    def setHistory(self, messages):
        """
        Public method to add a list of messages to the view.

        @param messages list of chat messages
        @type list[dict[str, str]]
        """
        self.__chatWidget.clear()

        for message in messages:
            self.__chatWidget.addMessage(
                role=message["role"], message=message["content"], scrollToBottom=False
            )