OllamaInterface/OllamaWidget.py

# -*- coding: utf-8 -*-
# Copyright (c) 2024 Detlev Offenbach <detlev@die-offenbachs.de>
#

"""
Module implementing the main ollama interface widget.
"""

import json
import os

from PyQt6.QtCore import Qt, pyqtSlot
from PyQt6.QtWidgets import QInputDialog, QLineEdit, QVBoxLayout, QWidget

from eric7 import Globals
from eric7.EricGui import EricPixmapCache
from eric7.EricWidgets import EricMessageBox

from .OllamaClient import OllamaClient
from .OllamaHistoryWidget import OllamaHistoryWidget
from .Ui_OllamaWidget import Ui_OllamaWidget


class OllamaWidget(QWidget, Ui_OllamaWidget):
    """
    Class implementing the main ollama interface widget.
    """

    OllamaHistoryFile = "ollama_history.json"

    def __init__(self, plugin, fromEric=True, parent=None):
        """
        Constructor

        @param plugin reference to the plug-in object
        @type PluginOllamaInterface
        @param fromEric flag indicating the eric-ide mode (defaults to True)
            (True = eric-ide mode, False = application mode)
        @type bool (optional)
        @param parent reference to the parent widget (defaults to None)
        @type QWidget (optional)
        """
        super().__init__(parent)
        self.setupUi(self)

        self.__plugin = plugin
        self.__client = OllamaClient(plugin, self)

        if fromEric:
            self.layout().setContentsMargins(0, 3, 0, 0)
        else:
            self.layout().setContentsMargins(0, 0, 0, 0)

        self.ollamaMenuButton.setIcon(EricPixmapCache.getIcon("superMenu"))
        self.reloadModelsButton.setIcon(EricPixmapCache.getIcon("reload"))
        self.newChatButton.setIcon(EricPixmapCache.getIcon("plus"))

        self.__chatHistoryLayout = QVBoxLayout()
        self.historyScrollWidget.setLayout(self.__chatHistoryLayout)
        self.__chatHistoryLayout.addStretch(1)

        self.mainSplitter.setSizes([200, 2000])

        self.newChatButton.setEnabled(False)
        self.__handleServerStateChanged(False)

        self.__connectClient()

        self.__loadHistory()

    def __connectClient(self):
        """
        Private method to connect the client signals.
        """
        self.__client.serverStateChanged.connect(self.__handleServerStateChanged)
        self.__client.serverVersion.connect(self.__setHeaderLabel)
        self.__client.modelsList.connect(self.__populateModelSelector)

    @pyqtSlot(bool)
    def __handleServerStateChanged(self, ok):
        """
        Private slot handling a change in the 'ollama' server responsiveness.

        @param ok flag indicating a responsive 'ollama' server
        @type bool
        """
        if ok:
            self.__finishSetup()
        else:
            self.ollamaVersionLabel.setText(
                self.tr("<b>Error: The configured server is not responding.</b>")
            )
        self.setEnabled(ok)

    @pyqtSlot()
    def __finishSetup(self):
        """
        Private slot to finish the UI setup.
        """
        self.__client.version()
        self.__client.list()

    @pyqtSlot()
    def on_reloadModelsButton_clicked(self):
        """
        Private slot to reload the list of available models.
        """
        self.__client.list()

    @pyqtSlot(str)
    def on_modelComboBox_currentTextChanged(self, model):
        """
        Private slot handling the selection of a model.

        @param model name of the selected model
        @type str
        """
        self.newChatButton.setEnabled(bool(model))

    @pyqtSlot()
    def on_newChatButton_clicked(self):
        """
        Private slot to start a new chat with the 'ollama' server.
        """
        model = self.modelComboBox.currentText()
        if not model:
            EricMessageBox.critical(
                self,
                self.tr("New Chat"),
                self.tr("""A model has to be selected first. Aborting..."""),
            )
            return

        title, ok = QInputDialog.getText(
            self,
            self.tr("New Chat"),
            self.tr("Enter a title for the new chat:"),
            QLineEdit.EchoMode.Normal,
        )
        if ok and title:
            self.__createHistoryWidget(title, model)
            # TODO: create an empty chat widget for new chat

    ############################################################################
    ## Methods handling signals from the 'ollama' client.
    ############################################################################

    @pyqtSlot(str)
    def __setHeaderLabel(self, version):
        """
        Private slot to receive the 'ollama' server version and set the header.

        @param version 'ollama' server version'
        @type str
        """
        self.ollamaVersionLabel.setText(
            self.tr("<b>ollama Server Version {0}</b>").format(version)
        )

    @pyqtSlot(list)
    def __populateModelSelector(self, modelNames):
        """
        Private slot to receive the list of available model names and populate
        the model selector with them.

        @param modelNames list of model names
        @type list[str]
        """
        self.modelComboBox.clear()

        self.modelComboBox.addItem("")
        self.modelComboBox.addItems(sorted(modelNames))

    ############################################################################
    ## Methods handling signals from the chat history widgets.
    ############################################################################

    def __createHistoryWidget(self, title, model, jsonStr=None):
        """
        Private method to create a chat history widget and insert it into the
        respective layout.

        @param title title of the chat
        @type str
        @param model name of the model
        @type str
        @param jsonStr string containing JSON serialize chat history data (defaults
            to None)
        @type str (optional)
        """
        history = OllamaHistoryWidget(title=title, model=model, jsonStr=jsonStr)
        self.__chatHistoryLayout.insertWidget(
            self.__chatHistoryLayout.count() - 1, history
        )

        scrollbar = self.historyScrollArea.verticalScrollBar()
        scrollbar.setMaximum(self.historyScrollWidget.height())
        scrollbar.setValue(scrollbar.maximum())

        history.deleteChatHistory.connect(self.__deleteHistory)
        history.dataChanged.connect(self.__saveHistory)
        history.newChatWithHistory.connect(self.__newChatWithHistory)

        self.__saveHistory()

    def __findHistoryWidgetIndex(self, uid):
        """
        Private method to find the index of the reference history widget.

        @param uid ID of the history widget
        @type str
        @return index of the history widget
        @rtype int
        """
        for index in range(self.__chatHistoryLayout.count() - 1):
            widget = self.__chatHistoryLayout.itemAt(index).widget()
            if widget.getId() == uid:
                return index

        return None

    def __historyFilePath(self):
        """
        Private method to get the path name of the chat history file.

        @return file path of the chat history file
        @rtype str
        """
        return os.path.join(Globals.getConfigDir(), OllamaWidget.OllamaHistoryFile)

    @pyqtSlot()
    def __saveHistory(self):
        """
        Private method to save the current chat history to the history file.
        """
        # step 1: collect all history entries
        entries = {}
        for index in range(self.__chatHistoryLayout.count() - 1):
            widget = self.__chatHistoryLayout.itemAt(index).widget()
            uid = widget.getId()
            entries[uid] = widget.saveToJson()

        # step 2: save the collected chat histories
        filePath = self.__historyFilePath()
        try:
            with open(filePath, "w") as f:
                json.dump(entries, f)
        except OSError as err:
            EricMessageBox.critical(
                self,
                self.tr("Save Chat History"),
                self.tr(
                    "<p>The chat history could not be saved to <b>{0}</b>.</p>"
                    "<p>Reason: {1}</p>"
                ).format(filePath, str(err)),
            )

    def __loadHistory(self):
        """
        Private method to load a previously saved history file.
        """
        # step 1: load the history file, if it exists
        filePath = self.__historyFilePath()
        if not os.path.exists(filePath):
            return

        try:
            with open(filePath, "r") as f:
                entries = json.load(f)
        except OSError as err:
            EricMessageBox.critical(
                self,
                self.tr("Load Chat History"),
                self.tr(
                    "<p>The chat history could not be loaded from <b>{0}</b>.</p>"
                    "<p>Reason: {1}</p>"
                ).format(filePath, str(err)),
            )
            return

        # step 2: create history widgets
        for uid in entries:
            self.__createHistoryWidget("", "", jsonStr=entries[uid])

    def clearHistory(self):
        """
        Public method to clear the history entries and close all chats.
        """
        while self.__chatHistoryLayout.count() > 1:
            # do not delete the spacer at the end of the list
            item = self.__chatHistoryLayout.takeAt(0)
            if item is not None:
                item.widget().deleteLater()

        self.__saveHistory()

    @pyqtSlot(str)
    def __deleteHistory(self, uid):
        """
        Private slot to delete the history with the given ID.

        @param uid ID of the history to be deleted
        @type str
        """
        widgetIndex = self.__findHistoryWidgetIndex(uid)
        if widgetIndex is not None:
            item = self.__chatHistoryLayout.takeAt(widgetIndex)
            if item is not None:
                item.widget().deleteLater()

        self.__saveHistory()

    @pyqtSlot(str)
    def __newChatWithHistory(self, uid):
        """
        Private slot to start a new chat using a previously saved history.

        @param uid ID of the history to be used
        @type str
        """
        # TODO: not implemented yet
        pass

    #######################################################################
    ## Menu related methods below
    #######################################################################

    def __initOllamaMenu(self):
        """
        Private method to create the super menu and attach it to the super
        menu button.
        """
        # TODO: implement the menu and menu methods
        #       * Clear Chat History
        #       * Show Model Details
        #       * Show Model Processes
        #       * Pull Model
        #       * Show Model Shop (via a web browser)
        #       * Remove Model