OllamaInterface/OllamaHistoryWidget.py

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

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

"""
Module implementing a widget showing a chat title and store a chat contents.
"""

import json
import uuid

from PyQt6.QtCore import pyqtSignal, pyqtSlot
from PyQt6.QtWidgets import QInputDialog, QLineEdit, QWidget

from eric7.EricGui import EricPixmapCache

from .Ui_OllamaHistoryWidget import Ui_OllamaHistoryWidget


class OllamaHistoryWidget(QWidget, Ui_OllamaHistoryWidget):
    """
    Class implementing a widget showing a chat title and store a chat contents.

    @signal deleteChatHistory(id:str) emitted to indicate, that this chat history
        should be deleted
    @signal newChatWithHistory(id:str) emitted to indicate, that a new chat using
        the saved history should be started
    @signal dataChanged(id:str) emitted to indicate a change of the chat history data
    """

    deleteChatHistory = pyqtSignal(str)
    newChatWithHistory = pyqtSignal(str)
    dataChanged = pyqtSignal(str)

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

        @param title title of the ollama chat
        @type str
        @param model name of the model used for the chat
        @type str
        @param jsonStr string containing JSON serialize chat history data
            (defaults to None)
        @type str (optional)
        @param parent reference to the parent widget (defaults to None)
        @type QWidget (optional)
        """
        super().__init__(parent)
        self.setupUi(self)

        self.newChatButton.setIcon(EricPixmapCache.getIcon("plus"))
        self.editButton.setIcon(EricPixmapCache.getIcon("editRename"))
        self.deleteButton.setIcon(EricPixmapCache.getIcon("trash"))

        if jsonStr is None:
            self.__title = title
            self.__model = model

            self.__id = str(uuid.uuid4())
            self.__messages = []
        else:
            self.loadFromJson(jsonStr)

        self.titleEdit.setText(self.__title)
        self.modelEdit.setText(self.__model)

    def getTitle(self):
        """
        Public method to get the chat title.

        @return chat title
        @rtype str
        """
        return self.__title

    def getModel(self):
        """
        Public method to get the model used by the chat.

        @return model name
        @rtype str
        """
        return self.__model

    def getId(self):
        """
        Public method to get the chat history ID.

        @return ID of the history entry
        @rtype str
        """
        return self.__id

    def getMessages(self):
        """
        Public method to get the list of messages.

        @return list of stored messages
        @rtype list[dict[str, str]]
        """
        return self.__messages

    @pyqtSlot()
    def on_deleteButton_clicked(self):
        """
        Private slot to delet this chat history entry..
        """
        self.deleteChatHistory.emit(self.__id)

    @pyqtSlot()
    def on_newChatButton_clicked(self):
        """
        Private slot to start a new chat using the saved chat history.
        """
        self.newChatWithHistory.emit(self.__id)

    @pyqtSlot()
    def on_editButton_clicked(self):
        """
        Private slot to edit the chat title.
        """
        title, ok = QInputDialog.getText(
            self,
            self.tr("Edit Chat Title"),
            self.tr("Enter the new title:"),
            QLineEdit.EchoMode.Normal,
            self.__title,
        )
        if ok and bool(title):
            self.__title = title
            self.titleEdit.setText(title)
            self.dataChanged.emit(self.__id)

    def loadFromJson(self, jsonStr):
        """
        Public method to load the chat history data from a JSON string.

        @param jsonStr JSON serialized chat data
        @type str
        """
        data = json.loads(jsonStr)
        self.__id = data["id"]
        self.__title = data["title"]
        self.__model = data["model"]
        self.__messages = data["messages"]

    def saveToJson(self):
        """
        Public method to serialize the chat history data to a JSON string.

        @return JSON serialized chat data
        @rtype str
        """
        return json.dumps(
            {
                "id": self.__id,
                "title": self.__title,
                "model": self.__model,
                "messages": self.__messages,
            }
        )

    def addToMessages(self, role, content):
        """
        Public method to add a chat message to the chat history.

        @param role chat role (one of 'system', 'user', 'assistant' or 'tool')
        @type str
        @param content content of the chat message
        @type str
        """
        self.__messages.append(
            {
                "role": role,
                "content": content,
            }
        )
        self.dataChanged.emit(self.__id)

    def checkModelAvailable(self, modelNames):
        """
        Public method to set the enabled state depending on available models.

        @param modelNames names of available models
        @type str
        """
        available = self.__model in modelNames

        self.newChatButton.setEnabled(available)
        self.editButton.setEnabled(available)