kateai/kateai.cpp

#include "kateai.h"

#include <QShortcut>
#include <KConfigGroup>
#include <KPluginFactory>
#include <KTextEditor/Document>
#include <KTextEditor/View>
#include <qdebug.h>
#include <qhash.h>
#include <qjsonobject.h>
#include <qmessagebox.h>
#include <qnamespace.h>
#include <QString>
#include <KActionCollection>
#include <KXMLGUIFactory>
#include <QJsonObject>
#include <QJsonDocument>
#include <QRandomGenerator>
#include <QTextCodec>
#include <limits>
#include <QMessageBox>
#include <QInputDialog>
#include <KLocalizedString>
#include <KSharedConfig>

#include "backend.h"
#include "exllama.h"
#include "kateaiconfigpage.h"

K_PLUGIN_FACTORY_WITH_JSON(KateAiPluginFactory, "kateai.json", registerPlugin<KateAiPlugin>();)

KateAiPlugin::KateAiPlugin(QObject *parent, const QList<QVariant> &)
	: KTextEditor::Plugin(parent), m_serverUrl(QStringLiteral("ws://localhost:8642"))
{
	m_ai = new ExLlama();
	KateAiPluginView::setAi(m_ai);
	readConfig();
}

KateAiPlugin::~KateAiPlugin() = default;

void KateAiPlugin::reconnect()
{
	m_ai->open(m_serverUrl);
}

QObject *KateAiPlugin::createView(KTextEditor::MainWindow *mainWindow)
{
	auto view = new KateAiPluginView(this, mainWindow);
	connect(view, &KateAiPluginView::reconnect, this, &KateAiPlugin::reconnect);
	return view;
}

KTextEditor::ConfigPage *KateAiPlugin::configPage(int number, QWidget *parent)
{
	if (number != 0)
		return nullptr;

	return new KateAiConfigPage(parent, this);
}

void KateAiPlugin::readConfig()
{
	KConfigGroup config(KSharedConfig::openConfig(), "Ai");
	m_serverUrl = QUrl(config.readEntry("Url", "ws://localhost:8642"));
	reconnect();

	KateAiPluginView::setInstruct(config.readEntry("Instruct", false));
	KateAiPluginView::setSystemPrompt(config.readEntry("SystemPrompt", "You are an intelligent programming assistant."));
	KateAiPluginView::setMaxContext(config.readEntry("Context", 1024));
}

KateAiPluginView::KateAiPluginView(KateAiPlugin *plugin, KTextEditor::MainWindow *mainwindow)
	: QObject(plugin)
	, m_mainWindow(mainwindow)
{
	KXMLGUIClient::setComponentName(QStringLiteral("kateaiplugin"), QStringLiteral("Kate Ai"));
	setXMLFile(QStringLiteral("ui.rc"));
	QAction *generateAction = actionCollection()->addAction(QStringLiteral("ai_generate"));
	generateAction->setText(QStringLiteral("Generate text using AI"));
	actionCollection()->setDefaultShortcut(generateAction, QKeySequence(Qt::CTRL | Qt::ALT | Qt::Key_A));
	m_mainWindow->guiFactory()->addClient(this);

	connect(generateAction, &QAction::triggered, this, &KateAiPluginView::generate);
	connect(m_ai, &AiBackend::gotResponse, this, &KateAiPluginView::gotResponse);
}

QPointer<KTextEditor::Document> KateAiPluginView::activeDocument() const
{
	KTextEditor::View *view = m_mainWindow->activeView();
	if(view && view->document())
		return view->document();
	return nullptr;
}

KTextEditor::Cursor KateAiPluginView::getCurrentCursor() const
{
	KTextEditor::View *view = m_mainWindow->activeView();
	if(view)
		return view->cursorPosition();
	return KTextEditor::Cursor();
}

void KateAiPluginView::gotResponse(AiBackend::Response response)
{
	QHash<uint32_t, Request>::iterator it = m_requests.find(response.getId());

	if(it != m_requests.end())
	{
		it.value().document->insertText(it.value().cursor, response.getText());
		if(response.isFinished())
			m_requests.erase(it);
		else
			it.value().cursor = getCurrentCursor();
	}
}

