chore: restructure llm plugin to be more coherent & fix various bugs in llm ui

Plugins/ZShell/Llm/message.cpp

@@ -1,25 +1,36 @@
 #include "message.hpp"
 
-#include "messagemodel.hpp"
 #include "session.hpp"
 
 namespace ZShell::llm {
 
-namespace {
+ChatMessage::ChatMessage(
+	ChatSession* session, Role role, qint64 timestamp, QObject* parent)
+	: QObject(parent)
+	, m_session(session)
+	, m_role(role)
+	, m_timestamp(timestamp) {}
 
-ChatSession* sessionOf(const ChatMessage* message) {
-	if (auto* model = qobject_cast<ChatMessageModel*>(message->parent()))
-		return model->session();
-	return nullptr;
+QString ChatMessage::content() const {
+	if (m_role == Role::User) return m_content;
+	const ChatGeneration* generation = activeGeneration();
+	return generation ? generation->content() : QString();
 }
 
-} // namespace
-
-ChatMessage::ChatMessage(Role role, qint64 timestamp, QObject* parent)
-	: QObject(parent), m_role(role), m_timestamp(timestamp) {}
+void ChatMessage::setContent(const QString& value) {
+	if (m_role != Role::User || m_content == value) return;
+	m_content = value;
+	Q_EMIT contentChanged();
+}
 
 ChatGeneration* ChatMessage::addGeneration(qint64 timestamp) {
-	auto* generation = new ChatGeneration(timestamp, this);
+	if (m_role != Role::Assistant) return nullptr;
+
+	auto* generation = new ChatGeneration(this, timestamp, this);
+	connect(
+		generation, &ChatGeneration::contentChanged, this, [this, generation]() {
+			if (generation == activeGeneration()) Q_EMIT contentChanged();
+		});
 	m_generations.append(generation);
 	if (m_active < 0) m_active = static_cast<int>(m_generations.size() - 1);
 	Q_EMIT generationsChanged();
@@ -27,7 +38,7 @@ ChatGeneration* ChatMessage::addGeneration(qint64 timestamp) {
 }
 
 ChatGeneration* ChatMessage::appendGeneration(qint64 timestamp) {
-	auto* generation = addGeneration(timestamp);
+	ChatGeneration* generation = addGeneration(timestamp);
 	setActiveInternal(static_cast<int>(m_generations.size() - 1));
 	return generation;
 }
@@ -44,13 +55,17 @@ void ChatMessage::removeGeneration(ChatGeneration* generation) {
 		m_active = static_cast<int>(m_generations.size() - 1);
 	}
 	Q_EMIT generationsChanged();
-	if (wasActive) Q_EMIT activeGenerationChanged();
+	if (wasActive) {
+		Q_EMIT activeGenerationChanged();
+		Q_EMIT contentChanged();
+	}
 }
 
 void ChatMessage::setActiveInternal(int index) {
 	if (index < 0 || index >= m_generations.size() || index == m_active) return;
 	m_active = index;
 	Q_EMIT activeGenerationChanged();
+	Q_EMIT contentChanged();
 }
 
 void ChatMessage::setActiveGeneration(int index) {
@@ -72,17 +87,25 @@ void ChatMessage::setAttachments(QList<Attachment*> attachments) {
 }
 
 void ChatMessage::edit(const QString& newContent) {
-	if (auto* generation = activeGeneration())
+	if (m_role == Role::User) {
+		setContent(newContent);
+	} else if (ChatGeneration* generation = activeGeneration()) {
 		generation->setContent(newContent);
-	if (auto* session = sessionOf(this)) session->persist();
+	} else {
+		return;
+	}
+
+	if (m_session) m_session->persist();
 }
 
 void ChatMessage::retry() {
-	if (auto* session = sessionOf(this)) session->retry(this);
+	if (m_role != Role::Assistant || !m_session) return;
+	m_session->retry(this);
 }
 
 void ChatMessage::generate() {
-	if (auto* session = sessionOf(this)) session->continueFrom(this);
+	if (m_role != Role::Assistant || !m_session) return;
+	m_session->continueFrom(this);
 }
 
 } // namespace ZShell::llm