Add Delete message(s) support

Fixes: cristianadam#8

Author: cristianadam

llamachateditor.cpp

@@ -156,6 +156,16 @@ ChatEditor::ChatEditor()
             &ChatManager::messageExtraUpdated,
             this,
             &ChatEditor::onMessageExtraUpdated);
+    connect(&ChatManager::instance(),
+            &ChatManager::messageDeleted,
+            this,
+            [this](const QString &convId) {
+                if (convId == m_viewingConvId) {
+                    // Refresh the entire view to update the tree
+                    ViewingChat chat = ChatManager::instance().getViewingChat(m_viewingConvId);
+                    refreshMessages(chat.messages, chat.conv.currNode);
+                }
+            });
 
     connect(m_input, &ChatInput::sendRequested, this, &ChatEditor::onSendRequested);
     connect(m_input, &ChatInput::stopRequested, this, &ChatEditor::onStopRequested);
@@ -425,6 +435,7 @@ void ChatEditor::refreshMessages(const QVector<Message> &messages, qint64 leafNo
         connect(w, &ChatMessage::editRequested, this, &ChatEditor::onEditRequested);
         connect(w, &ChatMessage::regenerateRequested, this, &ChatEditor::onRegenerateRequested);
         connect(w, &ChatMessage::siblingChanged, this, &ChatEditor::onSiblingChanged);
+        connect(w, &ChatMessage::deleteRequested, this, &ChatEditor::onDeleteMessageRequested);
         m_messageLayout->addWidget(w);
         m_messageWidgets.append(w);
     }
@@ -515,6 +526,7 @@ void ChatEditor::onMessageAppended(const Message &msg, qint64 pendingId)
         connect(w, &ChatMessage::editRequested, this, &ChatEditor::onEditRequested);
         connect(w, &ChatMessage::regenerateRequested, this, &ChatEditor::onRegenerateRequested);
         connect(w, &ChatMessage::siblingChanged, this, &ChatEditor::onSiblingChanged);
+        connect(w, &ChatMessage::deleteRequested, this, &ChatEditor::onDeleteMessageRequested);
 
         m_messageLayout->addWidget(w);
         m_messageWidgets.append(w);
@@ -565,6 +577,7 @@ void ChatEditor::onPendingMessageChanged(const Message &pm)
         connect(w, &ChatMessage::editRequested, this, &ChatEditor::onEditRequested);
         connect(w, &ChatMessage::regenerateRequested, this, &ChatEditor::onRegenerateRequested);
         connect(w, &ChatMessage::siblingChanged, this, &ChatEditor::onSiblingChanged);
+        connect(w, &ChatMessage::deleteRequested, this, &ChatEditor::onDeleteMessageRequested);
 
         // Insert speed label after the assistant widget
         if (settings().showTokensPerSecond.value()) {
@@ -678,6 +691,34 @@ void ChatEditor::onServerPropsUpdated()
     }
 }
 
+void ChatEditor::onDeleteMessageRequested(const Message &msg)
+{
+    auto [userMessages, assistantMessages] = ChatManager::instance().getBranchStats(msg.convId,
+                                                                                    msg.id);
+    int totalMessages = userMessages + assistantMessages;
+    if (totalMessages > 1) {
+        QString warningText = Tr::tr("This will delete %1 messages including: %2 user messages and "
+                                     "%3 assistant responses ...")
+                                  .arg(totalMessages)
+                                  .arg(userMessages)
+                                  .arg(assistantMessages);
+
+        if (QMessageBox::question(widget(), Tr::tr("Confirm Branch Deletion"), warningText)
+            != QMessageBox::Yes) {
+            return;
+        }
+    } else {
+        if (QMessageBox::question(widget(),
+                                  Tr::tr("Delete Message"),
+                                  Tr::tr("Are you sure you want to delete this message?"))
+            != QMessageBox::Yes) {
+            return;
+        }
+    }
+
+    ChatManager::instance().deleteMessageBranch(msg.convId, msg.id);
+}
+
 void ChatEditor::startSearch()
 {
     // Show the floating search toolbar

llamachateditor.h

@@ -62,6 +62,7 @@ public slots:
     void onRegenerateRequested(const LlamaCpp::Message &msg);
     void onSiblingChanged(qint64 siblingId);
     void onServerPropsUpdated();
+    void onDeleteMessageRequested(const LlamaCpp::Message &msg);
 
     void startSearch();
     void nextSearchResult();

llamachatmanager.cpp

@@ -917,4 +917,52 @@ void ChatManager::renameConversation(const QString &convId, const QString &name)
     m_storage->renameConversation(convId, name);
 }
 
