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

#include "turn.hpp"

#include "attachment.hpp"
#include "message.hpp"
#include "messagemodel.hpp"
#include "prompts.hpp"
#include "session.hpp"
#include "workspace.hpp"

#include <QFile>
#include <QSet>

namespace ZShell::llm {

namespace {

constexpr int kMinContextTokens = 512;
constexpr int kToolResultStubChars = 512;
constexpr double kPromptEstimateDiscount = 10.0 / 12.0;

qsizetype jsonValueChars(const QJsonValue& value) {
	switch (value.type()) {
	case QJsonValue::String:
		return value.toString().size();
	case QJsonValue::Array: {
		qsizetype total = 0;
		for (const QJsonValue& item : value.toArray())
			total += jsonValueChars(item);
		return total;
	}
	case QJsonValue::Object: {
		qsizetype total = 0;
		const QJsonObject obj = value.toObject();
		for (auto it = obj.constBegin(); it != obj.constEnd(); ++it)
			total += jsonValueChars(it.value());
		return total;
	}
	default:
		return 2;
	}
}

int estimateMessageTokens(const QJsonObject& message) {
	return int(jsonValueChars(message) / ToolRegistry::CharsPerToken) + 4;
}

int estimateArrayTokens(const QJsonArray& messages) {
	int total = 0;
	for (const QJsonValue& value : messages)
		total += estimateMessageTokens(value.toObject());
	return total;
}

int completionReserve(int contextSize) {
	return qBound(2048, contextSize / 8, 8192);
}

void shrinkUnitToBudget(QList<QJsonObject>& unit, int budgetTokens) {
	static const QStringList kTextKeys = {
		QStringLiteral("content"), QStringLiteral("reasoning_content")};
	for (int pass = 0; pass < 8; ++pass) {
		int tokens = 0;
		for (const QJsonObject& message : unit)
			tokens += estimateMessageTokens(message);
		if (tokens <= budgetTokens) return;
		qsizetype bestMessage = -1;
		QString bestKey;
		qsizetype bestLength = 0;
		for (qsizetype i = 0; i < unit.size(); ++i) {
			const QJsonObject message = unit.at(i);
			for (const QString& key : kTextKeys) {
				const QJsonValue value = message.value(key);
				if (value.isString() && value.toString().size() > bestLength) {
					bestMessage = i;
					bestKey = key;
					bestLength = value.toString().size();
				}
			}
		}
		if (bestMessage < 0) return;
		QString text = unit.at(bestMessage).value(bestKey).toString();
		text.truncate(qMax<qsizetype>(32, text.size() * 3 / 4));
		const QString marker = QStringLiteral("\n[…]");
		if (!text.endsWith(marker)) text += marker;
		unit[bestMessage][bestKey] = text;
	}
}

} // namespace

Turn::Turn(
	ChatSession* session,
	ChatGeneration* generation,
	const ToolRegistry& registry)
	: m_session(session), m_generation(generation), m_registry(registry) {}

int Turn::targetRow() const {
	if (!m_session) return -1;
	const ChatMessage* message = m_generation ? m_generation->message()
											  : nullptr;
	return message ? m_session->messagesModel()->rowOf(message) : -1;
}

int Turn::transcriptTokens() const {
	return estimateArrayTokens(m_transcript);
}

int Turn::promptBudgetTokens() const {
	const int contextSize = m_registry.contextSize();
	if (contextSize <= 0) return 0;
	return int(
		(contextSize - completionReserve(contextSize)) *
		kPromptEstimateDiscount);
}

int Turn::batchSize() const {
	int size = 0;
	for (const Call& call : m_calls)
		if (call.seen) ++size;
	return size;
}

int Turn::outputBudgetChars() const {
	const int batch = qMax(1, batchSize());
	const int promptLimit = promptBudgetTokens();
	if (promptLimit <= 0) return m_registry.outputBudgetChars(batch);
	const int available = promptLimit - transcriptTokens();
	const int tokens = qMax(available / batch, ToolRegistry::MinBudgetTokens);
	return tokens * ToolRegistry::CharsPerToken;
}

void Turn::resetRound() {
	roundDone = false;
	streamDone = false;
	toolPhase = false;
	pendingCalls = 0;
	finishReason.clear();
	m_calls.clear();
	m_results.clear();
}

void Turn::appendToTranscript(const QJsonObject& message) {
	m_transcript.append(message);
}

QJsonObject Turn::roundAssistantMessage() {
	QJsonObject assistant;
	assistant[QStringLiteral("role")] = QStringLiteral("assistant");

	const QString content = m_generation->content();
	const QString reasoning = m_generation->reasoning();
	if (content.size() > m_contentMark) {
		assistant[QStringLiteral("content")] = content.mid(m_contentMark);
		m_contentMark = content.size();
	}
	if (reasoning.size() > m_reasoningMark) {
		assistant[QStringLiteral("reasoning_content")] =
			reasoning.mid(m_reasoningMark);
		m_reasoningMark = reasoning.size();
	}

	QJsonArray calls;
	for (const Call& call : m_calls) {
		if (!call.seen) continue;
		QJsonObject function;
		function[QStringLiteral("name")] = call.name;
		function[QStringLiteral("arguments")] = call.arguments;
		QJsonObject entry;
		entry[QStringLiteral("id")] = call.id;
		entry[QStringLiteral("type")] = QStringLiteral("function");
		entry[QStringLiteral("function")] = function;
		calls.append(entry);
	}
	assistant[QStringLiteral("tool_calls")] = calls;
	return assistant;
}

QJsonArray Turn::toolMessages() const {
	QJsonArray messages;
	for (int i = 0; i < m_calls.size(); ++i) {
		const Call& call = m_calls.at(i);
		if (!call.seen) continue;
		QJsonObject toolMessage;
		toolMessage[QStringLiteral("role")] = QStringLiteral("tool");
		toolMessage[QStringLiteral("tool_call_id")] = call.id;
		toolMessage[QStringLiteral("content")] = m_results.at(i).text;
		messages.append(toolMessage);
	}
	return messages;
}

bool Turn::roundHasCalls() const {
	for (const Call& call : m_calls)
		if (call.seen) return true;
	return false;
}

void Turn::appendSandboxPreamble(QJsonArray& messages) const {
	if (!m_registry.sandboxEnabled()) return;
	const QString workspace = m_session->workspace()->path();
	if (workspace.isEmpty()) return;

	const auto* model = m_session->messagesModel();
	QStringList files;
	QSet<QString> seenPaths;
	for (int row = 0; row < model->rowCount(); ++row) {
		const auto* message = model->at(row);
		if (message->role() != ChatMessage::Role::User) continue;
		for (const auto* attachment : message->attachments()) {
			if (attachment->kind() != Attachment::Kind::File) continue;
			const QString relativePath =
				attachment->path().mid(workspace.size() + 1);
			if (seenPaths.contains(relativePath)) continue;
			seenPaths.insert(relativePath);
			files.append(QStringLiteral("- %1 (available as %2)")
							 .arg(attachment->name(), relativePath));
		}
	}
	if (files.isEmpty()) return;

	QJsonObject preamble;
	preamble[QStringLiteral("role")] = QStringLiteral("system");
	preamble[QStringLiteral("content")] =
		prompts::sandboxPreamble(workspace, files);
	messages.prepend(preamble);
}

QJsonArray Turn::contextMessages() const {
	const auto* model = m_session->messagesModel();
	const int stopBeforeRow = targetRow();

	QList<QList<QJsonObject>> units;
	for (int row = model->rowCount() - 1; row > stopBeforeRow; --row) {
		const auto* message = model->at(row);
		if (message->role() == ChatMessage::Role::User) {
			if (message->content().isEmpty() &&
				message->attachments().isEmpty())
				continue;

			QList<QJsonObject> unit;

			QJsonObject user;
			user[QStringLiteral("role")] = QStringLiteral("user");
			const QString text = message->content();
			const QList<Attachment*> attachments = message->attachments();
			if (attachments.isEmpty()) {
				user[QStringLiteral("content")] = text;
			} else {
				QJsonArray parts;
				if (!text.isEmpty()) {
					QJsonObject textPart;
					textPart[QStringLiteral("type")] = QStringLiteral("text");
					textPart[QStringLiteral("text")] = text;
					parts.append(textPart);
				}
				for (const auto* attachment : attachments) {
					if (attachment->kind() != Attachment::Kind::Image) continue;
					QFile file(attachment->path());
					if (!file.open(QIODevice::ReadOnly)) continue;
					const QByteArray base64 = file.readAll().toBase64();
					file.close();
					if (base64.isEmpty()) continue;
					QJsonObject imageUrl;
					imageUrl[QStringLiteral("url")] =
						QStringLiteral("data:%1;base64,%2")
							.arg(attachment->mimeType())
							.arg(QString::fromLatin1(base64));
					QJsonObject imagePart;
					imagePart[QStringLiteral("type")] =
						QLatin1String("image_url");
					imagePart[QStringLiteral("image_url")] = imageUrl;
					parts.append(imagePart);
				}
				if (parts.isEmpty())
					user[QStringLiteral("content")] = text;
				else
					user[QStringLiteral("content")] = parts;
			}
			unit.append(user);
			units.append(unit);
			continue;
		}

		const ChatGeneration* generation = message->activeGeneration();
		if (!generation) continue;
		const QList<LlmToolCall*> calls = generation->toolCalls();
		const bool hasAnyContent = !generation->content().isEmpty() ||
								   !generation->reasoning().isEmpty() ||
								   !calls.isEmpty();
		if (!hasAnyContent) continue;

		QJsonObject assistant;
		assistant[QStringLiteral("role")] = QStringLiteral("assistant");
		if (!calls.isEmpty() || !generation->content().isEmpty())
			assistant[QStringLiteral("content")] = generation->content();
		if (!generation->reasoning().isEmpty())
			assistant[QStringLiteral("reasoning_content")] =
				generation->reasoning();
		if (!calls.isEmpty()) {
			QJsonArray callArray;
			for (const auto* call : calls) {
				QJsonObject function;
				function[QStringLiteral("name")] = call->name();
				function[QStringLiteral("arguments")] = call->arguments();
				QJsonObject entry;
				entry[QStringLiteral("id")] = call->toolCallId();
				entry[QStringLiteral("type")] = QStringLiteral("function");
				entry[QStringLiteral("function")] = function;
				callArray.append(entry);
			}
			assistant[QStringLiteral("tool_calls")] = callArray;
		}

		QList<QJsonObject> unit;
		unit.append(assistant);
		for (const auto* call : calls) {
			QJsonObject toolMessage;
			toolMessage[QStringLiteral("role")] = QStringLiteral("tool");
			toolMessage[QStringLiteral("tool_call_id")] = call->toolCallId();
			toolMessage[QStringLiteral("content")] = call->result();
			unit.append(toolMessage);
		}
		units.append(unit);
	}

	if (m_registry.contextSize() > 0) {
		const int reserve = transcriptTokens();
		int budget = m_registry.contextSize() -
					 completionReserve(m_registry.contextSize()) - reserve;
		budget = qMax(budget, kMinContextTokens);

		int total = 0;
		for (const QList<QJsonObject>& unit : units)
			for (const QJsonObject& message : unit)
				total += estimateMessageTokens(message);
		while (units.size() > 1 && total > budget) {
			for (const QJsonObject& message : units.first())
				total -= estimateMessageTokens(message);
			units.removeFirst();
		}
		while (units.size() > 1 &&
			   units.first().first()[QStringLiteral("role")].toString() ==
				   QLatin1String("assistant")) {
			for (const QJsonObject& message : units.first())
				total -= estimateMessageTokens(message);
			units.removeFirst();
		}
		if (units.size() == 1) {
			QList<QJsonObject> single = units.first();
			shrinkUnitToBudget(single, budget);
			units[0] = single;
		}
	}

	QJsonArray messages;
	for (const QList<QJsonObject>& unit : units)
		for (const QJsonObject& message : unit)
			messages.append(message);

	for (const QJsonValue& value : m_transcript)
		messages.append(value);

	if (m_registry.contextSize() > 0) {
		const int promptLimit = promptBudgetTokens();
		while (estimateArrayTokens(messages) > promptLimit) {
			bool stubbed = false;
			for (qsizetype i = 0; i < messages.size(); ++i) {
				QJsonObject message = messages.at(i).toObject();
				if (message.value("role").toString() != QLatin1String("tool"))
					continue;
				const QString content = message.value("content").toString();
				if (content.size() <= kToolResultStubChars) continue;
				QString stub = content.left(kToolResultStubChars);
				stub += QStringLiteral(
							"\n[truncated to fit the context window; was "
							"%1 characters]")
							.arg(content.size());
				message["content"] = stub;
				messages[i] = message;
				stubbed = true;
				break;
			}
			if (!stubbed) break;
		}
	}

	appendSandboxPreamble(messages);
	return messages;
}

} // namespace ZShell::llm