|
1 | | -import * as fs from 'fs'; |
| 1 | +import lunr from 'lunr'; |
2 | 2 |
|
3 | | -class CodeCompletionAgent { |
4 | | - private codebase: string[]; |
| 3 | +interface Document { |
| 4 | + id: number; |
| 5 | + filename: string; |
| 6 | + content: string; |
| 7 | + identifier: number; |
| 8 | +} |
| 9 | + |
| 10 | +interface indexT{ |
| 11 | + isIndexed: boolean; |
| 12 | + lastIndexedTime?: number; |
| 13 | + reason?: string; |
| 14 | +} |
| 15 | + |
| 16 | +enum SupportedFileExtensions { |
| 17 | + solidity = '.sol', |
| 18 | + vyper = '.vy', |
| 19 | + circom = '.circom', |
| 20 | +} |
| 21 | + |
| 22 | +export class CodeCompletionAgent { |
| 23 | + props: any; |
| 24 | + indexer: any; |
| 25 | + Documents: Document[] = []; |
| 26 | + INDEX_THRESHOLD = 0.1; |
| 27 | + N_MATCHES = 1; |
| 28 | + indexed: indexT = { |
| 29 | + isIndexed: false, |
| 30 | + lastIndexedTime: 0, |
| 31 | + reason: 'Init', |
| 32 | + }; |
| 33 | + |
| 34 | + constructor(props) { |
| 35 | + this.props = props; |
| 36 | + this.listenForChanges(); |
| 37 | + this.indexer =lunr(function () { |
| 38 | + this.ref('id') |
| 39 | + this.field('filename') |
| 40 | + this.field('content') |
| 41 | + this.field('Identifier'); |
| 42 | + }); |
| 43 | + |
| 44 | + setInterval(() => { |
| 45 | + this.indexWorkspace() |
| 46 | + }, 60000) |
| 47 | + } |
| 48 | + |
| 49 | + listenForChanges() { |
| 50 | + this.props.on('fileManager', 'fileAdded', (path) => { this.indexed = { isIndexed: false, reason:"fileAdded" } }); |
| 51 | + this.props.on('fileManager', 'fileRemoved', (path) => { this.indexed = { isIndexed: false, reason:"fileRemoved" } }); |
| 52 | + this.props.on('filePanel', 'workspaceCreated', () => { this.indexed = { isIndexed: false, reason:"workspaceCreated" } }); |
| 53 | + this.props.on('filePanel', 'workspaceRenamed', () => { this.indexed = { isIndexed: false, reason:"workspaceRenamed" }}); |
| 54 | + this.props.on('filePanel', 'workspaceDeleted', () => { this.indexed = { isIndexed: false, reason:"workspaceDeleted" } }); |
| 55 | + } |
| 56 | + |
| 57 | + async getDcocuments() { |
| 58 | + try { |
| 59 | + const documents: Document[] = []; |
| 60 | + const jsonDirsContracts = await this.props.call('fileManager', 'copyFolderToJson', '/').then((res) => res.contracts); |
| 61 | + let c = 0; |
| 62 | + for (const file in jsonDirsContracts.children) { |
| 63 | + if (!Object.values(SupportedFileExtensions).some(ext => file.endsWith(ext))) continue; |
| 64 | + documents.push({ |
| 65 | + id: ++c, |
| 66 | + filename: file, |
| 67 | + content: jsonDirsContracts.children[file].content, |
| 68 | + identifier: c - 1, |
| 69 | + }); |
| 70 | + } |
| 71 | + return documents; |
| 72 | + } catch (error) { |
| 73 | + return []; |
| 74 | + } |
| 75 | + } |
5 | 76 |
|
6 | | - constructor(codebasePath: string) { |
7 | | - // git or fs |
8 | | - this.codebase = this.loadCodebase(codebasePath); |
| 77 | + async getLocalImports(fileContent: string, currentFile: string) { |
| 78 | + try { |
| 79 | + const lines = fileContent.split('\n'); |
| 80 | + const imports = []; |
| 81 | + |
| 82 | + for (const line of lines) { |
| 83 | + const trimmedLine = line.trim(); |
| 84 | + if (trimmedLine.startsWith('import')) { |
| 85 | + const parts = trimmedLine.split(' '); |
| 86 | + if (parts.length >= 2) { |
| 87 | + const importPath = parts[1].replace(/['";]/g, ''); |
| 88 | + imports.push(importPath); |
| 89 | + } |
| 90 | + } |
| 91 | + } |
| 92 | + // Only local imports are those files that are in the workspace |
| 93 | + const localImports = this.Documents.length >0 ? imports.filter((imp) => {return this.Documents.find((doc) => doc.filename === imp);}) : []; |
| 94 | + |
| 95 | + return localImports; |
| 96 | + } catch (error) { |
| 97 | + return []; |
| 98 | + } |
9 | 99 | } |
10 | 100 |
|
11 | | - private loadCodebase(path: string): string[] { |
12 | | - const files = fs.readdirSync(path); |
13 | | - return files |
14 | | - .filter(file => file.endsWith('.ts')) |
15 | | - .flatMap(file => fs.readFileSync(`${path}/${file}`, 'utf-8').split('\n')); |
| 101 | + indexWorkspace() { |
| 102 | + this.getDcocuments().then((documents) => { |
| 103 | + this.indexer =lunr(function () { |
| 104 | + this.ref('id') |
| 105 | + this.field('filename') |
| 106 | + this.field('content') |
| 107 | + this.field('Identifier'); |
| 108 | + |
| 109 | + documents.forEach(doc => { |
| 110 | + this.add(doc); |
| 111 | + }); |
| 112 | + }); |
| 113 | + this.Documents = documents; |
| 114 | + }); |
| 115 | + |
| 116 | + this.indexed = { isIndexed: true, lastIndexedTime: Date.now(), reason: 'init Indexing' }; |
16 | 117 | } |
17 | 118 |
|
18 | | - public getSuggestions(currentLine: string, numSuggestions: number = 3): string[] { |
19 | | - const suggestions: string[] = []; |
20 | | - // get `numSuggestions` from the llm |
21 | | - return suggestions; |
| 119 | + async getContextFiles(prompt) { |
| 120 | + try { |
| 121 | + if (!this.indexed.isIndexed) { |
| 122 | + await this.indexWorkspace(); |
| 123 | + } |
| 124 | + |
| 125 | + const currentFile = await this.props.call('fileManager', 'getCurrentFile'); |
| 126 | + const content = prompt; |
| 127 | + const searchResult = this.indexer.search(content) |
| 128 | + const fcps = await this.processResults(searchResult, currentFile); |
| 129 | + const resolvedFcps = await Promise.all(fcps); |
| 130 | + return resolvedFcps; |
| 131 | + } catch (error) { |
| 132 | + return []; |
| 133 | + } |
22 | 134 | } |
| 135 | + |
| 136 | + async processResults(results: any, currentFile: string) { |
| 137 | + |
| 138 | + // remove the current file name from the results list |
| 139 | + const rmResults = await results.filter(result => { |
| 140 | + return this.Documents.find(doc => doc.id === Number(result.ref)).filename !== currentFile; |
| 141 | + }); |
| 142 | + |
| 143 | + // filter out the results which have the same extension as the current file. |
| 144 | + // Do not mix and match file extensions as this will lead to incorrect completions |
| 145 | + const extResults = await rmResults.filter(result => { |
| 146 | + return this.Documents.find(doc => doc.id === Number(result.ref)).filename.split('.').pop() === currentFile.split('.').pop(); |
| 147 | + }); |
| 148 | + |
| 149 | + // filter out the results which have a score less than the INDEX_THRESHOLD |
| 150 | + const topResults = await extResults.filter(result => result.score >= this.INDEX_THRESHOLD).slice(0, this.N_MATCHES); |
| 151 | + |
| 152 | + // get the LATEST content of the top results in case the file has been modified and not indexed yet |
| 153 | + const fileContentPairs = topResults.map(async result => { |
| 154 | + const document = this.Documents.find(doc => doc.id === Number(result.ref)); |
| 155 | + const currentContent = await this.props.call('fileManager', 'readFile', document.filename); |
| 156 | + return { file: document.filename, content: currentContent }; |
| 157 | + }); |
| 158 | + |
| 159 | + const localImports = await this.getLocalImports(await this.props.call('fileManager', 'readFile', currentFile), currentFile); |
| 160 | + // check if the local import is in fileContentPairs file |
| 161 | + for (const li of localImports) { |
| 162 | + if (fileContentPairs.find(fcp => fcp.file === li)) continue; |
| 163 | + const currentContent = await this.props.call('fileManager', 'readFile', li); |
| 164 | + fileContentPairs.push({ file: li, content: currentContent }); |
| 165 | + } |
| 166 | + return fileContentPairs; |
| 167 | + } |
| 168 | + |
23 | 169 | } |
0 commit comments