OllamaInterface/RunOllamaServerDialog.py

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

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

"""
Module implementing a dialog to run the ollama server locally.
"""

from PyQt6.QtCore import QProcess, Qt, QTimer, pyqtSignal, pyqtSlot
from PyQt6.QtWidgets import QDialog, QDialogButtonBox

from eric7.EricWidgets import EricMessageBox

from .Ui_RunOllamaServerDialog import Ui_RunOllamaServerDialog


class RunOllamaServerDialog(QDialog, Ui_RunOllamaServerDialog):
    """
    Class implementing a dialog to run the ollama server locally.

    @signal serverStarted() emitted after the start of the 'ollama' server
    @signal serverStopped() emitted after the 'ollama' server was stopped
    """

    serverStarted = pyqtSignal()
    serverStopped = pyqtSignal()

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

        @param ollamaClient reference to the 'ollama' client object
        @type OllamaClient
        @param plugin reference to the plug-in object
        @type PluginOllamaInterface
        @param parent reference to the parent widget
        @type QWidget
        """
        super().__init__(parent)
        self.setupUi(self)

        self.__plugin = plugin
        self.__ui = parent
        self.__client = ollamaClient

        self.__process = None

        self.buttonBox.button(QDialogButtonBox.StandardButton.Close).setEnabled(True)
        self.buttonBox.button(QDialogButtonBox.StandardButton.Close).setDefault(True)

        self.__defaultTextFormat = self.outputEdit.currentCharFormat()

    def startServer(self):
        """
        Public method to start the ollama server process.

        @return flag indicating success
        @rtype bool
        """
        env = self.__ui.prepareServerRuntimeEnvironment()
        self.__process = QProcess()
        self.__process.setProcessEnvironment(env)
        self.__process.setProcessChannelMode(QProcess.ProcessChannelMode.MergedChannels)

        self.__process.readyReadStandardOutput.connect(self.__readStdOut)
        self.__process.finished.connect(self.__processFinished)

        self.outputEdit.clear()

        command = "ollama"
        args = ["serve"]

        self.__process.start(command, args)
        ok = self.__process.waitForStarted(10000)
        if not ok:
            EricMessageBox.critical(
                None,
                self.tr("Run Local ollama Server"),
                self.tr("""The loacl ollama server process could not be started."""),
            )
        else:
            self.buttonBox.button(QDialogButtonBox.StandardButton.Close).setEnabled(
                False
            )
            self.stopServerButton.setEnabled(True)
            self.stopServerButton.setDefault(True)
            self.restartServerButton.setEnabled(True)

            self.serverStarted.emit()

        return ok

    def closeEvent(self, evt):
        """
        Protected method handling a close event.

        @param evt reference to the close event
        @type QCloseEvent
        """
        self.on_stopServerButton_clicked()
        evt.accept()

    @pyqtSlot()
    def __readStdOut(self):
        """
        Private slot to add the server process output to the output pane.
        """
        if self.__process is not None:
            out = str(self.__process.readAllStandardOutput(), "utf-8")
            self.outputEdit.insertPlainText(out)

    @pyqtSlot()
    def __processFinished(self):
        """
        Private slot handling the finishing of the server process.
        """
        if (
            self.__process is not None
            and self.__process.state() != QProcess.ProcessState.NotRunning
        ):
            self.__process.terminate()
            QTimer.singleShot(2000, self.__process.kill)
            self.__process.waitForFinished(3000)

        self.__process = None

        self.restartServerButton.setEnabled(True)
        self.stopServerButton.setEnabled(False)
        self.buttonBox.button(QDialogButtonBox.StandardButton.Close).setEnabled(True)
        self.buttonBox.button(QDialogButtonBox.StandardButton.Close).setDefault(True)
        self.buttonBox.button(QDialogButtonBox.StandardButton.Close).setFocus(
            Qt.FocusReason.OtherFocusReason
        )

        self.serverStopped.emit()

    @pyqtSlot()
    def on_stopServerButton_clicked(self):
        """
        Private slot to stop the running server.
        """
        self.__process.terminate()

    @pyqtSlot()
    def on_restartServerButton_clicked(self):
        """
        Private slot to re-start the ollama server.
        """
        # step 1: stop the current server
        if self.__process is not None:
            self.on_stopServerButton_clicked()

        # step 2: start a new server
        self.startServer()