reducing ai context to imported files

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

@@ -113,10 +113,10 @@ export class RemixAIPlugin extends Plugin {
   }
 
   async code_completion(prompt: string, promptAfter: string, params:IParams=CompletionParams): Promise<any> {
-    if (this.completionAgent.indexer == null || this.completionAgent.indexer == undefined) await this.completionAgent.indexWorkspace()
     params.provider = 'mistralai' // default provider for code completion
     const currentFileName = await this.call('fileManager', 'getCurrentFile')
-    const contextfiles = await this.completionAgent.getContextFiles(prompt)
+    const contextfiles = await this.completionAgent.getContextFiles()
+    console.log('completion file context is:', contextfiles)
     if (this.isOnDesktop && !this.useRemoteInferencer) {
       return await this.call(this.remixDesktopPluginName, 'code_completion', prompt, promptAfter, contextfiles, currentFileName, params)
     } else {
@@ -281,11 +281,11 @@ export class RemixAIPlugin extends Plugin {
   }
 
   async code_insertion(msg_pfx: string, msg_sfx: string, params:IParams=CompletionParams): Promise<any> {
-    if (this.completionAgent.indexer == null || this.completionAgent.indexer == undefined) await this.completionAgent.indexWorkspace()
 
     params.provider = 'mistralai' // default provider for code completion
     const currentFileName = await this.call('fileManager', 'getCurrentFile')
-    const contextfiles = await this.completionAgent.getContextFiles(msg_pfx)
+    const contextfiles = await this.completionAgent.getContextFiles()
+    console.log('insertion file context is:', contextfiles)
     if (this.isOnDesktop && !this.useRemoteInferencer) {
       return await this.call(this.remixDesktopPluginName, 'code_insertion', msg_pfx, msg_sfx, contextfiles, currentFileName, params)
     } else {

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

@@ -1,82 +1,9 @@
-import lunr from 'lunr';
 import { extractImportsFromFile } from "../helpers/localImportsExtractor";
 
-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',
-  javascript = '.js',
-  typescript = '.ts',
-  tests_ts = '.test.ts',
-  tests_js = '.test.js',
-}
-
 export class CodeCompletionAgent {
   props: any;
-  indexer: any;
-  Documents: Document[] = [];
-  INDEX_THRESHOLD = 0.05;
-  N_MATCHES = 1;
-  indexed: indexT = {
-    isIndexed: false,
-    lastIndexedTime: 0,
-    reason: 'Init',
-  };
-
   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" } });
-  }
-
-  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 [];
-    }
   }
 
   async getLocalImports(fileContent: string, currentFile: string) {
@@ -90,72 +17,28 @@ export class CodeCompletionAgent {
     }
   }
 
-  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) {
+  async getContextFiles() {
     try {
-      if (!this.indexed.isIndexed) {
-        await this.indexWorkspace();
+      const currentFile = await this.props.call('fileManager', 'getCurrentFile');
+      const currentFileContent = await this.props.call('fileManager', 'readFile', currentFile);
+
+      const localImports = await this.getLocalImports(currentFileContent, currentFile);
+
+      // Only return context files that are actually imported by the current file
+      const fileContentPairs = [];
+      for (const importPath of localImports) {
+        try {
+          const content = await this.props.call('fileManager', 'readFile', importPath);
+          fileContentPairs.push({ file: importPath, content: content });
+        } catch (error) {
+          continue;
+        }
       }
 
-      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;
+      return fileContentPairs;
     } 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 };
-    });
-
-    const localImports = await this.getLocalImports(await this.props.call('fileManager', 'readFile', currentFile), currentFile);
-    // check if the local import is in fileContentPairs file
-    for (const li of localImports) {
-      if (fileContentPairs.find(fcp => fcp.file === li)) continue;
-      const currentContent = await this.props.call('fileManager', 'readFile', li);
-      fileContentPairs.push({ file: li, content: currentContent });
-    }
-    return fileContentPairs;
-  }
-
-}
+}