Mercurial Repositories > ollama Interface / file diff

--- a/OllamaInterface/OllamaWidget.py	Tue Aug 27 09:19:39 2024 +0200
+++ b/OllamaInterface/OllamaWidget.py	Tue Aug 27 14:06:50 2024 +0200
@@ -86,6 +86,9 @@
         self.newChatButton.setEnabled(False)
         self.__handleServerStateChanged(False)
 
+        self.__pullProgressDialog = None
+        self.__pulling = False
+
         self.__localServerDialog = None
         self.__localServerProcess = None
 
@@ -110,6 +113,11 @@
         self.__client.modelsList.connect(self.__populateModelSelector)
         self.__client.modelsList.connect(self.__checkHistoryModels)
         self.__client.replyReceived.connect(self.__handleServerMessage)
+        self.__client.pullStatus.connect(self.__handlePullStatus)
+        self.__client.pullError.connect(self.__handlePullError)
+
+        self.__client.errorOccurred.connect(self.__handleClientError)
+        self.__client.finished.connect(self.__handleClientFinished)
 
     @pyqtSlot(bool)
     def __handleServerStateChanged(self, ok):
@@ -605,8 +613,6 @@
         Private method to create the super menu and attach it to the super
         menu button.
         """
-        # TODO: implement the menu and menu methods
-        #       * Pull Model
         ###################################################################
         ## Menu with Chat History related actions
         ###################################################################
@@ -635,7 +641,9 @@
             self.tr("Show Model Library"), self.__showModelLibrary
         )
         self.__modelMenu.addSeparator()
-        self.__modelMenu.addAction(self.tr("Download Model"), self.__pullModel)
+        self.__pullModelAct = self.__modelMenu.addAction(
+            self.tr("Install Model"), self.__pullModel
+        )
         self.__removeModelAct = self.__modelMenu.addAction(
             self.tr("Remove Model"), self.__removeModel
         )
@@ -690,6 +698,7 @@
             self.__localServerProcess is not None and self.__localServerDialog is None
         )
 
+        self.__pullModelAct.setEnabled(not self.__pulling)
         self.__removeModelAct.setEnabled(bool(self.__availableModels))
 
     @pyqtSlot()
@@ -929,8 +938,58 @@
         """
         Private slot to download a model from the 'ollama' model library.
         """
-        # TODO: not implemented yet
-        pass
+        from .OllamaPullProgressDialog import OllamaPullProgressDialog
+
+        if self.__pulling:
+            # only one pull operation supported
+            return
+
+        model, ok = QInputDialog.getText(
+            self,
+            self.tr("Install Model"),
+            self.tr("Enter the name of the model to be installed:"),
+            QLineEdit.EchoMode.Normal,
+        )
+        if ok and model:
+            self.__pulling = True
+
+            if self.__pullProgressDialog is None:
+                self.__pullProgressDialog = OllamaPullProgressDialog(self)
+                self.__pullProgressDialog.abortPull.connect(self.__client.abortPull)
+
+            self.__pullProgressDialog.setModel(model)
+            self.__pullProgressDialog.clear()
+            self.__pullProgressDialog.show()
+
+            self.__client.pull(model)
+
+    @pyqtSlot(str, str, "unsigned long int", "unsigned long int")
+    def __handlePullStatus(self, status, idStr, total, completed):
+        """
+        Private slot to handle a pull 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
+        """
+        if self.__pullProgressDialog is not None:
+            self.__pullProgressDialog.setStatus(status, idStr, total, completed)
+
+    @pyqtSlot(str)
+    def __handlePullError(self, errMsg):
+        """
+        Private slot to handle an error during a pull operation.
+
+        @param errMsg error message
+        @type str
+        """
+        if self.__pullProgressDialog is not None:
+            self.__pullProgressDialog.showError(errMsg)
 
     @pyqtSlot()
     def __removeModel(self):
@@ -966,3 +1025,27 @@
                             " 'ollama' server.</p>"
                         ).format(modelName),
                     )
+
+    @pyqtSlot(str)
+    def __handleClientError(self, errMsg):
+        """
+        Private slot to handle an error message sent by the server.
+
+        @param errMsg error message
+        @type str
+        """
+        EricMessageBox.warning(
+            self,
+            self.tr("Network Error"),
+            errMsg,
+        )
+
+    @pyqtSlot()
+    def __handleClientFinished(self):
+        """
+        Private slot to handle the end of a client server interaction.
+        """
+        if self.__pullProgressDialog is not None and self.__pulling:
+            self.__pullProgressDialog.setFinished(True)
+            self.__pulling = False
+            self.__client.list()