Skip to content

Commit

Permalink
better handling of chat messages
Browse files Browse the repository at this point in the history
adds copy / delete controls to each message
  • Loading branch information
w4ffl35 committed Oct 25, 2024
1 parent 98ca98a commit a5c79b6
Show file tree
Hide file tree
Showing 8 changed files with 285 additions and 123 deletions.
32 changes: 25 additions & 7 deletions src/airunner/styles/dark_theme/styles.qss
Original file line number Diff line number Diff line change
Expand Up @@ -560,13 +560,6 @@ QTableWidget::item:selected {
border: 1px solid rgba(0, 132, 185, 50);
}

QWidget#message QPlainTextEdit {
border-radius: 5px;
border: 5px solid #1f1f1f;
background-color: #1f1f1f;
color: #ffffff;
}

QSplitter::handle::vertical {
border-top: 4px solid rgba(0, 132, 185, 50);
}
Expand Down Expand Up @@ -734,3 +727,28 @@ QMenu {
SwitchWidget:enabled[checked="true"] {
background-color: rgba(0, 132, 185, 100);
}

QScrollArea#chat_container #message_container {
background-color: #1f1f1f;
}

QScrollArea#chat_container #message_container[class="alternate"] {
background-color: #000;
}

QScrollArea#chat_container #message_container QTextEdit {
border: none;
background-color: transparent;
color: #ffffff;
}

QScrollArea#chat_container #message_container QPushButton {
border: none;
background-color: transparent;
color: #ffffff;
}

QScrollArea#chat_container #message_container QPushButton:hover {
background-color: #000;
border: 1px solid #1f1f1f;
}
15 changes: 12 additions & 3 deletions src/airunner/widgets/llm/chat_prompt_widget.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,8 +96,10 @@ def _set_conversation_widgets(self, messages):
name=message["name"],
message=message["content"],
is_bot=message["is_bot"],
message_id=message["id"],
first_message=True
)
QTimer.singleShot(100, self.scroll_to_bottom)

