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

#include "session.hpp"

#include "chatstore.hpp"
#include "llmclient.hpp"

#include <QDebug>
#include <QDateTime>
#include <QtGlobal>

namespace ZShell::llm {

namespace {

QString titleFrom(const QString& content) {
	const QString flat = content.simplified();
	if (flat.isEmpty()) return QString();
	if (flat.size() <= 48) return flat;
	return flat.left(47) + QStringLiteral("…");
}

} // namespace

ChatSession::ChatSession(const QString& id, QObject* parent)
	: QObject(parent), m_id(id) {
	m_model = new ChatMessageModel(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();
}

void ChatSession::setUpdatedAt(qint64 value) {
	if (m_updatedAt == value) return;
	m_updatedAt = value;
	Q_EMIT updatedAtChanged();
}

void ChatSession::setPinned(bool value) {
	if (m_pinned == value) return;
	m_pinned = value;
	Q_EMIT pinnedChanged();
}

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

void ChatSession::setMeta(
	const QString& title, qint64 createdAt, qint64 updatedAt, int messageCount) {
	m_title = title;
	m_createdAt = createdAt;
	m_updatedAt = updatedAt;
	m_messageCount = messageCount;
}

void ChatSession::setLastTokenCount(int value) {
	m_lastTokenCount = value;
}

LlmClient* ChatSession::client() const {
	if (auto* store = qobject_cast<ChatStore*>(parent()))
		return store->llmClient();
	return nullptr;
}

void ChatSession::ensureLoaded() {
	if (m_loadRequested) return;
	m_loadRequested = true;
	if (auto* store = qobject_cast<ChatStore*>(parent()))
		store->loadMessagesInto(this);
}

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

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

bool ChatSession::takeClearPending() {
	const bool pending = m_clearPending;
	m_clearPending = false;
	return pending;
}

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

void ChatSession::persist() {
	if (auto* store = qobject_cast<ChatStore*>(parent())) 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) {
	auto* generation = target->activeGeneration();
	auto* clientObject = client();
	if (generation && clientObject) {
		if (isLoaded()) {
			clientObject->startGeneration(this, generation);
		} else {
			// The store load is still in flight; the request needs the
			// full history, so start once it lands.
			connect(
				this, &ChatSession::loaded, clientObject,
				[this, generation, clientObject]() {
					clientObject->startGeneration(this, generation);
				});
		}
	}
	persist();
}

void ChatSession::sendMessage(const QString& text) {
	const QString trimmed = text.trimmed();
	if (trimmed.isEmpty()) return;
	if (auto* clientObject = client()) {
		if (clientObject->busy()) return;
	}

	const qint64 now = QDateTime::currentMSecsSinceEpoch();
	appendNewest(ChatMessage::Role::User, trimmed, now);
	while (m_model->rowCount() > 200)
		m_model->removeMessage(m_model->at(m_model->rowCount() - 1));
	if (m_title.isEmpty()) {
		m_title = titleFrom(trimmed);
		Q_EMIT titleChanged();
		qInfo() << "ChatSession:" << m_id << "new conversation,"
				<< "fallback title" << m_title
				<< "- requesting generated title and icon";
		if (auto* clientObject = client()) {
			clientObject->requestTitle(this, trimmed);
			clientObject->requestIcon(this, trimmed);
		}
	}
	auto* assistant =
		appendNewest(ChatMessage::Role::Assistant, QString(), now);
	startGeneration(assistant);
}

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

	// Drop everything newer than the target, then regenerate from the
	// context ending at the user message before it.
	m_model->removeRange(0, row - 1);
	if (m_model->rowCount() < 2) return;

	const qint64 now = QDateTime::currentMSecsSinceEpoch();
	target->appendGeneration(now);
	startGeneration(target);
}

void ChatSession::continueFrom(ChatMessage* message) {
	if (!message) return;
	if (auto* clientObject = client()) {
		if (clientObject->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 {
		auto* assistant =
			appendNewest(ChatMessage::Role::Assistant, QString(), now);
		startGeneration(assistant);
	}
}

void ChatSession::clear() {
	if (auto* clientObject = client()) {
		if (clientObject->busy() && clientObject->streamingSession() == this) {
			clientObject->clearOnFinish(this);
			clientObject->stop();
			return;
		}
	}
	if (!isLoaded()) {
		// The load lands shortly; drop everything once it does.
		m_clearPending = true;
		return;
	}
	m_model->clear();
	persist();
}

} // namespace ZShell::llm