diff --git a/apps/remix-ide/src/app/plugins/remixAIPlugin.tsx b/apps/remix-ide/src/app/plugins/remixAIPlugin.tsx index 56ddc22d894..cd6e470839a 100644 --- a/apps/remix-ide/src/app/plugins/remixAIPlugin.tsx +++ b/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 { - 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 { - 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 { diff --git a/libs/remix-ai-core/src/agents/completionAgent.ts b/libs/remix-ai-core/src/agents/completionAgent.ts index 692de5ca36d..83045ec259a 100644 --- a/libs/remix-ai-core/src/agents/completionAgent.ts +++ b/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; - } - -} +} \ No newline at end of file