Workspace completion

apps/remix-ide/src/app/plugins/remixAIPlugin.tsx

@@ -6,6 +6,7 @@ import React, { useCallback } from 'react';
 import { ICompletions, IModel, RemoteInferencer, IRemoteModel, IParams, GenerationParams, CodeExplainAgent, SecurityAgent } from '@remix/remix-ai-core';
 import { CustomRemixApi } from '@remix-api'
 import { PluginViewWrapper } from '@remix-ui/helper'
+import { CodeCompletionAgent } from '@remix/remix-ai-core';
 const _paq = (window._paq = window._paq || [])
 
 type chatRequestBufferT<T> = {
@@ -41,6 +42,7 @@ export class RemixAIPlugin extends ViewPlugin {
   securityAgent: SecurityAgent
   useRemoteInferencer:boolean = false
   dispatch: any
+  completionAgent: CodeCompletionAgent
 
   constructor(inDesktop:boolean) {
     super(profile)
@@ -62,7 +64,7 @@ export class RemixAIPlugin extends ViewPlugin {
       this.useRemoteInferencer = true
       this.initialize()
     }
-
+    this.completionAgent = new CodeCompletionAgent(this)
     this.securityAgent = new SecurityAgent(this)
   }
 
@@ -107,10 +109,14 @@ export class RemixAIPlugin extends ViewPlugin {
   }
 
   async code_completion(prompt: string, promptAfter: string): Promise<any> {
+    if (this.completionAgent.indexer == null || this.completionAgent.indexer == undefined) await this.completionAgent.indexWorkspace()
+
+    const currentFile = await this.call('fileManager', 'getCurrentFile')
+    const contextfiles = await this.completionAgent.getContextFiles(prompt)
     if (this.isOnDesktop && !this.useRemoteInferencer) {
       return await this.call(this.remixDesktopPluginName, 'code_completion', prompt, promptAfter)
     } else {
-      return await this.remoteInferencer.code_completion(prompt, promptAfter)
+      return await this.remoteInferencer.code_completion(prompt, promptAfter, contextfiles, currentFile)
     }
   }
 
@@ -169,10 +175,15 @@ export class RemixAIPlugin extends ViewPlugin {
   }
 
   async code_insertion(msg_pfx: string, msg_sfx: string): Promise<any> {
+    if (this.completionAgent.indexer == null || this.completionAgent.indexer == undefined) await this.completionAgent.indexWorkspace()
+
+    const currentFile = await this.call('fileManager', 'getCurrentFile')
+    const contextfiles = await this.completionAgent.getContextFiles(msg_pfx)
+
     if (this.isOnDesktop && !this.useRemoteInferencer) {
       return await this.call(this.remixDesktopPluginName, 'code_insertion', msg_pfx, msg_sfx)
     } else {
-      return await this.remoteInferencer.code_insertion(msg_pfx, msg_sfx)
+      return await this.remoteInferencer.code_insertion( msg_pfx, msg_sfx, contextfiles, currentFile)
     }
   }
 
@@ -195,7 +206,7 @@ export class RemixAIPlugin extends ViewPlugin {
     else {
       console.log("chatRequestBuffer is not empty. First process the last request.", this.chatRequestBuffer)
     }
-    _paq.push(['trackEvent', 'ai', 'remixAI_chat', 'askFromTerminal'])
+    _paq.push(['trackEvent', 'ai', 'remixAI', 'remixAI_chat'])
   }
 
   async ProcessChatRequestBuffer(params:IParams=GenerationParams){

libs/remix-ai-core/src/agents/completionAgent.ts

@@ -1,23 +1,137 @@
-import * as fs from 'fs';
+import lunr from 'lunr';
 
-class CodeCompletionAgent {
-  private codebase: string[];
+interface Document {
+  id: number;
+  filename: string;
+  content: string;
+  identifier: number;
+}
+
+interface indexT{
+  isIndexed: boolean;
+  lastIndexedTime?: number;
+  reason?: string;
+}
+
+enum SupportedFileExtensions {
+  solidity = '.sol',
+  vyper = '.vy',
+  circom = '.circom',
+}
+
+export class CodeCompletionAgent {
+  props: any;
+  indexer: any;
+  Documents: Document[] = [];
+  INDEX_THRESHOLD = 0.1;
+  N_MATCHES = 1;
+  indexed: indexT = {
+    isIndexed: false,
+    lastIndexedTime: 0,
+    reason: 'Init',
+  };
 
-  constructor(codebasePath: string) {
-    // git or fs
-    this.codebase = this.loadCodebase(codebasePath);
+  constructor(props) {
+    this.props = props;
+    this.listenForChanges();
+    this.indexer =lunr(function () {
+      this.ref('id')
+      this.field('filename')
+      this.field('content')
+      this.field('Identifier');
+    });
+
+    setInterval(() => {
+      this.indexWorkspace()
+    }, 60000)
+  }
+
+  listenForChanges() {
+    this.props.on('fileManager', 'fileAdded', (path) => { this.indexed = { isIndexed: false, reason:"fileAdded" } });
+    this.props.on('fileManager', 'fileRemoved', (path) => { this.indexed = { isIndexed: false, reason:"fileRemoved" } });
+    this.props.on('filePanel', 'workspaceCreated', () => { this.indexed = { isIndexed: false, reason:"workspaceCreated" } });
+    this.props.on('filePanel', 'workspaceRenamed', () => { this.indexed = { isIndexed: false, reason:"workspaceRenamed" }});
+    this.props.on('filePanel', 'workspaceDeleted', () => { this.indexed = { isIndexed: false, reason:"workspaceDeleted" } });
   }
 
-  private loadCodebase(path: string): string[] {
-    const files = fs.readdirSync(path);
-    return files
-      .filter(file => file.endsWith('.ts'))
-      .flatMap(file => fs.readFileSync(`${path}/${file}`, 'utf-8').split('\n'));
+  async getDcocuments() {
+    try {
+      const documents: Document[] = [];
+      const jsonDirsContracts = await this.props.call('fileManager', 'copyFolderToJson', '/').then((res) => res.contracts);
+      let c = 0;
+      for (const file in jsonDirsContracts.children) {
+        if (!Object.values(SupportedFileExtensions).some(ext => file.endsWith(ext))) continue;
+        documents.push({
+          id: ++c,
+          filename: file,
+          content: jsonDirsContracts.children[file].content,
+          identifier: c - 1,
+        });
+      }
+      return documents;
+    } catch (error) {
+      return [];
+    }
   }
 
-  public getSuggestions(currentLine: string, numSuggestions: number = 3): string[] {
-    const suggestions: string[] = [];
-    // get `numSuggestions` from the llm
-    return suggestions;
+  indexWorkspace() {
+    this.getDcocuments().then((documents) => {
+      this.indexer =lunr(function () {
+        this.ref('id')
+        this.field('filename')
+        this.field('content')
+        this.field('Identifier');
+
+        documents.forEach(doc => {
+          this.add(doc);
+        });
+      });
+      this.Documents = documents;
+    });
+
+    this.indexed = { isIndexed: true, lastIndexedTime: Date.now(), reason: 'init Indexing' };
   }
+
+  async getContextFiles(prompt) {
+    try {
+      if (!this.indexed.isIndexed) {
+        await this.indexWorkspace();
+      }
+
+      const currentFile = await this.props.call('fileManager', 'getCurrentFile');
+      const content = prompt;
+      const searchResult = this.indexer.search(content)
+      const fcps = await this.processResults(searchResult, currentFile);
+      const resolvedFcps = await Promise.all(fcps);
+      return resolvedFcps;
+    } catch (error) {
+      return [];
+    }
+  }
+
+  async processResults(results: any, currentFile: string) {
+
+    // remove the current file name from the results list
+    const rmResults = await results.filter(result => {
+      return this.Documents.find(doc => doc.id === Number(result.ref)).filename !== currentFile;
+    });
+
+    // filter out the results which have the same extension as the current file.
+    // Do not mix and match file extensions as this will lead to incorrect completions
+    const extResults = await rmResults.filter(result => {
+      return this.Documents.find(doc => doc.id === Number(result.ref)).filename.split('.').pop() === currentFile.split('.').pop();
+    });
+
+    // filter out the results which have a score less than the INDEX_THRESHOLD
+    const topResults = await extResults.filter(result => result.score >= this.INDEX_THRESHOLD).slice(0, this.N_MATCHES);
+
+    // get the LATEST content of the top results in case the file has been modified and not indexed yet
+    const fileContentPairs = topResults.map(async result => {
+      const document = this.Documents.find(doc => doc.id === Number(result.ref));
+      const currentContent = await this.props.call('fileManager', 'readFile', document.filename);
+      return { file: document.filename, content: currentContent };
+    });
+    return fileContentPairs;
+  }
+
 }

libs/remix-ai-core/src/index.ts

@@ -22,4 +22,5 @@ export {
 export * from './types/types'
 export * from './helpers/streamHandler'
 export * from './agents/codeExplainAgent'
-export * from './agents/securityAgent'
+export * from './agents/completionAgent'
+export * from './agents/securityAgent'

libs/remix-ai-core/src/inferencers/remote/remoteInference.ts

@@ -110,13 +110,15 @@ export class RemoteInferencer implements ICompletions {
     }
   }
 
-  async code_completion(prompt, promptAfter, options:IParams=CompletionParams): Promise<any> {
-    const payload = { prompt, 'context':promptAfter, "endpoint":"code_completion", ...options }
+  async code_completion(prompt, promptAfter, ctxFiles, fileName, options:IParams=CompletionParams): Promise<any> {
+    const payload = { prompt, 'context':promptAfter, "endpoint":"code_completion",
+      'ctxFiles':ctxFiles, 'currentFileName':fileName, ...options }
     return this._makeRequest(payload, AIRequestType.COMPLETION)
   }
 
-  async code_insertion(msg_pfx, msg_sfx, options:IParams=InsertionParams): Promise<any> {
-    const payload = { "endpoint":"code_insertion", msg_pfx, msg_sfx, ...options, prompt: '' }
+  async code_insertion(msg_pfx, msg_sfx, ctxFiles, fileName, options:IParams=InsertionParams): Promise<any> {
+    const payload = { "endpoint":"code_insertion", msg_pfx, msg_sfx, 'ctxFiles':ctxFiles,
+      'currentFileName':fileName, ...options, prompt: '' }
     return this._makeRequest(payload, AIRequestType.COMPLETION)
   }
 

libs/remix-ai-core/src/types/types.ts

@@ -50,8 +50,8 @@ export interface InferenceModel {
 }
 
 export interface ICompletions{
-  code_completion(context, params:IParams): Promise<any>;
-  code_insertion(msg_pfx, msg_sfx, params:IParams): Promise<any>;
+  code_completion(context, ctxFiles, fileName, params:IParams): Promise<any>;
+  code_insertion(msg_pfx, msg_sfx, ctxFiles, fileName, params:IParams): Promise<any>;
 }
 
 export interface IParams {

package.json

@@ -171,6 +171,7 @@
     "jszip": "^3.6.0",
     "just-once": "^2.2.0",
     "latest-version": "^5.1.0",
+    "lunr": "^2.3.9",
     "merge": "^2.1.1",
     "npm-install-version": "^6.0.2",
     "octokit": "^3.1.2",

yarn.lock

@@ -16080,6 +16080,11 @@ flatted@^3.1.0:
   resolved "https://registry.yarnpkg.com/flatted/-/flatted-3.2.7.tgz#609f39207cb614b89d0765b477cb2d437fbf9787"
   integrity sha512-5nqDSxl8nn5BSNxyR3n4I6eDmbolI6WT+QqR547RwxQapgjQBmtktdP+HTBb/a/zLsbzERTONyUB5pefh5TtjQ==
 
+flexsearch@^0.7.43:
+  version "0.7.43"
+  resolved "https://registry.yarnpkg.com/flexsearch/-/flexsearch-0.7.43.tgz#34f89b36278a466ce379c5bf6fb341965ed3f16c"
+  integrity sha512-c5o/+Um8aqCSOXGcZoqZOm+NqtVwNsvVpWv6lfmSclU954O3wvQKxxK8zj74fPaSJbXpSLTs4PRhh+wnoCXnKg==
+
 flora-colossus@^1.0.0:
   version "1.0.1"
   resolved "https://registry.yarnpkg.com/flora-colossus/-/flora-colossus-1.0.1.tgz#aba198425a8185341e64f9d2a6a96fd9a3cbdb93"
@@ -21348,6 +21353,11 @@ lru_map@^0.3.3:
   resolved "https://registry.yarnpkg.com/lru_map/-/lru_map-0.3.3.tgz#b5c8351b9464cbd750335a79650a0ec0e56118dd"
   integrity sha512-Pn9cox5CsMYngeDbmChANltQl+5pi6XmTrraMSzhPmMBbmgcxmqWry0U3PGapCU1yB4/LqCcom7qhHZiF/jGfQ==
 
+lunr@^2.3.9:
+  version "2.3.9"
+  resolved "https://registry.yarnpkg.com/lunr/-/lunr-2.3.9.tgz#18b123142832337dd6e964df1a5a7707b25d35e1"
+  integrity sha512-zTU3DaZaF3Rt9rhN3uBMGQD3dD2/vFQqnvZCDv4dl5iOzq2IZQqTxu90r4E5J+nP70J3ilqVCrbho2eWaeW8Ow==
+
 lz-string@^1.4.4:
   version "1.4.4"
   resolved "https://registry.yarnpkg.com/lz-string/-/lz-string-1.4.4.tgz#c0d8eaf36059f705796e1e344811cf4c498d3a26"