z-bar-qt/Plugins/ZShell/Llm/llmclient.hpp

#pragma once

#include "segment.hpp"
#include "tool.hpp"

#include <QByteArray>
#include <QJsonArray>
#include <QNetworkAccessManager>
#include <QObject>
#include <QPointer>
#include <QString>
#include <QStringList>

#include <functional>

class QJsonObject;
class QNetworkReply;

namespace ZShell::llm {

class ChatGeneration;
class ChatSession;
class LlmSegment;
class LlmTool;

// The only component that talks to the LLM server: owns the network
// manager, the in-flight streaming state, the SSE parsing and the
// tool-calling loop.
class LlmClient : public QObject {
	Q_OBJECT

	public:
	explicit LlmClient(QObject* parent = nullptr);
	~LlmClient() override;

	[[nodiscard]] QString endpoint() const { return m_endpoint; }
	[[nodiscard]] QString model() const { return m_model; }
	[[nodiscard]] double temperature() const { return m_temperature; }
	[[nodiscard]] QStringList availableModels() const {
		return m_availableModels;
	}
	[[nodiscard]] int contextSize() const { return m_contextSize; }
	[[nodiscard]] ToolRegistry* tools() const { return m_tools; }
	void setEndpoint(const QString& value);
	void setModel(const QString& value);
	void setTemperature(double value);
	void setContextSize(int size);
	void probeContextSize();

	[[nodiscard]] bool busy() const { return m_busy; }
	[[nodiscard]] bool toolsEnabled() const { return m_tools->enabled(); }
	void setToolsEnabled(bool value) { m_tools->setEnabled(value); }
	[[nodiscard]] QString streamingChatId() const { return m_streamingChatId; }
	[[nodiscard]] ChatSession* streamingSession() const { return m_active; }

	// Streams a new assistant reply into `target` (the active generation
	// of a message of `session`). The context sent to the model is every
	// message of the session that is older than the target message.
	// Tool calls made by the model are executed and fed back
	// transparently until the model produces its final answer.
	void startGeneration(ChatSession* session, ChatGeneration* target);
	void stop();
	void endStream();
	// Clears the session's conversation once the current stream ends.
	void clearOnFinish(ChatSession* session);
	// A session is about to be destroyed; drop any state pointing at it.
	void sessionRemoved(ChatSession* session);

	void refreshModels();
	void requestTitle(ChatSession* session, const QString& userText);
	void requestIcon(ChatSession* session, const QString& userText);

	Q_SIGNALS:
	void busyChanged();
	void endpointChanged();
	void modelChanged();
	void availableModelsChanged();
	void contextSizeChanged();
	void toolsEnabledChanged();
	void streamingChatIdChanged();
	void errorOccurred(const QString& message);
	void titleSuggested(ZShell::llm::ChatSession* session, const QString& title);
	void iconSuggested(ZShell::llm::ChatSession* session, const QString& icon);

	private:
	struct ToolCallBuilder {
		QString id;
		QString name;
		QString arguments;
		QPointer<LlmSegment> segment;
		bool seen = false;
	};

	// Sends one streaming round: context + transcript so far.
	void sendRound();
	// The session context for a round, oldest first, ending just before
	// `stopBeforeRow` (the message of the generation being streamed).
	QJsonArray buildContextMessages(
		ChatSession* session, int stopBeforeRow) const;
	void applyToolCallDelta(const QJsonObject& call);
	// One round's stream ended; either ends the turn or executes the
	// requested tool calls and sends the next round. Runs at most once
	// per round ([DONE] and the reply's finished signal both reach it).
	void roundFinished();
	// Dispatches every call of the round; tools run concurrently.
	void executeAllCalls();
	// All results in: appends the tool messages (in call order) and
	// sends the next round.
	void flushCallResults();
	// Ends the current turn gracefully and persists the session.
	void finishTurn();
	void fail(const QString& message);
	void handleLine(const QByteArray& line);
	void drainBuffer();
	void updateTokenUsage(const QJsonObject& data);
	void setBusy(bool value);
	void setStreamingChatId(const QString& id);
	void shortRequest(
		const QString& tag,
		const QString& systemPrompt,
		const QString& userText,
		std::function<void(QString result)> onResult);
	static QString completionsPath(const QString& endpoint, const QString& subpath);
	static QString serverErrorMessage(
		const QByteArray& body, const QString& fallback);

	QNetworkAccessManager m_manager;
	ToolRegistry* m_tools = nullptr;
	QNetworkReply* m_reply = nullptr;
	QByteArray m_buffer;
	ChatSession* m_active = nullptr;
	ChatGeneration* m_streaming = nullptr;
	QPointer<ChatSession> m_pendingClear;
	bool m_busy = false;
	QString m_streamingChatId;
	QString m_endpoint;
	QString m_model;
	QStringList m_availableModels;
	double m_temperature = 0.7;
	int m_contextSize = 0;

	// State of the multi-round tool loop of the current turn.
	QJsonArray m_transcript;
	QList<ToolCallBuilder> m_callBuilders;
	struct ToolCallResult {
		QString content;
		bool success = false;
	};
	QList<ToolCallResult> m_callResults;
	QString m_finishReason;
	int m_round = 0;
	qsizetype m_contentMark = 0;
	qsizetype m_reasoningMark = 0;
	bool m_roundDone = false;
	bool m_toolPhase = false;
	int m_pendingCalls = 0;
	static constexpr int kMaxToolRounds = 12;
};

} // namespace ZShell::llm