z-bar-qt/Plugins/ZShell/Llm/session.cpp

#include "session.hpp"

#include "attachment.hpp"
#include "chatstore.hpp"
#include "clipboard.hpp"
#include "llmclient.hpp"
#include "prompts.hpp"
#include "workspace.hpp"

#include <QDebug>
#include <QDateTime>
#include <QDir>
#include <QFile>

namespace ZShell::llm {

namespace {

constexpr int kMaxStoredMessages = 200;

} // namespace

ChatSession::ChatSession(ChatStore* store, const QString& id, QObject* parent)
	: QObject(parent), m_store(store), m_id(id) {
	m_model = new ChatMessageModel(this, this);
	m_workspace = new Workspace(this, this);

	connect(
		m_model,
		&QAbstractListModel::rowsInserted,
		this,
		&ChatSession::onModelRowsChanged);
	connect(
		m_model,
		&QAbstractListModel::rowsRemoved,
		this,
		&ChatSession::onModelRowsChanged);
	connect(
		m_model,
		&QAbstractListModel::modelReset,
		this,
		&ChatSession::onModelRowsChanged);
}

void ChatSession::onModelRowsChanged() {
	setCount(m_model->rowCount());
}

void ChatSession::setTitle(const QString& value) {
	if (m_title == value) return;
	m_title = value;
	Q_EMIT titleChanged();
	persist();
}

void ChatSession::setIcon(const QString& value) {
	if (m_icon == value) return;
	m_icon = value;
	Q_EMIT iconChanged();
	if (m_store) m_store->saveFlags(this);
}

void ChatSession::setUpdatedAt(qint64 value) {
	if (m_updatedAt == value) return;
	m_updatedAt = value;
	Q_EMIT updatedAtChanged();
	if (m_store) m_store->reposition(this);
}

void ChatSession::setPinned(bool value) {
	if (m_pinned == value) return;
	m_pinned = value;
	Q_EMIT pinnedChanged();
	if (m_store) m_store->saveFlags(this);
}

void ChatSession::setUsageTokens(int value) {
	if (m_usageTokens == value) return;
	m_usageTokens = value;
	Q_EMIT usageTokensChanged();
}

void ChatSession::setCount(int value) {
	if (m_messageCount == value) return;
	m_messageCount = value;
	Q_EMIT messageCountChanged();
}

void ChatSession::hydrate(
	const QString& title,
	const QString& icon,
	qint64 createdAt,
	qint64 updatedAt,
	int messageCount,
	bool pinned) {
	m_title = title;
	m_icon = icon;
	m_createdAt = createdAt;
	m_updatedAt = updatedAt;
	m_messageCount = messageCount;
	m_pinned = pinned;
}

void ChatSession::ensureLoaded() {
	if (m_loadRequested) return;
	m_loadRequested = true;
	if (m_store) m_store->loadMessagesInto(this);
}

ChatMessageModel* ChatSession::messagesModel() {
	ensureLoaded();
	return m_model;
}

void ChatSession::markLoaded() {
	if (m_loaded) return;
	m_loaded = true;
	Q_EMIT loaded();
}

void ChatSession::adoptMessages(QList<ChatMessage*> messages) {
	m_model->loadMessages(std::move(messages));
}

void ChatSession::persist() {
	if (m_store) m_store->persist(this);
}

ChatMessage* ChatSession::appendNewest(
	ChatMessage::Role role, const QString& content, qint64 timestamp) {
	return m_model->appendNewest(role, content, timestamp);
}

void ChatSession::removeMessage(ChatMessage* message) {
	m_model->removeMessage(message);
}

void ChatSession::clearMessages() {
	m_model->clear();
}

void ChatSession::startGeneration(ChatMessage* target) {
	ChatGeneration* generation = target->activeGeneration();
	LlmClient* clientObject = client();
	if (!generation || !clientObject) return;

	if (isLoaded()) {
		clientObject->startGeneration(this, generation);
		return;
	}

	connect(
		this,
		&ChatSession::loaded,
		clientObject,
		[this, generation, clientObject]() {
			clientObject->startGeneration(this, generation);
		});
	persist();
}

LlmClient* ChatSession::client() const {
	return m_store ? m_store->client() : nullptr;
}

void ChatSession::setAttachmentError(const QString& value) {
	if (m_attachmentError == value) return;
	m_attachmentError = value;
	Q_EMIT attachmentErrorChanged();
}

QVariantList ChatSession::addAttachments(const QStringList& paths) {
	if (!m_workspace) return {};
	setAttachmentError(QString());

	QVariantList result;
	QString error;
	for (Attachment* attachment : m_workspace->importFiles(paths, error))
		result.append(QVariant::fromValue(attachment));
	if (!error.isEmpty()) setAttachmentError(error);
	return result;
}

bool ChatSession::clipboardHasAttachment() {
	return clipboard::hasAttachment();
}

QVariantList ChatSession::pasteAttachments() {
	QString error;
	const QStringList paths = clipboard::attachmentPaths(error);
	if (paths.isEmpty()) {
		if (!error.isEmpty()) setAttachmentError(error);
		return {};
	}

	const QVariantList attachments = addAttachments(paths);

	for (const QString& path : paths)
		if (path.startsWith(QDir::tempPath())) QFile::remove(path);

	return attachments;
}

void ChatSession::sendMessage(const QString& text) {
	sendMessage(text, QVariantList());
}

void ChatSession::sendMessage(
	const QString& text, const QVariantList& attachments) {
	QList<Attachment*> attachmentList;
	for (const QVariant& variant : attachments)
		if (auto* attachment =
				qobject_cast<Attachment*>(variant.value<QObject*>()))
			attachmentList.append(attachment);
	const QString trimmed = text.trimmed();
	if (trimmed.isEmpty() && attachmentList.isEmpty()) return;
	if (client() && client()->busy()) return;

	if (!m_committed && m_store) m_store->commitDraft(this);

	const qint64 now = QDateTime::currentMSecsSinceEpoch();
	ChatMessage* userMessage =
		appendNewest(ChatMessage::Role::User, trimmed, now);
	if (!attachmentList.isEmpty()) userMessage->setAttachments(attachmentList);
	while (m_model->rowCount() > kMaxStoredMessages)
		m_model->removeMessage(m_model->at(m_model->rowCount() - 1));

	if (m_title.isEmpty()) {
		QString titleSource = trimmed;
		if (titleSource.isEmpty() && !attachmentList.isEmpty())
			titleSource = attachmentList.first()->name();
		m_title = prompts::cleanTitle(titleSource);
		Q_EMIT titleChanged();
		persist();

		if (LlmClient* clientObject = client()) {
			clientObject->requestTitle(this, titleSource);
			clientObject->requestIcon(this, titleSource);
		}
	}

	ChatMessage* assistant =
		appendNewest(ChatMessage::Role::Assistant, QString(), now);
	startGeneration(assistant);
}

void ChatSession::retry(ChatMessage* target) {
	if (!target || target->role() != ChatMessage::Role::Assistant) return;
	if (client() && client()->busy()) return;
	const int row = m_model->rowOf(target);
	if (row < 0) return;

	m_model->removeRange(0, row - 1);
	if (m_model->rowCount() < 2) return;

	target->appendGeneration(QDateTime::currentMSecsSinceEpoch());
	startGeneration(target);
}

void ChatSession::continueFrom(ChatMessage* message) {
	if (!message) return;
	if (client() && client()->busy()) return;
	const int row = m_model->rowOf(message);
	if (row < 0) return;

	m_model->removeRange(0, row - 1);

	const qint64 now = QDateTime::currentMSecsSinceEpoch();
	if (message->role() == ChatMessage::Role::Assistant) {
		if (m_model->rowCount() < 2) return;
		message->appendGeneration(now);
		startGeneration(message);
	} else {
		ChatMessage* assistant =
			appendNewest(ChatMessage::Role::Assistant, QString(), now);
		startGeneration(assistant);
	}
}

void ChatSession::clear() {
	if (!m_store) return;
	m_store->requestClear(this);
}

} // namespace ZShell::llm