chore: restructure llm plugin to be more coherent & fix various bugs in llm ui

Plugins/ZShell/Llm/turn.cpp

@@ -0,0 +1,391 @@
+#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