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

#include "generation.hpp"

#include <QDateTime>

namespace ZShell::llm {

ChatGeneration::ChatGeneration(qint64 timestamp, QObject* parent)
: QObject(parent), m_timestamp(timestamp) {
	m_timer.setParent(this);
	m_timer.setInterval(500);
	m_timer.setTimerType(Qt::CoarseTimer);
	connect(&m_timer, &QTimer::timeout, this, [this]() {
		bool anyRunning = false;
		for (auto* segment : m_segments) {
			if (!segment->running())
				continue;
			anyRunning = true;
			segment->elapsedMsChanged();
		}
		if (anyRunning)
			Q_EMIT elapsedMsChanged();
		if (!anyRunning && !m_streaming)
			m_timer.stop();
	});
}

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

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

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;
}

int ChatGeneration::toolCallCount() const {
	int count = 0;
	for (const auto* segment : m_segments)
		if (segment->type() == LlmSegment::Type::ToolCall)
			++count;
	return count;
}

bool ChatGeneration::hasRunningTool() const {
	for (const auto* segment : m_segments)
		if (segment->type() == LlmSegment::Type::ToolCall &&
		    segment->running())
			return true;
	return false;
}

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

void ChatGeneration::setContent(const QString& value) {
	// Replaces the entire answer: the first content segment takes the
	// new text, any later content bursts are cleared.
	LlmSegment* first = nullptr;
	for (auto* segment : m_segments) {
		if (segment->type() != LlmSegment::Type::Content)
			continue;
		if (!first)
			first = segment;
		else
			segment->setText(QString());
	}
	if (!first) {
		// Do not materialize an empty content segment (e.g. the assistant
		// placeholder created before the stream starts).
		if (value.isEmpty())
			return;
		first = new LlmSegment(
			LlmSegment::Type::Content,
			QDateTime::currentMSecsSinceEpoch(),
			this);
		addSegment(first);
	}
	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();
	}
	openContentSegment()->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();
	}
	openReasoningSegment()->appendText(piece);
}

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 elapsedMsChanged();
	updateReasoningActive();
}

LlmSegment* ChatGeneration::openContentSegment() {
	for (auto* segment : m_segments)
		if (segment->type() == LlmSegment::Type::Content &&
		    segment->running())
			return segment;
	auto* segment = new LlmSegment(
		LlmSegment::Type::Content,
		QDateTime::currentMSecsSinceEpoch(),
		this);
	segment->begin();
	addSegment(segment);
	return segment;
}

LlmSegment* ChatGeneration::openReasoningSegment() {
	for (auto* segment : m_segments)
		if (segment->type() == LlmSegment::Type::Reasoning &&
		    segment->running())
			return segment;
	auto* segment = new LlmSegment(
		LlmSegment::Type::Reasoning,
		QDateTime::currentMSecsSinceEpoch(),
		this);
	segment->begin();
	addSegment(segment);
	return segment;
}

LlmSegment* ChatGeneration::beginToolCall(
	const QString& name, const QString& toolCallId) {
	closeOpenSegments();
	auto* segment = new LlmSegment(
		LlmSegment::Type::ToolCall,
		QDateTime::currentMSecsSinceEpoch(),
		this);
	segment->setName(name);
	segment->setToolCallId(toolCallId);
	segment->setStatus(LlmSegment::Status::Running);
	segment->begin();
	addSegment(segment);
	Q_EMIT toolStateChanged();
	return segment;
}

void ChatGeneration::addSegment(LlmSegment* segment) {
	if (!segment || m_segments.contains(segment))
		return;
	segment->setParent(this);
	connect(
		segment, &LlmSegment::textChanged, this, [this, segment]() {
			if (segment->type() == LlmSegment::Type::Reasoning)
				Q_EMIT reasoningChanged();
			else if (segment->type() == LlmSegment::Type::Content)
				Q_EMIT contentChanged();
			updateReasoningActive();
		});
	connect(
		segment, &LlmSegment::statusChanged, this, [this]() {
			Q_EMIT toolStateChanged();
		});
	connect(
		segment, &LlmSegment::resultChanged, this, [this]() {
			Q_EMIT toolStateChanged();
		});
	connect(
		segment, &LlmSegment::runningChanged, this, [this]() {
			Q_EMIT elapsedMsChanged();
			Q_EMIT toolStateChanged();
			updateReasoningActive();
		});
	m_segments.append(segment);
	Q_EMIT segmentsChanged();
}

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

} // namespace ZShell::llm