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) {
	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) {
		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