--- a/OllamaInterface/OllamaWidget.py Mon Aug 05 18:37:16 2024 +0200
+++ b/OllamaInterface/OllamaWidget.py Tue Aug 06 18:18:39 2024 +0200
@@ -9,13 +9,15 @@
import json
import os
-from PyQt6.QtCore import Qt, pyqtSlot
+from PyQt6.QtCore import Qt, QTimer, pyqtSlot
from PyQt6.QtWidgets import QInputDialog, QLineEdit, QVBoxLayout, QWidget
from eric7 import Globals
from eric7.EricGui import EricPixmapCache
from eric7.EricWidgets import EricMessageBox
+from eric7.EricWidgets.EricApplication import ericApp
+from .OllamaChatWidget import OllamaChatWidget
from .OllamaClient import OllamaClient
from .OllamaHistoryWidget import OllamaHistoryWidget
from .Ui_OllamaWidget import Ui_OllamaWidget
@@ -51,9 +53,16 @@
else:
self.layout().setContentsMargins(0, 0, 0, 0)
+ iconSuffix = "-dark" if ericApp().usesDarkPalette() else "-light"
+
self.ollamaMenuButton.setIcon(EricPixmapCache.getIcon("superMenu"))
self.reloadModelsButton.setIcon(EricPixmapCache.getIcon("reload"))
self.newChatButton.setIcon(EricPixmapCache.getIcon("plus"))
+ self.sendButton.setIcon(
+ EricPixmapCache.getIcon(
+ os.path.join("OllamaInterface", "icons", "send{0}".format(iconSuffix))
+ )
+ )
self.__chatHistoryLayout = QVBoxLayout()
self.historyScrollWidget.setLayout(self.__chatHistoryLayout)
@@ -66,7 +75,11 @@
self.__connectClient()
+ self.sendButton.clicked.connect(self.__sendMessage)
+ self.messageEdit.returnPressed.connect(self.__sendMessage)
+
self.__loadHistory()
+ self.__updateMessageEditState()
def __connectClient(self):
"""
@@ -75,6 +88,7 @@
self.__client.serverStateChanged.connect(self.__handleServerStateChanged)
self.__client.serverVersion.connect(self.__setHeaderLabel)
self.__client.modelsList.connect(self.__populateModelSelector)
+ self.__client.replyReceived.connect(self.__handleServerMessage)
@pyqtSlot(bool)
def __handleServerStateChanged(self, ok):
@@ -117,30 +131,6 @@
"""
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.
############################################################################
@@ -187,35 +177,46 @@
@param jsonStr string containing JSON serialize chat history data (defaults
to None)
@type str (optional)
+ @return reference to the created history widget
+ @rtype OllamaHistoryWidget
"""
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.
+ QTimer.singleShot(0, self.__scrollHistoryToBottom)
+
+ return history
- @param uid ID of the history widget
+ @pyqtSlot()
+ def __scrollHistoryToBottom(self):
+ """
+ Private slot to scroll the history widget to the bottom.
+ """
+ scrollbar = self.historyScrollArea.verticalScrollBar()
+ scrollbar.setMaximum(self.historyScrollWidget.height())
+ scrollbar.setValue(scrollbar.maximum())
+
+ def __findHistoryWidget(self, hid):
+ """
+ Private method to find the widget of a given chat history ID.
+
+ @param hid ID of the chat history
@type str
- @return index of the history widget
- @rtype int
+ @return reference to the chat history widget
+ @rtype OllamaHistoryWidget
"""
for index in range(self.__chatHistoryLayout.count() - 1):
widget = self.__chatHistoryLayout.itemAt(index).widget()
- if widget.getId() == uid:
- return index
+ if widget.getId() == hid:
+ return widget
return None
@@ -237,8 +238,8 @@
entries = {}
for index in range(self.__chatHistoryLayout.count() - 1):
widget = self.__chatHistoryLayout.itemAt(index).widget()
- uid = widget.getId()
- entries[uid] = widget.saveToJson()
+ hid = widget.getId()
+ entries[hid] = widget.saveToJson()
# step 2: save the collected chat histories
filePath = self.__historyFilePath()
@@ -279,8 +280,8 @@
return
# step 2: create history widgets
- for uid in entries:
- self.__createHistoryWidget("", "", jsonStr=entries[uid])
+ for hid in entries:
+ self.__createHistoryWidget("", "", jsonStr=entries[hid])
def clearHistory(self):
"""
@@ -295,31 +296,207 @@
self.__saveHistory()
@pyqtSlot(str)
- def __deleteHistory(self, uid):
+ def __deleteHistory(self, hid):
"""
Private slot to delete the history with the given ID.
- @param uid ID of the history to be deleted
+ @param hid ID of the history to be deleted
@type str
"""
- widgetIndex = self.__findHistoryWidgetIndex(uid)
- if widgetIndex is not None:
+ widget = self.__findHistoryWidget(hid)
+ if widget is not None:
+ widgetIndex = self.__chatHistoryLayout.indexOf(widget)
item = self.__chatHistoryLayout.takeAt(widgetIndex)
if item is not None:
item.widget().deleteLater()
- self.__saveHistory()
+ self.__saveHistory()
+
+ self.__removeChatWidget(hid)
+
+ #######################################################################
+ ## Chat related methods below
+ #######################################################################
+
+ def __findChatWidget(self, hid):
+ """
+ Private method to find a chat widget given a chat history ID.
+
+ @param hid chat history ID
+ @type str
+ @return reference to the chat widget related to the given ID
+ @rtype OllamaChatWidget
+ """
+ for index in range(self.chatStackWidget.count()):
+ widget = self.chatStackWidget.widget(index)
+ if widget.getHistoryId() == hid:
+ return widget
+
+ return None
+
+ @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:
+ historyWidget = self.__createHistoryWidget(title, model)
+ hid = historyWidget.getId()
+ chatWidget = OllamaChatWidget(hid=hid, title=title, model=model)
+ index = self.chatStackWidget.addWidget(chatWidget)
+ self.chatStackWidget.setCurrentIndex(index)
+
+ self.__updateMessageEditState()
+ self.messageEdit.setFocus(Qt.FocusReason.OtherFocusReason)
@pyqtSlot(str)
- def __newChatWithHistory(self, uid):
+ def __newChatWithHistory(self, hid):
"""
Private slot to start a new chat using a previously saved history.
- @param uid ID of the history to be used
+ @param hid ID of the history to be used
+ @type str
+ """
+ chatWidget = self.__findChatWidget(hid)
+ if chatWidget is None:
+ historyWidget = self.__findHistoryWidget(hid)
+ if historyWidget is None:
+ # Oops, treat it as a new chat.
+ self.on_newChatButton_clicked()
+ return
+
+ chatWidget = OllamaChatWidget(
+ hid=hid, title=historyWidget.getTitle(), model=historyWidget.getModel()
+ )
+ index = self.chatStackWidget.addWidget(chatWidget)
+ self.chatStackWidget.setCurrentIndex(index)
+ for message in historyWidget.getMessages():
+ chatWidget.addMessage(role=message["role"], message=message["content"])
+ else:
+ # simply switch to the already existing chatWidget
+ self.chatStackWidget.setCurrentWidget(chatWidget)
+
+ self.__updateMessageEditState()
+ self.messageEdit.setFocus(Qt.FocusReason.OtherFocusReason)
+
+ def __removeChatWidget(self, hid):
+ """
+ Private method to remove a chat widget given its chat history ID.
+
+ @param hid chat history ID
+ @type str
+ """
+ widget = self.__findChatWidget(hid)
+ if widget is not None:
+ self.chatStackWidget.removeWidget(widget)
+
+ @pyqtSlot()
+ def __updateMessageEditState(self):
+ """
+ Private slot to set the enabled state of the message line edit and the send
+ button.
+ """
+ chatActive = bool(self.chatStackWidget.count())
+ hasText = bool(self.messageEdit.text())
+
+ self.messageEdit.setEnabled(chatActive)
+ self.sendButton.setEnabled(chatActive and hasText)
+
+ @pyqtSlot(str)
+ def on_messageEdit_textChanged(self, msg):
+ """
+ Private slot to handle a change of the entered message.
+
+ @param msg text of the message line edit
@type str
"""
- # TODO: not implemented yet
- pass
+ self.sendButton.setEnabled(bool(msg))
+
+ @pyqtSlot()
+ def __sendMessage(self):
+ """
+ Private method to send the given message of the current chat to the
+ 'ollama' server.
+
+ This sends the message with context (i.e. the history of the current chat).
+ """
+ msg = self.messageEdit.text()
+ if not msg:
+ # empty message => ignore
+ return
+
+ if not bool(self.chatStackWidget.count()):
+ # no current stack => ignore
+ return
+
+ # 1. determine hid of the current chat via chat stack widget
+ chatWidget = self.chatStackWidget.currentWidget()
+ hid = chatWidget.getHistoryId()
+
+ # 2. get chat history widget via hid from chat history widget
+ historyWidget = self.__findHistoryWidget(hid)
+ if historyWidget is not None:
+ # 3. append the message to the history
+ historyWidget.addToMessages("user", msg)
+
+ # 4. get the complete messages list from the history
+ messages = historyWidget.getMessages()
+
+ # 5. add the message to the current chat and an empty one
+ # for the response
+ chatWidget.addMessage("user", msg)
+ chatWidget.addMessage("assistant", "")
+
+ # 6. send the request via the client (non-streaming (?))
+ model = historyWidget.getModel()
+ self.__client.chat(
+ model=model,
+ messages=messages,
+ streaming=self.__plugin.getPreferences("StreamingChatResponse"),
+ )
+
+ # 7. clear the message editor and give input focus back
+ self.messageEdit.clear()
+ self.messageEdit.setFocus(Qt.FocusReason.OtherFocusReason)
+
+ @pyqtSlot(str, str, bool)
+ def __handleServerMessage(self, content, role, done):
+ """
+ Private slot handling an 'ollama' server chat response.
+
+ @param content message sent by the server
+ @type str
+ @param role role name
+ @type str
+ @param done flag indicating the last chat response
+ @type bool
+ """
+ if not bool(self.chatStackWidget.count()):
+ # no current stack => ignore
+ return
+
+ chatWidget = self.chatStackWidget.currentWidget()
+ chatWidget.appendMessage(content)
+ if done:
+ hid = chatWidget.getHistoryId()
+ historyWidget = self.__findHistoryWidget(hid)
+ if historyWidget is not None:
+ historyWidget.addToMessages(role, chatWidget.getRecentMessage())
#######################################################################
## Menu related methods below
