OllamaInterface/OllamaPullProgressDialog.py

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

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

"""
Module implementing a dialog showing the progress of a model pull action..
"""

from PyQt6.QtCore import Qt, pyqtSignal, pyqtSlot
from PyQt6.QtWidgets import (
    QAbstractButton,
    QDialog,
    QDialogButtonBox,
    QProgressBar,
    QTreeWidgetItem,
)

from eric7 import Globals

from .Ui_OllamaPullProgressDialog import Ui_OllamaPullProgressDialog


class OllamaPullProgressBar(QProgressBar):
    """
    Class implementing a progress bar allowing values outside the standard range.
    """

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

        @param parent reference to the parent widget (defaults to None)
        @type QWidget (optional)
        """
        super().__init__(parent)

        self.__maximum = 100
        self.__minimum = 0
        self.__value = 0

    def maximum(self):
        """
        Public method to get the maximum value.

        @return maximum value
        @rtype int
        """
        return self.__maximum

    def setMaximum(self, value):
        """
        Public method to set the maximum value.

        @param value new maximum value
        @type int
        """
        if value != self.__maximum:
            self.__maximum = value
            self.setValue(self.__value)

    def value(self):
        """
        Public method to get the current value.

        @return current value
        @rtype int
        """
        return self.__value

    def setValue(self, value):
        """
        Public method to set the current value.

        @param value new value
        @type int
        """
        if value != self.__value:
            self.__value = value
            super().setValue(self.__value * 100 // self.__maximum)


class OllamaPullProgressDialog(QDialog, Ui_OllamaPullProgressDialog):
    """
    Class implementing a dialog showing the progress of a model pull action.

    @signal abortPull() emitted to abort the current model pull operation
    """

    abortPull = pyqtSignal()

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

        @param parent reference to the parent widget (defaults to None)
        @type QWidget (optional)
        """
        super().__init__(parent)
        self.setupUi(self)

        self.header.clear()
        self.__progressBarItems = {}

        self.setFinished(True)

    @pyqtSlot(QAbstractButton)
    def on_buttonBox_clicked(self, button):
        """
        Private slot handling a button of the button box being clicked.

        @param button reference to the clicked button
        @type QAbstractButton
        """
        if button == self.buttonBox.button(QDialogButtonBox.StandardButton.Cancel):
            self.abortPull.emit()
        elif button == self.buttonBox.button(QDialogButtonBox.StandardButton.Close):
            self.close()

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

        @param evt reference to the close event object
        @type QCloseEvent
        """
        if not self.__finished:
            evt.ignore()

    def setModel(self, model):
        """
        Public method to show the model name in the header.

        @param model model name
        @type str
        """
        self.header.setText(
            self.tr("<p>Installing model <b>{0}</b>.</p>").format(model)
        )

    def clear(self):
        """
        Public method to clear the progress information.
        """
        self.__progressBarItems.clear()
        self.progressList.clear()

    def setStatus(self, status, idStr, total, completed):
        """
        Public method to show the status update.

        @param status status message reported by the 'ollama' server
        @type str
        @param idStr ID of the file being pulled or empty
        @type str
        @param total size of the file being pulled or 0 in case of an empty ID
        @type int
        @param completed downloaded bytes or 0 in case of an empty ID
        @type int
        """
        self.setFinished(False)

        if idStr:
            try:
                itm = self.__progressBarItems[idStr]
            except KeyError:
                itm = QTreeWidgetItem(self.progressList, [status, "", ""])
                itm.setTextAlignment(2, Qt.AlignmentFlag.AlignCenter)
                self.progressList.resizeColumnToContents(0)

                bar = OllamaPullProgressBar()
                bar.setMaximum(total)
                self.progressList.setItemWidget(itm, 1, bar)

                self.__progressBarItems[idStr] = itm

            if completed == total:
                itm.setText(2, Globals.dataString(total))
            else:
                itm.setText(
                    2,
                    self.tr("{0} / {1}", "completed / total").format(
                        Globals.dataString(completed), Globals.dataString(total)
                    ),
                )
            self.progressList.itemWidget(itm, 1).setValue(completed)
        else:
            itm = QTreeWidgetItem(self.progressList, [status])
            itm.setFirstColumnSpanned(True)

            if status == "success":
                self.setFinished(True)

    def showError(self, errMsg):
        """
        Public method to show an error message reported by the server.

        @param errMsg error message
        @type str
        """
        itm = QTreeWidgetItem(self.progressList, [self.tr("Error: {0}").format(errMsg)])
        itm.setFirstColumnSpanned(True)

    def setFinished(self, finished):
        """
        Public method to set the finished state.

        @param finished flag indicating the finished state
        @type bool
        """
        self.__finished = finished

        self.buttonBox.button(QDialogButtonBox.StandardButton.Close).setEnabled(
            finished
        )
        self.buttonBox.button(QDialogButtonBox.StandardButton.Cancel).setEnabled(
            not finished
        )