feat: [Ai] Add basic RAG capabilities, optimize Codegeex

Log: as title

src/plugins/codegeex/codegeex.cpp

@@ -59,11 +59,15 @@ bool CodeGeex::start()
  connect(&dpf::Listener::instance(), &dpf::Listener::pluginsStarted, [=] {
  QTimer::singleShot(5000, windowService, [=] {
  bool ret = CodeGeeXManager::instance()->isLoggedIn();
- if (ret)
- return;
+ if (!ret) {
+ QStringList actions { "codegeex_login_default", CodeGeex::tr("Login") };
+ windowService->notify(0, "CodeGeex", CodeGeex::tr("Please login to use CodeGeeX."), actions);
+ }
 
- QStringList actions { "codegeex_login_default", CodeGeex::tr("Login") };
- windowService->notify(0, "CodeGeex", CodeGeex::tr("Please login to use CodeGeeX."), actions);
+ if (!CodeGeeXManager::instance()->checkCondaInstalled()) {
+ QStringList actions { "ai_rag_install", CodeGeex::tr("Install") };
+ windowService->notify(0, "AI", CodeGeex::tr("Install a Python Conda virtual environment for using the RAG feature. Without RAG, there may be abnormalities in the @codebase and some AI functionalities."), actions);
+ }
  });
  });
 

src/plugins/codegeex/codegeex/askapi.cpp

@@ -4,10 +4,13 @@
 #include "askapi.h"
 #include "codegeexmanager.h"
 #include "src/common/supportfile/language.h"
+#include "services/project/projectservice.h"
+#include "services/window/windowservice.h"
 
 #include <QNetworkAccessManager>
 #include <QNetworkReply>
 #include <QDesktopServices>
+#include <QtConcurrent>
 
 #include <QJsonObject>
 #include <QJsonArray>