QStringList KateAiPluginView::getIncludePaths(const QString& text)
{
	QStringList lines = text.split(U'\n');
	QStringList paths;

	for(const QString& line : lines)
	{
		if(line.trimmed().startsWith(QStringLiteral("#include")))
		{
			int start = line.indexOf(U'<');
			int end;
			if(start != -1)
			{
				end = line.indexOf(U'>', start+1);
				if(end == -1)
					continue;
			}
			else
			{
				start = line.indexOf(U'"');
				if(start == -1)
					continue;

				end = line.indexOf(U'"', start+1);
				if(end == -1)
					continue;
			}
			paths.push_back(line.mid(start+1, (end-(start+1))).trimmed());
		}
	}
	return paths;
}

QString KateAiPluginView::assembleContext(QPointer<KTextEditor::Document> document,
                                          const KTextEditor::Cursor& cursor, const QString& instruction)
{
	QString mime = document->mimeType();
	QString context;
	QString documentText = document->text(KTextEditor::Range(KTextEditor::Cursor(0, 0), cursor));

	if(m_useInstruct)
	{
		context.append(QStringLiteral("### System Prompt\n"));
		context.append(m_systemPrompt);
		context.push_back(u'\n');
		context.append(QStringLiteral("### User Message\n"));
		context.append(instruction);
		context.push_back(QStringLiteral("\n### Assistant\n"));
	}

	if(mime ==  QStringLiteral("text/x-c++src") || mime ==  QStringLiteral("text/x-csrc"))
	{
		QFileInfo documentFileInfo(document->url().path());
		QString directory = documentFileInfo.absolutePath();
		QStringList paths = getIncludePaths(documentText);
		qDebug()<<__func__<<"Directory:"<<directory<<"Paths:"<<paths;

		for(QString& path : paths)
		{
			path = directory + QDir::separator() + path;
			qDebug()<<path;
			QFile file(path);
			if(!file.isOpen())
				continue;
			QByteArray fileData = file.readAll();
			QString fileText = QString::fromUtf8(fileData);
			context.append(fileText);
		}
	}

	context.append(documentText);

	return context;
}

void KateAiPluginView::generate()
{
	qDebug()<<activeDocument()->mimeType();
	if(m_ai->ready())
	{
		KTextEditor::Cursor cursor = getCurrentCursor();
		QPointer<KTextEditor::Document> document = activeDocument();
		QString instruction;

		if(m_useInstruct)
		{
			bool ok;
			instruction = QInputDialog::getMultiLineText(m_mainWindow->activeView(), i18n("Input instruction"),
			                                             i18n("Instruction"), m_lastInstruct, &ok);
			if(!ok)
				return;
		}

		QString text = assembleContext(document, cursor, instruction);

		AiBackend::Request request(text, AiBackend::Request::INFERENCE);
		bool ret = m_ai->generate(request);
		if(!ret)
			QMessageBox::warning(m_mainWindow->activeView(), i18n("Failure"),
								 i18n("The Ai backend was unable to process this request"));
		else
			m_requests.insert(request.getId(), {cursor, document});
	}
	else
	{
		QMessageBox box(m_mainWindow->activeView());
		box.setText(i18n("The AI server is not connected."));
		box.setInformativeText(i18n("would you like to try and reconnect?"));
		box.setStandardButtons(QMessageBox::Yes | QMessageBox::No);
		box.setDefaultButton(QMessageBox::Yes);
		int ret = box.exec();
		if(ret == QMessageBox::Yes)
			reconnect();
	}
}

KateAiPluginView::~KateAiPluginView()
{
	m_mainWindow->guiFactory()->removeClient(this);
}

void KateAiPluginView::setInstruct(bool instruct)
{
	m_useInstruct = instruct;
}

void KateAiPluginView::setSystemPrompt(const QString& in)
{
	m_systemPrompt = in;
}

void KateAiPluginView::setAi(QPointer<AiBackend> ai)
{
	m_ai = ai;
}

void KateAiPluginView::setMaxContext(int maxContext)
{
	m_maxContext = maxContext;
}

#include "kateai.moc"
#include "moc_kateai.cpp"