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 os
import uuid

from PyQt6.QtCore import pyqtSignal, pyqtSlot
from PyQt6.QtWidgets import QDialog, QWidget

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

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 viewChatHistory(id:str) emitted to indicate, that this chat history
        should be shown in a separate window
    @signal dataChanged(id:str) emitted to indicate a change of the chat history data
    @signal parametersChanged(id:str, title:str, model:str) emitted to indicate a
        change of editable chat parameters
    """

    deleteChatHistory = pyqtSignal(str)
    newChatWithHistory = pyqtSignal(str)
    viewChatHistory = pyqtSignal(str)
    dataChanged = pyqtSignal(str)
    parametersChanged = pyqtSignal(str, str, str)

    def __init__(self, title, model, mainWidget, 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 mainWidget reference to the Ollama main widget
        @type OllamaWidget
        @param jsonStr string containing JSON serialized 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)

        iconSuffix = "-dark" if ericApp().usesDarkPalette() else "-light"

        self.newChatButton.setIcon(EricPixmapCache.getIcon("plus"))
        self.editButton.setIcon(EricPixmapCache.getIcon("editRename"))
        self.viewButton.setIcon(
            EricPixmapCache.getIcon(
                os.path.join("OllamaInterface", "icons", "view{0}".format(iconSuffix))
            )
        )
        self.deleteButton.setIcon(EricPixmapCache.getIcon("trash"))

        self.__ollamaWidget = mainWidget

        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.
        """
        from .OllamaHistoryEditDialog import OllamaHistoryEditDialog

        dlg = OllamaHistoryEditDialog(
            title=self.__title,
            model=self.__model,
            selectableModels=self.__ollamaWidget.getSelectableModels(),
        )
        if dlg.exec() == QDialog.DialogCode.Accepted:
            self.__title, self.__model = dlg.getData()
            self.titleEdit.setText(self.__title)
            self.modelEdit.setText(self.__model)
            self.dataChanged.emit(self.__id)
            self.parametersChanged.emit(self.__id, self.__title, self.__model)

    @pyqtSlot()
    def on_viewButton_clicked(self):
        """
        Private slot to view the current chat history in a separate window..
        """
        self.viewChatHistory.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)