@@ -100,12 +103,22 @@ void AskApi::postSSEChat(const QString &url,
  const QString &talkId)
 {
  QJsonArray jsonArray = convertHistoryToJSONArray(history);
- QByteArray body = assembleSSEChatBody(prompt, machineId, jsonArray, talkId);
- QNetworkReply *reply = postMessage(url, token, body);
- connect(this, &AskApi::stopReceive, reply, [reply]() {
- reply->close();
+ auto impl = CodeGeeXManager::instance();
+ if (impl->isReferenceCodebase() && impl->checkCondaInstalled()) {
+ QStringList actions { "ai_rag_install", tr("Install") };
+ dpfservice::WindowService *windowService = dpfGetService(dpfservice::WindowService);
+ windowService->notify(0, "AI", tr("The RAG feature is not available. Please install it (the installation process may take several minutes)."), actions);
+ }
+ QtConcurrent::run([prompt, machineId, jsonArray, talkId, url, token, this]() {
+ QByteArray body = assembleSSEChatBody(prompt, machineId, jsonArray, talkId);
+ QMetaObject::invokeMethod(this, [url, token, body, this]() {
+ QNetworkReply *reply = postMessage(url, token, body);
+ connect(this, &AskApi::stopReceive, reply, [reply]() {
+ reply->close();
+ });
+ processResponse(reply);
+ }, Qt::QueuedConnection);
  });
- processResponse(reply);
 }
 
 void AskApi::postNewSession(const QString &url,
@@ -213,7 +226,6 @@ QNetworkReply *AskApi::postMessage(const QString &url, const QString &token, con
  QNetworkRequest request(url);
  request.setHeader(QNetworkRequest::ContentTypeHeader, "application/json");
  request.setRawHeader("code-token", token.toUtf8());
-
  return manager->post(request, body);
 }
 
@@ -322,9 +334,32 @@ QByteArray AskApi::assembleSSEChatBody(const QString &prompt,
 
  jsonObject.insert("prompt", prompt);
  jsonObject.insert("machineId", machineId);
- jsonObject.insert("client", "deepin-unioncode");
+ //jsonObject.insert("client", "deepin-unioncode");
  jsonObject.insert("history", history);
  jsonObject.insert("locale", locale);
+ jsonObject.insert("model", model);
+
+ if (CodeGeeXManager::instance()->isReferenceCodebase()) {
+ using dpfservice::ProjectService;
+ ProjectService *prjService = dpfGetService(ProjectService);
+ auto currentProjectPath = prjService->getActiveProjectInfo().workspaceFolder();
+ QJsonObject result = CodeGeeXManager::instance()->query(currentProjectPath, prompt, 20);
+ QJsonArray chunks = result["Chunks"].toArray();
+ if (!chunks.isEmpty()) {
+ CodeGeeXManager::instance()->cleanHistoryMessage(); // incase history is too big
+ jsonObject["history"] = QJsonArray();
+ QString context;
+ context += prompt;
+ context += "\n 参考下面这些代码片段，回答上面的问题。不要参考其他的代码和上下文，数据不够充分的情况下提示用户\n";
+ for (auto chunk : chunks) {
+ context += chunk.toObject()["fileName"].toString();
+ context += '\n';
+ context += chunk.toObject()["content"].toString();
+ context += "\n\n";
+ }
+ jsonObject["prompt"] = context;
+ }
+ }
 
  if (!CodeGeeXManager::instance()->getReferenceFiles().isEmpty()) {
  auto fileDatas = parseFile(CodeGeeXManager::instance()->getReferenceFiles());
@@ -334,8 +369,6 @@ QByteArray AskApi::assembleSSEChatBody(const QString &prompt,
  jsonObject.insert("files", files);
  } else if (CodeGeeXManager::instance()->isConnectToNetWork())
  jsonObject.insert("command", "online_search");
- else
- jsonObject.insert("model", model);
 
  if (!talkId.isEmpty())
  jsonObject.insert("talkId", talkId);

src/plugins/codegeex/codegeex/copilotapi.cpp

@@ -4,13 +4,15 @@
 #include "copilotapi.h"
 #include "src/common/supportfile/language.h"
 #include "src/services/editor/editorservice.h"
+#include "src/services/project/projectservice.h"
 
 #include <QNetworkAccessManager>
 #include <QNetworkRequest>
 #include <QNetworkReply>
 
 #include <QJsonObject>
 #include <QJsonArray>
+#include <QJsonObject>
 #include <QJsonDocument>
 
 #include <QString>
@@ -26,13 +28,17 @@ CopilotApi::CopilotApi(QObject *parent)
 
 void CopilotApi::postGenerate(const QString &url, const QString &prefix, const QString &suffix, GenerateType type)
 {
- if (completionReply)
- completionReply->close();
- QByteArray body = assembleGenerateBody(prefix, suffix, type);
- QNetworkReply *reply = postMessage(url, CodeGeeXManager::instance()->getSessionId(), body);
- completionReply = reply;
- reply->setProperty("responseType", CopilotApi::inline_completions);
- processResponse(reply);
+ QtConcurrent::run([prefix, suffix, type, url, this](){
+ QByteArray body = assembleGenerateBody(prefix, suffix, type);
+ QMetaObject::invokeMethod(this, [url, body, this](){
+ if (completionReply)
+ completionReply->close();
+ QNetworkReply *reply = postMessage(url, CodeGeeXManager::instance()->getSessionId(), body);
+ completionReply = reply;
+ reply->setProperty("responseType", CopilotApi::inline_completions);
+ processResponse(reply);
+ });
+ });
 }
 
 void CopilotApi::postTranslate(const QString &url,
@@ -108,12 +114,27 @@ QByteArray CopilotApi::assembleGenerateBody(const QString &prefix, const QString
  activeDocument.insert("suffix", suffix);
  activeDocument.insert("lang", file.second);
 
- QJsonObject contextItem;
- contextItem.insert("kind", "active_document");
- contextItem.insert("active_document", activeDocument);
+ QJsonObject activeContextItem;
+ activeContextItem.insert("kind", "active_document");
+ activeContextItem.insert("active_document", activeDocument);
 
+ ProjectService *prjSrv = dpfGetService(ProjectService);
  QJsonArray context;
- context.append(contextItem);
+ context.append(activeContextItem);
+ QJsonObject queryResults = CodeGeeXManager::instance()->query(prjSrv->getActiveProjectInfo().workspaceFolder(), prefix, 5);
+ QJsonArray chunks = queryResults["Chunks"].toArray();
+
+ for (auto chunk : chunks) {
+ QJsonObject document;
+ document.insert("path", chunk.toObject()["fileName"].toString());
+ document.insert("text", chunk.toObject()["content"].toString());
+ document.insert("lang", file.second);
+
+ QJsonObject contextItem;
+ contextItem.insert("kind", "document");
+ contextItem.insert("document", document);
+ context.append(contextItem);
+ }
 
  QJsonObject json;
  json.insert("context", context);

src/plugins/codegeex/codegeexmanager.cpp

@@ -5,6 +5,9 @@
 #include "codegeexmanager.h"
 #include "copilot.h"
 #include "common/util/custompaths.h"
+#include "services/window/windowservice.h"
+#include "services/window/windowelement.h"
+#include "services/terminal/terminalservice.h"
 
 #include <QDebug>
 #include <QFile>
@@ -14,6 +17,7 @@
 #include <QUuid>
 #include <QTimer>
 #include <QDateTime>
+#include <QtConcurrent>
 
 static const char *kUrlSSEChat = "https://codegeex.cn/prod/code/chatCodeSseV3/chat";
 //static const char *kUrlSSEChat = "https://codegeex.cn/prod/code/chatGlmSse/chat";
@@ -23,6 +27,7 @@ static const char *kUrlQuerySession = "https://codegeex.cn/prod/code/chatGlmTalk
 static const char *kUrlQueryMessage = "https://codegeex.cn/prod/code/chatGmlMsg/selectList";
 
 using namespace CodeGeeX;
+using dpfservice::WindowService;
 
 CodeGeeXManager *CodeGeeXManager::instance()
 {
@@ -49,6 +54,19 @@ bool CodeGeeXManager::isLoggedIn() const
  return isLogin;
 }
 
+bool CodeGeeXManager::checkCondaInstalled()
+{
+ QProcess process;
+ QStringList arguments;
+ arguments << "env" << "list";
+
+ process.start(condaRootPath() + "/miniforge/condabin/conda", arguments);
+ process.waitForFinished();
+
+ QString output = process.readAll();
+ return output.contains("deepin_unioncode_env");
+}
+
 void CodeGeeXManager::saveConfig(const QString &sessionId, const QString &userId)
 {
  QJsonObject config;
@@ -390,3 +408,85 @@ QString CodeGeeXManager::modifiedData(const QString &data)
 
  return retData;
 }
+
+QString CodeGeeXManager::condaRootPath() const
+{
+ return QDir::homePath() + "/.unioncode";
+}
+
+void CodeGeeXManager::installConda()
+{
+ QString scriptPath = CustomPaths::CustomPaths::global(CustomPaths::Scripts) + "/rag/install.sh";
+// QProcess process;
+// process.setProgram("bash");
+// process.setArguments(QStringList() << scriptPath << condaRootPath());
+// process.startDetached();
+
+ QMetaObject::invokeMethod(this, [=, scriptPath]() {
+ auto terminalServ = dpfGetService(dpfservice::TerminalService);
+ WindowService *windowService = dpfGetService(WindowService);
+ windowService->switchContextWidget(dpfservice::TERMINAL_TAB_TEXT);
+ terminalServ->executeCommand("install", "bash", QStringList() << scriptPath << condaRootPath(), condaRootPath(), QStringList());
+ });
+}
+
+void CodeGeeXManager::generateRag(const QString &projectPath)
+{
+ if (indexingProject.contains(projectPath))
+ return;
+ bool failed = false;
+ indexingProject.append(projectPath);
+ QProcess process;
+ QObject::connect(&process, &QProcess::readyReadStandardError, &process, [&, projectPath]() {
+ if (!failed) // only notify once
+ QMetaObject::invokeMethod(this, [&, projectPath](){
+ WindowService *windowService = dpfGetService(WindowService);
+ windowService->notify(2, "Ai", tr("The error occurred when performing rag on project %1.").arg(projectPath), QStringList{});
+ });
+ failed = true;
+ qInfo() << "Error:" << process.readAllStandardError() << "\n";
+ });
+ QObject::connect(&process, QOverload<int, QProcess::ExitStatus>::of(&QProcess::finished),
+ this, [&](int exitCode, QProcess::ExitStatus exitStatus) {
+ qInfo() << "Python script finished with exit code" << exitCode << "Exit!!!";
+ });
+ qInfo() << "start rag project:" << projectPath;
+ QString ragPath = CustomPaths::CustomPaths::global(CustomPaths::Scripts) + "/rag";
+ process.setWorkingDirectory(ragPath);
+ auto generatePyPath = ragPath + "/generate.py";
+ auto pythonPath = condaRootPath() + "/miniforge/envs/deepin_unioncode_env/bin/python";
+ if (!QFileInfo(pythonPath).exists())
+ return;
+ process.start(pythonPath, QStringList() << generatePyPath << projectPath);
+ process.waitForFinished(-1);
+ indexingProject.removeOne(projectPath);
+}
+
+/*
+ JsonObject:
+ Query: str
+ Chunks: Arr[fileName:str, content:str]
+ Instructions: obj{name:str, description:str, content:str}
+*/
+QJsonObject CodeGeeXManager::query(const QString &projectPath, const QString &query, int topItems)
+{
+ QProcess process;
+ QObject::connect(&process, &QProcess::readyReadStandardError, &process, [&]() {
+ qInfo() << "Error:" << process.readAllStandardError() << "\n";
+ });
+
+ // If modified, update the Python file accordingly.
+ auto pythonPath = condaRootPath() + "/miniforge/envs/deepin_unioncode_env/bin/python";
+ if (!QFileInfo(pythonPath).exists() || !QFileInfo(condaRootPath() +"/index.sqlite").exists())
+ return {};
+
+ QString ragPath = CustomPaths::CustomPaths::global(CustomPaths::Scripts) + "/rag";
+ process.setWorkingDirectory(ragPath);
+ auto queryPyPath = ragPath + "/query.py";
+ process.start(pythonPath, QStringList() << queryPyPath << projectPath << query << QString::number(topItems));
+ process.waitForFinished();
+ auto test = process.readAll();
+ QJsonDocument document = QJsonDocument::fromJson(test);
+ auto obj = document.object();
+ return obj;
+}

src/plugins/codegeex/codegeexmanager.h

@@ -51,6 +51,7 @@ class CodeGeeXManager : public QObject
 
  Q_INVOKABLE void login();
  bool isLoggedIn() const;
+ bool checkCondaInstalled();
 
  void saveConfig(const QString &sessionId, const QString &userId);
  void loadConfig();
@@ -81,7 +82,21 @@ class CodeGeeXManager : public QObject
  void connectToNetWork(bool connecting);
  bool isConnectToNetWork() { return isConnecting; }
  QStringList getReferenceFiles() { return referenceFiles; }
- void setRefereceFiles(const QStringList &files) { referenceFiles = files; }
+ void setReferenceCodebase(bool on) { referenceCodebase = on; }
+ bool isReferenceCodebase() { return referenceCodebase; }
+ void setReferenceFiles(const QStringList &files) { referenceFiles = files; }
+
+ // Rag
+ QString condaRootPath() const;
+ void installConda();
+ void generateRag(const QString &projectPath);
+ /*
+ JsonObject:
+ Query: str
+ Chunks: Arr[fileName:str, content:str]
+ Instructions: obj{name:str, description:str, content:str}
+ */
+ QJsonObject query(const QString &projectPath, const QString &query, int topItems);
 
 Q_SIGNALS:
  void loginSuccessed();
@@ -132,7 +147,9 @@ public Q_SLOTS:
  bool isLogin { false };
  bool isRunning { false };
  bool isConnecting { false };
- QStringList referenceFiles;
+ bool referenceCodebase { false };
+ QStringList indexingProject {};
+ QStringList referenceFiles {};
 };
 
 #endif // CODEGEEXMANAGER_H

src/plugins/codegeex/copilot.cpp

@@ -30,6 +30,8 @@ Copilot::Copilot(QObject *parent)
  if (!editorService) {
  qFatal("Editor service is null!");
  }
+ generateTimer = new QTimer(this);
+ generateTimer->setSingleShot(true);
 
  connect(&copilotApi, &CopilotApi::response, [this](CopilotApi::ResponseType responseType, const QString &response, const QString &dstLang) {
  switch (responseType) {
@@ -65,6 +67,7 @@ Copilot::Copilot(QObject *parent)
 
  connect(&copilotApi, &CopilotApi::responseByStream, this, &Copilot::response);
  connect(&copilotApi, &CopilotApi::messageSended, this, &Copilot::messageSended);
+ connect(generateTimer, &QTimer::timeout, this, &Copilot::generateCode);
 }
 
 QString Copilot::selectedText() const
@@ -164,9 +167,11 @@ void Copilot::setCurrentModel(CodeGeeX::languageModel model)
 
 void Copilot::handleTextChanged()
 {
- // start generate code.
+ editorService->setCompletion("", QIcon::fromTheme("codegeex_anwser_icon"), QKeySequence(Qt::CTRL | Qt::Key_T));
  QMetaObject::invokeMethod(this, [this]() {
- generateCode();
+ if (generateTimer->isActive())
+ generateTimer->stop();
+ generateTimer->start(500);
  });
 }
 

src/plugins/codegeex/copilot.h

@@ -62,6 +62,7 @@ public slots:
 
  CodeGeeX::CopilotApi copilotApi;
  dpfservice::EditorService *editorService = nullptr;
+ QTimer *generateTimer = nullptr;
  QStringList generateCache {};
  QString generatedCode {};
  QString extractSingleLine();