diff --git a/apps/remix-ide/src/app/plugins/remixAIPlugin.tsx b/apps/remix-ide/src/app/plugins/remixAIPlugin.tsx index 0159b7d938..5853a1e6f9 100644 --- a/apps/remix-ide/src/app/plugins/remixAIPlugin.tsx +++ b/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 = { @@ -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 { + 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 { + 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){ diff --git a/libs/remix-ai-core/src/agents/completionAgent.ts b/libs/remix-ai-core/src/agents/completionAgent.ts index 1bf14005c7..ae2866f91c 100644 --- a/libs/remix-ai-core/src/agents/completionAgent.ts +++ b/libs/remix-ai-core/src/agents/completionAgent.ts @@ -1,23 +1,169 @@ -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(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 []; + } + } - constructor(codebasePath: string) { - // git or fs - this.codebase = this.loadCodebase(codebasePath); + async getLocalImports(fileContent: string, currentFile: string) { + try { + const lines = fileContent.split('\n'); + const imports = []; + + for (const line of lines) { + const trimmedLine = line.trim(); + if (trimmedLine.startsWith('import')) { + const parts = trimmedLine.split(' '); + if (parts.length >= 2) { + const importPath = parts[1].replace(/['";]/g, ''); + imports.push(importPath); + } + } + } + // Only local imports are those files that are in the workspace + const localImports = this.Documents.length >0 ? imports.filter((imp) => {return this.Documents.find((doc) => doc.filename === imp);}) : []; + + return localImports; + } catch (error) { + return []; + } } - 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')); + 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' }; } - public getSuggestions(currentLine: string, numSuggestions: number = 3): string[] { - const suggestions: string[] = []; - // get `numSuggestions` from the llm - return suggestions; + 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 }; + }); + + 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; + } + } diff --git a/libs/remix-ai-core/src/index.ts b/libs/remix-ai-core/src/index.ts index 130f040c48..d24ac6abae 100644 --- a/libs/remix-ai-core/src/index.ts +++ b/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' \ No newline at end of file +export * from './agents/completionAgent' +export * from './agents/securityAgent' diff --git a/libs/remix-ai-core/src/inferencers/remote/remoteInference.ts b/libs/remix-ai-core/src/inferencers/remote/remoteInference.ts index dbee3e37ab..a8fa702fc7 100644 --- a/libs/remix-ai-core/src/inferencers/remote/remoteInference.ts +++ b/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 { - const payload = { prompt, 'context':promptAfter, "endpoint":"code_completion", ...options } + async code_completion(prompt, promptAfter, ctxFiles, fileName, options:IParams=CompletionParams): Promise { + 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 { - const payload = { "endpoint":"code_insertion", msg_pfx, msg_sfx, ...options, prompt: '' } + async code_insertion(msg_pfx, msg_sfx, ctxFiles, fileName, options:IParams=InsertionParams): Promise { + const payload = { "endpoint":"code_insertion", msg_pfx, msg_sfx, 'ctxFiles':ctxFiles, + 'currentFileName':fileName, ...options, prompt: '' } return this._makeRequest(payload, AIRequestType.COMPLETION) } diff --git a/libs/remix-ai-core/src/types/types.ts b/libs/remix-ai-core/src/types/types.ts index 1e5d9fb603..3e0417031f 100644 --- a/libs/remix-ai-core/src/types/types.ts +++ b/libs/remix-ai-core/src/types/types.ts @@ -50,8 +50,8 @@ export interface InferenceModel { } export interface ICompletions{ - code_completion(context, params:IParams): Promise; - code_insertion(msg_pfx, msg_sfx, params:IParams): Promise; + code_completion(context, ctxFiles, fileName, params:IParams): Promise; + code_insertion(msg_pfx, msg_sfx, ctxFiles, fileName, params:IParams): Promise; } export interface IParams { diff --git a/libs/remix-ui/editor/src/lib/providers/inlineCompletionProvider.ts b/libs/remix-ui/editor/src/lib/providers/inlineCompletionProvider.ts index 5f15022b10..8c9443cdbe 100644 --- a/libs/remix-ui/editor/src/lib/providers/inlineCompletionProvider.ts +++ b/libs/remix-ui/editor/src/lib/providers/inlineCompletionProvider.ts @@ -106,7 +106,6 @@ export class RemixInLineCompletionProvider implements monacoTypes.languages.Inli const generatedText = output // no need to clean it. should already be this.task = 'code_insertion' - _paq.push(['trackEvent', 'ai', 'remixAI', this.task]) const item: monacoTypes.languages.InlineCompletion = { insertText: generatedText, range: new monaco.Range(position.lineNumber, position.column, position.lineNumber, position.column) @@ -163,7 +162,6 @@ export class RemixInLineCompletionProvider implements monacoTypes.languages.Inli let clean = data // if clean starts with a comment, remove it if (clean.startsWith('//') || clean.startsWith('/*') || clean.startsWith('*') || clean.startsWith('*/')){ - console.log("clean starts with comment") return "" } diff --git a/package.json b/package.json index 5ba5c0c00a..385a5d9a50 100644 --- a/package.json +++ b/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", diff --git a/yarn.lock b/yarn.lock index 89c2ba8cd6..57ae2d08ae 100644 --- a/yarn.lock +++ b/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"