Skip to content

Commit fbc4029

Browse files
STetsingSTetsing
committed
reducing ai context to imported files
1 parent d553f90 commit fbc4029

File tree

2 files changed

+46
-140
lines changed

2 files changed

+46
-140
lines changed

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

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -113,10 +113,9 @@ export class RemixAIPlugin extends Plugin {
113113
}
114114

115115
async code_completion(prompt: string, promptAfter: string, params:IParams=CompletionParams): Promise<any> {
116-
if (this.completionAgent.indexer == null || this.completionAgent.indexer == undefined) await this.completionAgent.indexWorkspace()
117116
params.provider = 'mistralai' // default provider for code completion
118117
const currentFileName = await this.call('fileManager', 'getCurrentFile')
119-
const contextfiles = await this.completionAgent.getContextFiles(prompt)
118+
const contextfiles = await this.completionAgent.getContextFiles()
120119
if (this.isOnDesktop && !this.useRemoteInferencer) {
121120
return await this.call(this.remixDesktopPluginName, 'code_completion', prompt, promptAfter, contextfiles, currentFileName, params)
122121
} else {
@@ -268,11 +267,10 @@ export class RemixAIPlugin extends Plugin {
268267
}
269268

270269
async code_insertion(msg_pfx: string, msg_sfx: string, params:IParams=CompletionParams): Promise<any> {
271-
if (this.completionAgent.indexer == null || this.completionAgent.indexer == undefined) await this.completionAgent.indexWorkspace()
272270

273271
params.provider = 'mistralai' // default provider for code completion
274272
const currentFileName = await this.call('fileManager', 'getCurrentFile')
275-
const contextfiles = await this.completionAgent.getContextFiles(msg_pfx)
273+
const contextfiles = await this.completionAgent.getContextFiles()
276274
if (this.isOnDesktop && !this.useRemoteInferencer) {
277275
return await this.call(this.remixDesktopPluginName, 'code_insertion', msg_pfx, msg_sfx, contextfiles, currentFileName, params)
278276
} else {

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

Lines changed: 44 additions & 136 deletions
Original file line numberDiff line numberDiff line change
@@ -1,81 +1,8 @@
1-
import lunr from 'lunr';
2-
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-
javascript = '.js',
21-
typescript = '.ts',
22-
tests_ts = '.test.ts',
23-
tests_js = '.test.js',
24-
}
25-
261
export class CodeCompletionAgent {
272
props: any;
28-
indexer: any;
29-
Documents: Document[] = [];
30-
INDEX_THRESHOLD = 0.05;
31-
N_MATCHES = 1;
32-
indexed: indexT = {
33-
isIndexed: false,
34-
lastIndexedTime: 0,
35-
reason: 'Init',
36-
};
373

384
constructor(props) {
395
this.props = props;
40-
this.listenForChanges();
41-
this.indexer =lunr(function () {
42-
this.ref('id')
43-
this.field('filename')
44-
this.field('content')
45-
this.field('Identifier');
46-
});
47-
48-
setInterval(() => {
49-
this.indexWorkspace()
50-
}, 60000)
51-
}
52-
53-
listenForChanges() {
54-
this.props.on('fileManager', 'fileAdded', (path) => { this.indexed = { isIndexed: false, reason:"fileAdded" } });
55-
this.props.on('fileManager', 'fileRemoved', (path) => { this.indexed = { isIndexed: false, reason:"fileRemoved" } });
56-
this.props.on('filePanel', 'workspaceCreated', () => { this.indexed = { isIndexed: false, reason:"workspaceCreated" } });
57-
this.props.on('filePanel', 'workspaceRenamed', () => { this.indexed = { isIndexed: false, reason:"workspaceRenamed" }});
58-
this.props.on('filePanel', 'workspaceDeleted', () => { this.indexed = { isIndexed: false, reason:"workspaceDeleted" } });
59-
}
60-
61-
async getDcocuments() {
62-
try {
63-
const documents: Document[] = [];
64-
const jsonDirsContracts = await this.props.call('fileManager', 'copyFolderToJson', '/').then((res) => res.contracts);
65-
let c = 0;
66-
for (const file in jsonDirsContracts.children) {
67-
if (!Object.values(SupportedFileExtensions).some(ext => file.endsWith(ext))) continue;
68-
documents.push({
69-
id: ++c,
70-
filename: file,
71-
content: jsonDirsContracts.children[file].content,
72-
identifier: c - 1,
73-
});
74-
}
75-
return documents;
76-
} catch (error) {
77-
return [];
78-
}
796
}
807

818
async getLocalImports(fileContent: string, currentFile: string) {
@@ -86,88 +13,69 @@ export class CodeCompletionAgent {
8613
for (const line of lines) {
8714
const trimmedLine = line.trim();
8815
if (trimmedLine.startsWith('import')) {
89-
const parts = trimmedLine.split(' ');
90-
if (parts.length >= 2) {
91-
const importPath = parts[1].replace(/['";]/g, '');
16+
const importMatch = trimmedLine.match(/import\s+(?:.*\s+from\s+)?["']([^"']+)["']/);
17+
if (importMatch) {
18+
let importPath = importMatch[1];
19+
20+
// Skip library imports (npm packages, node_modules, etc.)
21+
if (importPath.startsWith('@') ||
22+
(!importPath.startsWith('./') && !importPath.startsWith('../') && !importPath.startsWith('/'))) {
23+
continue;
24+
}
25+
26+
// Handle relative imports
27+
if (importPath.startsWith('./') || importPath.startsWith('../')) {
28+
const currentDir = currentFile.includes('/') ? currentFile.substring(0, currentFile.lastIndexOf('/')) : '';
29+
importPath = this.resolvePath(currentDir, importPath);
30+
}
31+
9232
imports.push(importPath);
9333
}
9434
}
9535
}
96-
// Only local imports are those files that are in the workspace
97-
const localImports = this.Documents.length >0 ? imports.filter((imp) => {return this.Documents.find((doc) => doc.filename === imp);}) : [];
9836

99-
return localImports;
37+
return imports;
10038
} catch (error) {
10139
return [];
10240
}
10341
}
10442

105-
indexWorkspace() {
106-
this.getDcocuments().then((documents) => {
107-
this.indexer =lunr(function () {
108-
this.ref('id')
109-
this.field('filename')
110-
this.field('content')
111-
this.field('Identifier');
43+
resolvePath(currentDir: string, relativePath: string): string {
44+
const parts = currentDir.split('/').filter(part => part !== '');
45+
const relativeParts = relativePath.split('/').filter(part => part !== '');
11246

113-
documents.forEach(doc => {
114-
this.add(doc);
115-
});
116-
});
117-
this.Documents = documents;
118-
});
47+
for (const part of relativeParts) {
48+
if (part === '..') {
49+
parts.pop();
50+
} else if (part !== '.') {
51+
parts.push(part);
52+
}
53+
}
11954

120-
this.indexed = { isIndexed: true, lastIndexedTime: Date.now(), reason: 'init Indexing' };
55+
return parts.length > 0 ? parts.join('/') : relativePath;
12156
}
12257

123-
async getContextFiles(prompt) {
58+
async getContextFiles() {
12459
try {
125-
if (!this.indexed.isIndexed) {
126-
await this.indexWorkspace();
60+
const currentFile = await this.props.call('fileManager', 'getCurrentFile');
61+
const currentFileContent = await this.props.call('fileManager', 'readFile', currentFile);
62+
63+
const localImports = await this.getLocalImports(currentFileContent, currentFile);
64+
65+
// Only return context files that are actually imported by the current file
66+
const fileContentPairs = [];
67+
for (const importPath of localImports) {
68+
try {
69+
const content = await this.props.call('fileManager', 'readFile', importPath);
70+
fileContentPairs.push({ file: importPath, content: content });
71+
} catch (error) {
72+
continue;
73+
}
12774
}
12875

129-
const currentFile = await this.props.call('fileManager', 'getCurrentFile');
130-
const content = prompt;
131-
const searchResult = this.indexer.search(content)
132-
const fcps = await this.processResults(searchResult, currentFile);
133-
const resolvedFcps = await Promise.all(fcps);
134-
return resolvedFcps;
76+
return fileContentPairs;
13577
} catch (error) {
13678
return [];
13779
}
13880
}
139-
140-
async processResults(results: any, currentFile: string) {
141-
142-
// remove the current file name from the results list
143-
const rmResults = await results.filter(result => {
144-
return this.Documents.find(doc => doc.id === Number(result.ref)).filename !== currentFile;
145-
});
146-
147-
// filter out the results which have the same extension as the current file.
148-
// Do not mix and match file extensions as this will lead to incorrect completions
149-
const extResults = await rmResults.filter(result => {
150-
return this.Documents.find(doc => doc.id === Number(result.ref)).filename.split('.').pop() === currentFile.split('.').pop();
151-
});
152-
153-
// filter out the results which have a score less than the INDEX_THRESHOLD
154-
const topResults = await extResults.filter(result => result.score >= this.INDEX_THRESHOLD).slice(0, this.N_MATCHES);
155-
156-
// get the LATEST content of the top results in case the file has been modified and not indexed yet
157-
const fileContentPairs = topResults.map(async result => {
158-
const document = this.Documents.find(doc => doc.id === Number(result.ref));
159-
const currentContent = await this.props.call('fileManager', 'readFile', document.filename);
160-
return { file: document.filename, content: currentContent };
161-
});
162-
163-
const localImports = await this.getLocalImports(await this.props.call('fileManager', 'readFile', currentFile), currentFile);
164-
// check if the local import is in fileContentPairs file
165-
for (const li of localImports) {
166-
if (fileContentPairs.find(fcp => fcp.file === li)) continue;
167-
const currentContent = await this.props.call('fileManager', 'readFile', li);
168-
fileContentPairs.push({ file: li, content: currentContent });
169-
}
170-
return fileContentPairs;
171-
}
172-
17381
}

0 commit comments

Comments
 (0)