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

#include "generation.hpp"

#include "llmclient.hpp"
#include "message.hpp"
#include "session.hpp"

#include <QDateTime>

namespace ZShell::llm {

ChatGeneration::ChatGeneration(
	ChatMessage* message, qint64 timestamp, QObject* parent)
	: QObject(parent), m_message(message), m_timestamp(timestamp) {
	m_timer.setParent(this);
	m_timer.setInterval(500);
	m_timer.setTimerType(Qt::CoarseTimer);
	connect(&m_timer, &QTimer::timeout, this, [this]() {
		bool runningText = false;
		bool runningTool = false;
		for (auto* segment : m_segments) {
			if (!segment->running()) continue;
			Q_EMIT segment->elapsedMsChanged();
			if (segment->type() == LlmSegment::Type::ToolCall)
				runningTool = true;
			else
				runningText = true;
		}
		if (runningText) {
			Q_EMIT reasoningElapsedMsChanged();
			Q_EMIT contentElapsedMsChanged();
		}
		if (runningTool) Q_EMIT toolsElapsedMsChanged();
		if (!runningText && !runningTool && !m_streaming) m_timer.stop();
	});
}

ChatMessage* ChatGeneration::message() const {
	return m_message;
}

QString ChatGeneration::content() const {
	QStringList parts;
	for (const auto* segment : m_segments) {
		if (segment->type() != LlmSegment::Type::Content) continue;
		const auto* text = qobject_cast<const LlmTextSegment*>(segment);
		if (!text || text->text().isEmpty()) continue;
		parts.append(text->text());
	}
	return parts.join(QStringLiteral("\n\n"));
}

QString ChatGeneration::reasoning() const {
	QStringList parts;
	for (const auto* segment : m_segments) {
		if (segment->type() != LlmSegment::Type::Reasoning) continue;
		const auto* text = qobject_cast<const LlmTextSegment*>(segment);
		if (!text || text->text().isEmpty()) continue;
		parts.append(text->text());
	}
	return parts.join(QStringLiteral("\n\n"));
}

bool ChatGeneration::reasoningActive() const {
	if (!m_streaming) return false;
	if (!content().isEmpty()) return false;
	return !hasRunningTool() && !hasPendingTool();
}

qint64 ChatGeneration::reasoningElapsedMs() const {
	qint64 total = 0;
	for (const auto* segment : m_segments)
		if (segment->type() == LlmSegment::Type::Reasoning)
			total += segment->elapsedMs();
	return total;
}

qint64 ChatGeneration::contentElapsedMs() const {
	qint64 total = 0;
	for (const auto* segment : m_segments)
		if (segment->type() == LlmSegment::Type::Content)
			total += segment->elapsedMs();
	return total;
}

qint64 ChatGeneration::toolsElapsedMs() const {
	qint64 total = 0;
	for (const auto* segment : m_segments)
		if (segment->type() == LlmSegment::Type::ToolCall)
			total += segment->elapsedMs();
	return total;
}

QList<LlmToolCall*> ChatGeneration::toolCalls() const {
	QList<LlmToolCall*> calls;
	for (auto* segment : m_segments)
		if (auto* call = qobject_cast<LlmToolCall*>(segment))
			calls.append(call);
	return calls;
}

int ChatGeneration::toolCallCount() const {
	return static_cast<int>(toolCalls().size());
}

bool ChatGeneration::hasRunningTool() const {
	for (const auto* call : toolCalls())
		if (call->running()) return true;
	return false;
}

bool ChatGeneration::hasPendingTool() const {
	for (const auto* call : toolCalls())
		if (call->status() == LlmToolCall::Status::Pending) return true;
	return false;
}

void ChatGeneration::setApprovalPending(bool value) {
	if (m_approvalPending == value) return;
	m_approvalPending = value;
	Q_EMIT toolApprovalPendingChanged();
	Q_EMIT pendingToolCallsChanged();
}

QVariantList ChatGeneration::pendingToolCalls() const {
	QVariantList out;
	if (!m_approvalPending) return out;
	for (const auto* call : toolCalls())
		if (call->status() == LlmToolCall::Status::Pending)
			out.append(QVariant::fromValue(call));
	return out;
}

void ChatGeneration::approveTools() {
	if (auto* llmClient = client()) llmClient->approveTools(this);
}

void ChatGeneration::denyTools() {
	if (auto* llmClient = client()) llmClient->denyTools(this);
}

void ChatGeneration::dismissError() {
	if (!m_hasError && m_errorMessage.isEmpty()) return;
	m_hasError = false;
	m_errorMessage.clear();
	Q_EMIT errorChanged();
	Q_EMIT errorMessageChanged();
}

LlmClient* ChatGeneration::client() const {
	if (m_message && m_message->session())
		return m_message->session()->client();
	return nullptr;
}

void ChatGeneration::updateReasoningActive() {
	const bool active = reasoningActive();
	if (m_reasoningActive == active) return;
	m_reasoningActive = active;
	Q_EMIT reasoningActiveChanged();
}

void ChatGeneration::setContent(const QString& value) {
	LlmTextSegment* first = nullptr;
	for (auto* segment : m_segments) {
		if (segment->type() != LlmSegment::Type::Content) continue;
		auto* text = qobject_cast<LlmTextSegment*>(segment);
		if (!text) continue;
		if (!first)
			first = text;
		else
			text->setText(QString());
	}
	if (!first) {
		if (value.isEmpty()) return;
		first = openTextSegment(LlmSegment::Type::Content);
	}
	first->setText(value);
}

void ChatGeneration::appendContent(const QString& piece) {
	if (piece.isEmpty()) return;
	for (auto* segment : m_segments) {
		if (segment->type() == LlmSegment::Type::Reasoning &&
			segment->running())
			segment->close();
	}
	openTextSegment(LlmSegment::Type::Content)->appendText(piece);
}

void ChatGeneration::appendReasoning(const QString& piece) {
	if (piece.isEmpty()) return;
	for (auto* segment : m_segments) {
		if (segment->type() == LlmSegment::Type::Content && segment->running())
			segment->close();
	}
	openTextSegment(LlmSegment::Type::Reasoning)->appendText(piece);
}

void ChatGeneration::setError(const QString& value) {
	const bool hasError = !value.isEmpty();
	if (m_hasError == hasError && m_errorMessage == value) return;
	m_hasError = hasError;
	m_errorMessage = value;
	Q_EMIT errorChanged();
	Q_EMIT errorMessageChanged();
}

void ChatGeneration::setStreaming(bool value) {
	if (m_streaming == value) return;
	m_streaming = value;
	Q_EMIT streamingChanged();
	if (value) {
		if (!m_timer.isActive()) m_timer.start();
	} else {
		closeOpenSegments();
		m_timer.stop();
	}
	Q_EMIT reasoningElapsedMsChanged();
	Q_EMIT contentElapsedMsChanged();
	Q_EMIT toolsElapsedMsChanged();
	updateReasoningActive();
}

LlmTextSegment* ChatGeneration::openTextSegment(LlmSegment::Type type) {
	for (auto* segment : m_segments) {
		if (segment->type() != type) continue;
		auto* text = qobject_cast<LlmTextSegment*>(segment);
		if (text && text->running()) return text;
	}
	auto* text =
		new LlmTextSegment(type, QDateTime::currentMSecsSinceEpoch(), this);
	text->begin();
	addSegment(text);
	return text;
}

LlmToolCall* ChatGeneration::beginToolCall(
	const QString& name, const QString& toolCallId) {
	closeOpenSegments();
	auto* call = new LlmToolCall(QDateTime::currentMSecsSinceEpoch(), this);
	call->setName(name);
	call->setToolCallId(toolCallId);
	call->setStatus(LlmToolCall::Status::Pending);
	call->begin();
	addSegment(call);
	Q_EMIT toolCallCountChanged();
	return call;
}

void ChatGeneration::addSegment(LlmSegment* segment) {
	if (!segment || m_segments.contains(segment)) return;
	segment->setParent(this);
	connect(segment, &LlmSegment::runningChanged, this, [this, segment]() {
		Q_EMIT toolCallCountChanged();
		if (segment->type() == LlmSegment::Type::ToolCall)
			Q_EMIT hasRunningToolChanged();
		updateReasoningActive();
	});
	if (auto* text = qobject_cast<LlmTextSegment*>(segment))
		connect(text, &LlmTextSegment::textChanged, this, [this, text]() {
			if (text->type() == LlmSegment::Type::Reasoning)
				Q_EMIT reasoningChanged();
			else
				Q_EMIT contentChanged();
			updateReasoningActive();
		});
	if (auto* call = qobject_cast<LlmToolCall*>(segment)) {
		connect(call, &LlmToolCall::statusChanged, this, [this, call]() {
			Q_EMIT hasRunningToolChanged();
			Q_EMIT pendingToolCallsChanged();
			if (call->status() == LlmToolCall::Status::Pending)
				Q_EMIT toolApprovalPendingChanged();
			updateReasoningActive();
		});
	}
	m_segments.append(segment);
	Q_EMIT segmentsChanged();
	Q_EMIT toolCallCountChanged();
}

void ChatGeneration::closeOpenSegments() {
	for (auto* segment : m_segments)
		if (segment->running()) segment->close();
	updateReasoningActive();
}

} // namespace ZShell::llm