+void ChatManager::deleteMessageBranch(const QString &convId, qint64 msgId)
+{
+    if (m_storage->deleteMessageBranch(msgId))
+        emit messageDeleted(convId);
+}
+
+QPair<int, int> ChatManager::getBranchStats(const QString &convId, qint64 msgId) const
+{
+    QPair<int, int> stats{0, 0};
+    int &userMessages = stats.first;
+    int &assistantMessages = stats.second;
+
+    QVector<Message> allMessages = m_storage->getMessages(convId);
+
+    // Map for quick lookup
+    QHash<qint64, Message> map;
+    for (const auto &m : std::as_const(allMessages)) {
+        map.insert(m.id, m);
+    }
+
+    if (!map.contains(msgId))
+        return stats;
+
+    // Traverse the branch using a stack (DFS)
+    QVector<qint64> stack;
+    stack.push_back(msgId);
+
+    while (!stack.isEmpty()) {
+        qint64 currentId = stack.takeLast();
+        if (!map.contains(currentId))
+            continue;
+
+        const Message &m = map[currentId];
+        if (m.role == "user") {
+            ++userMessages;
+        } else {
+            ++assistantMessages;
+        }
+
+        // Add children to stack to continue traversal
+        for (qint64 childId : m.children) {
+            stack.push_back(childId);
+        }
+    }
+
+    return stats;
+}
+
 } // namespace LlamaCpp

llamachatmanager.h

@@ -55,6 +55,8 @@ class ChatManager : public QObject
     Conversation createConversation(const QString &name);
     void deleteConversation(const QString &convId);
     void renameConversation(const QString &convId, const QString &name);
+    void deleteMessageBranch(const QString &convId, qint64 msgId);
+    QPair<int,int> getBranchStats(const QString &convId, qint64 msgId) const;
 
     QList<Conversation> allConversations();
 
@@ -88,6 +90,7 @@ class ChatManager : public QObject
                                    qint64 leafNodeId,
                                    const QStringList &quetions);
     void messageExtraUpdated(const LlamaCpp::Message &msg, const QList<QVariantMap> &newExtra);
+    void messageDeleted(const QString &convId);
 
 private:
     explicit ChatManager(QObject *parent = nullptr);

llamachatmessage.cpp

@@ -222,6 +222,13 @@ void ChatMessage::buildUI()
     connect(m_copyButton, &QToolButton::clicked, this, &ChatMessage::onCopyClicked);
     actionLayout->addWidget(m_copyButton);
 
+    m_deleteButton = new QToolButton(this);
+    m_deleteButton->setText("K");
+    m_deleteButton->setToolTip(Tr::tr("Delete this message"));
+    m_deleteButton->setVisible(!m_haveToolCalls);
+    connect(m_deleteButton, &QToolButton::clicked, this, &ChatMessage::onDeleteClicked);
+    actionLayout->addWidget(m_deleteButton);
+
     m_mainLayout->addWidget(m_bubble);
     m_mainLayout->addLayout(actionLayout);
 
@@ -271,6 +278,7 @@ void ChatMessage::messageCompleted(bool completed)
         // Normal assistant – show buttons only when the answer is finished.
         m_regenButton->setVisible(completed && !haveToolCalls());
         m_copyButton->setVisible(completed && !haveToolCalls());
+        m_deleteButton->setVisible(completed && !haveToolCalls());
 
         renderMarkdown(m_msg.content, completed);
     } else if (m_isTool) {
@@ -577,6 +585,11 @@ void ChatMessage::onSaveToDisk(const QString &fileName, const QString &verbatimC
     sourceFile.writeFileContents(verbatimCode.toUtf8());
 }
 
+void ChatMessage::onDeleteClicked()
+{
+    emit deleteRequested(m_msg);
+}
+
 void ChatMessage::updateUI()
 {
     m_siblingLabel->setText(QString("%1/%2").arg(m_siblingIdx).arg(m_siblingLeafIds.size()));

llamachatmessage.h

@@ -43,6 +43,7 @@ class ChatMessage : public QWidget
     void regenerateRequested(const Message &msg);
     void editRequested(const Message &msg);
     void siblingChanged(qint64 siblingId);
+    void deleteRequested(const Message &msg);
 
 private slots:
     void onCopyClicked();
@@ -52,6 +53,7 @@ private slots:
     void onNextSiblingClicked();
     void onCopyToClipboard(const QString &verbatimCode, const QString &highlightedCode);
     void onSaveToDisk(const QString &fileName, const QString &verbatimCode);
+    void onDeleteClicked();
 
 private:
     void buildUI();
@@ -77,6 +79,7 @@ private slots:
     QToolButton *m_prevButton{nullptr};
     QToolButton *m_nextButton{nullptr};
     QToolButton *m_attachedFiles{nullptr};
+    QToolButton *m_deleteButton{nullptr};
     QLabel *m_siblingLabel{nullptr};
     MarkdownLabel *m_markdownLabel{nullptr};
     QVBoxLayout *m_mainLayout{nullptr};