def on_hear_signal(self, data: dict):
transcription = data["transcription"]
Expand Down Expand Up @@ -189,11 +191,14 @@ def do_generate(self, image_override=None, prompt_override=None, callback=None,
return
self.generating = True

self.add_message_to_conversation(
widget = self.add_message_to_conversation(
name=self.chatbot.username,
message=self.prompt,
is_bot=False
)
if widget is not None:
QTimer.singleShot(100, widget.set_content_size)
QTimer.singleShot(100, self.scroll_to_bottom)

self.clear_prompt()
self.start_progress_bar()
Expand Down Expand Up @@ -312,7 +317,8 @@ def add_message_to_conversation(
name,
message,
is_bot,
first_message=True
first_message=True,
message_id=None
):
if not first_message:
# get the last widget from the scrollAreaWidgetContents.layout()
Expand All @@ -330,15 +336,18 @@ def add_message_to_conversation(

self.remove_spacer()

widget = None
if message != "":
widget = MessageWidget(name=name, message=message, is_bot=is_bot)
widget = MessageWidget(name=name, message=message, is_bot=is_bot, message_id=message_id, conversation_id=self.conversation_id)
self.ui.scrollAreaWidgetContents.layout().addWidget(widget)

self.add_spacer()

# automatically scroll to the bottom of the scrollAreaWidgetContents
self.scroll_to_bottom()

return widget

def remove_spacer(self):
if self.spacer is not None:
self.ui.scrollAreaWidgetContents.layout().removeItem(self.spacer)
Expand Down
66 changes: 49 additions & 17 deletions src/airunner/widgets/llm/message_widget.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
from airunner.data.models.settings_models import Message, Conversation
from airunner.enums import SignalCode
from airunner.widgets.base_widget import BaseWidget
from airunner.widgets.llm.templates.message_ui import Ui_message

from PySide6.QtGui import QTextCursor, QFontDatabase, QFont
from PySide6.QtWidgets import QTextEdit
from PySide6.QtWidgets import QTextEdit, QApplication, QWidget
from PySide6.QtGui import QFontMetrics
from PySide6.QtCore import Qt, QSize
from PySide6.QtCore import Qt, QSize, Slot, QEvent, QTimer
from PySide6.QtCore import Signal


Expand All @@ -22,30 +23,48 @@ class MessageWidget(BaseWidget):
def __init__(self, *args, **kwargs):
self.name = kwargs.pop("name")
self.message = kwargs.pop("message")
self.conversation_id = kwargs.pop("conversation_id")
self.message_id = kwargs.pop("message_id")
if self.message_id is None:
session = self.session
message = session.query(Message).filter(Message.conversation_id == self.conversation_id).order_by(
Message.id.desc()).first()
if message is not None:
self.message_id = message.id
self.is_bot = kwargs.pop("is_bot")
super().__init__(*args, **kwargs)
self.ui.content.setReadOnly(True)
self.ui.content.insertPlainText(self.message)
self.ui.content.document().contentsChanged.connect(self.sizeChange)
name = self.name
if self.is_bot:
self.ui.bot_name.show()
self.ui.bot_name.setText(f"{name}")
self.ui.bot_name.setStyleSheet("font-weight: normal;")
self.ui.user_name.hide()
else:
self.ui.user_name.show()
self.ui.user_name.setText(f"{name}")
self.ui.user_name.setStyleSheet("font-weight: normal;")
self.ui.bot_name.hide()

self.ui.content.setStyleSheet("border-radius: 5px; border: 5px solid #1f1f1f; background-color: #1f1f1f; color: #ffffff;")

self.ui.user_name.setText(f"{self.name}")
self.register(SignalCode.APPLICATION_SETTINGS_CHANGED_SIGNAL, self.on_application_settings_changed_signal)
self.font_family = None
self.font_size = None
self.set_chat_font()

if self.is_bot:
self.ui.message_container.setProperty("class", "alternate")

self.ui.copy_button.setVisible(False)
self.ui.delete_button.setVisible(False)
self.ui.message_container.installEventFilter(self)
self.set_cursor(Qt.CursorShape.ArrowCursor)

def set_cursor(self, cursor_type):
self.ui.message_container.setCursor(cursor_type)
for child in self.ui.message_container.findChildren(QWidget):
child.setCursor(cursor_type)

def eventFilter(self, obj, event):
if obj == self.ui.message_container:
if event.type() == QEvent.Type.Enter:
self.ui.copy_button.setVisible(True)
self.ui.delete_button.setVisible(True)
elif event.type() == QEvent.Type.Leave:
self.ui.copy_button.setVisible(False)
self.ui.delete_button.setVisible(False)
return super().eventFilter(obj, event)

def on_application_settings_changed_signal(self):
self.set_chat_font()

Expand All @@ -71,7 +90,7 @@ def set_chat_font(self):
def set_content_size(self):
doc_height = self.ui.content.document().size().height()
doc_width = self.ui.content.document().size().width()
self.setMinimumHeight(int(doc_height) + 25)
self.setMinimumHeight(int(doc_height) + 45)
self.setMinimumWidth(int(doc_width))

def sizeChange(self):
Expand All @@ -98,3 +117,16 @@ def update_message(self, text):

self.ui.content.setPlainText(self.message)

@Slot()
def delete(self):
session = self.session
message = session.query(Message).filter(Message.id == self.message_id).first()
session.delete(message)
session.commit()
self.setParent(None)
self.deleteLater()

@Slot()
def copy(self):
clipboard = QApplication.clipboard()
clipboard.setText(self.message)
6 changes: 6 additions & 0 deletions src/airunner/widgets/llm/templates/chat_prompt.ui
Original file line number Diff line number Diff line change
Expand Up @@ -225,6 +225,9 @@
</property>
<widget class="QWidget" name="verticalLayoutWidget">
<layout class="QVBoxLayout" name="verticalLayout">
<property name="spacing">
<number>0</number>
</property>
<item>
<widget class="QScrollArea" name="chat_container">
<property name="frameShape">
Expand Down Expand Up @@ -255,6 +258,9 @@
<property name="bottomMargin">
<number>0</number>
</property>
<property name="spacing">
<number>0</number>
</property>
<item row="0" column="0">
<widget class="QTextEdit" name="conversation">
<property name="font">
Expand Down
2 changes: 2 additions & 0 deletions src/airunner/widgets/llm/templates/chat_prompt_ui.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,7 @@ def setupUi(self, chat_prompt):
self.verticalLayoutWidget = QWidget(self.chat_prompt_splitter)
self.verticalLayoutWidget.setObjectName(u"verticalLayoutWidget")
self.verticalLayout = QVBoxLayout(self.verticalLayoutWidget)
self.verticalLayout.setSpacing(0)
self.verticalLayout.setObjectName(u"verticalLayout")
self.verticalLayout.setContentsMargins(0, 0, 0, 0)
self.chat_container = QScrollArea(self.verticalLayoutWidget)
Expand All @@ -125,6 +126,7 @@ def setupUi(self, chat_prompt):
self.scrollAreaWidgetContents.setObjectName(u"scrollAreaWidgetContents")
self.scrollAreaWidgetContents.setGeometry(QRect(0, 0, 593, 321))
self.gridLayout_2 = QGridLayout(self.scrollAreaWidgetContents)
self.gridLayout_2.setSpacing(0)
self.gridLayout_2.setObjectName(u"gridLayout_2")
self.gridLayout_2.setContentsMargins(0, 0, 0, 0)
self.conversation = QTextEdit(self.scrollAreaWidgetContents)
Expand Down
Loading

0 comments on commit a5c79b6

Please sign in to comment.