Merge pull request #6164 from ethereum/ai_local_imports

Ai local imports

Author: STetsing

apps/remix-ide-e2e/src/tests/ai_panel.test.ts

@@ -130,7 +130,7 @@ const tests = {
         timeout: 120000
       })
       .waitForElementPresent('*[data-id="remix-ai-assistant-ready"]')
-      .assistantGenerate('a simple ERC20 contract', 'mistralai')
+      .assistantGenerate('a simple ERC20 contract', 'anthropic')
       .waitForElementVisible({
         locateStrategy: 'xpath',
         selector: '//div[contains(@class,"chat-bubble") and contains(.,"New workspace created:")]',
@@ -231,39 +231,6 @@ const tests = {
         selector: "//*[@data-id='remix-ai-streaming' and @data-streaming='false']"
       })
   },
-  'Generate new workspaces code with all AI assistant providers #group1': function (browser: NightwatchBrowser) {
-    browser
-      .assistantClearChat()
-      .waitForCompilerLoaded()
-      .clickLaunchIcon('remixaiassistant')
-
-      .waitForElementPresent('*[data-id="remix-ai-assistant-ready"]')
-
-      .assistantGenerate('a simple ERC20 contract', 'openai')
-      .waitForElementVisible({
-        locateStrategy: 'xpath',
-        selector: '//div[contains(@class,"chat-bubble") and contains(.,"New workspace created:")]',
-        timeout: 60000
-      })
-      .waitForElementPresent({
-        locateStrategy: 'xpath',
-        selector: "//*[@data-id='remix-ai-streaming' and @data-streaming='false']"
-      })
-      .assistantClearChat()
-
-      .clickLaunchIcon('remixaiassistant')
-
-      .assistantGenerate('a simple ERC20 contract', 'anthropic')
-      .waitForElementVisible({
-        locateStrategy: 'xpath',
-        selector: '//div[contains(@class,"chat-bubble") and contains(.,"New workspace created:")]',
-        timeout: 60000
-      })
-      .waitForElementPresent({
-        locateStrategy: 'xpath',
-        selector: "//*[@data-id='remix-ai-streaming' and @data-streaming='false']"
-      })
-  },
   "Should close the AI assistant #group1": function (browser: NightwatchBrowser) {
     browser
       .click('*[data-id="movePluginToLeft"]')

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

@@ -154,6 +154,16 @@ export class RemixAIPlugin extends Plugin {
 
   async error_explaining(prompt: string, params: IParams=GenerationParams): Promise<any> {
     let result
+    let localFilesImports = ""
+
+    // Get local imports from the workspace restrict to 5 most relevant files
+    const relevantFiles = this.workspaceAgent.getRelevantLocalFiles(prompt, 5);
+
+    for (const file in relevantFiles) {
+      localFilesImports += `\n\nFileName: ${file}\n\n${relevantFiles[file]}`
+    }
+    localFilesImports = localFilesImports + "\n End of local files imports.\n\n"
+    prompt = localFilesImports ? `Using the following local imports: ${localFilesImports}\n\n` + prompt : prompt
     if (this.isOnDesktop && !this.useRemoteInferencer) {
       result = await this.call(this.remixDesktopPluginName, 'error_explaining', prompt)
     } else {
@@ -183,17 +193,20 @@ export class RemixAIPlugin extends Plugin {
    * Generates a new remix IDE workspace based on the provided user prompt, optionally using Retrieval-Augmented Generation (RAG) context.
    * - If `useRag` is `true`, the function fetches additional context from a RAG API and prepends it to the user prompt.
    */
-  async generate(userPrompt: string, params: IParams=AssistantParams, newThreadID:string="", useRag:boolean=false): Promise<any> {
+  async generate(prompt: string, params: IParams=AssistantParams, newThreadID:string="", useRag:boolean=false, statusCallback?: (status: string) => Promise<void>): Promise<any> {
     params.stream_result = false // enforce no stream result
     params.threadId = newThreadID
-    params.provider = this.assistantProvider
+    params.provider = 'anthropic' // enforce all generation to be only on anthropic
+    useRag = false
     _paq.push(['trackEvent', 'ai', 'remixAI', 'GenerateNewAIWorkspace'])
+    let userPrompt = ''
 
     if (useRag) {
+      statusCallback?.('Fetching RAG context...')
       try {
         let ragContext = ""
         const options = { headers: { 'Content-Type': 'application/json', } }
-        const response = await axios.post(endpointUrls.rag, { query: userPrompt, endpoint:"query" }, options)
+        const response = await axios.post(endpointUrls.rag, { query: prompt, endpoint:"query" }, options)
         if (response.data) {
           ragContext = response.data.response
           userPrompt = "Using the following context: ```\n\n" + JSON.stringify(ragContext) + "```\n\n" + userPrompt
@@ -203,32 +216,37 @@ export class RemixAIPlugin extends Plugin {
       } catch (error) {
         console.log('RAG context error:', error)
       }
+    } else {
+      userPrompt = prompt
     }
     // Evaluate if this function requires any context
     // console.log('Generating code for prompt:', userPrompt, 'and threadID:', newThreadID)
-    let result
-    if (this.isOnDesktop && !this.useRemoteInferencer) {
-      result = await this.call(this.remixDesktopPluginName, 'generate', userPrompt, params)
-    } else {
-      result = await this.remoteInferencer.generate(userPrompt, params)
-    }
+    await statusCallback?.('Generating new workspace with AI...\nThis might take some minutes. Please be patient!')
+    const result = await this.remoteInferencer.generate(userPrompt, params)
+
+    await statusCallback?.('Creating contracts and files...')
+    const genResult = await this.contractor.writeContracts(result, userPrompt, statusCallback)
 
-    const genResult = await this.contractor.writeContracts(result, userPrompt)
     if (genResult.includes('No payload')) return genResult
     await this.call('menuicons', 'select', 'filePanel')
     return genResult
   }
 
   /**
-   * Performs any user action on the entire curren workspace or updates the workspace based on a user prompt, optionally using Retrieval-Augmented Generation (RAG) for additional context.
+   * Performs any user action on the entire curren workspace or updates the workspace based on a user prompt,
+   * optionally using Retrieval-Augmented Generation (RAG) for additional context.
    *
    */
-  async generateWorkspace (userPrompt: string, params: IParams=AssistantParams, newThreadID:string="", useRag:boolean=false): Promise<any> {
+  async generateWorkspace (userPrompt: string, params: IParams=AssistantParams, newThreadID:string="", useRag:boolean=false, statusCallback?: (status: string) => Promise<void>): Promise<any> {
     params.stream_result = false // enforce no stream result
     params.threadId = newThreadID
     params.provider = this.assistantProvider
+    useRag = false
     _paq.push(['trackEvent', 'ai', 'remixAI', 'WorkspaceAgentEdit'])
+
+    await statusCallback?.('Performing workspace request...')
     if (useRag) {
+      await statusCallback?.('Fetching RAG context...')
       try {
         let ragContext = ""
         const options = { headers: { 'Content-Type': 'application/json', } }
@@ -244,25 +262,20 @@ export class RemixAIPlugin extends Plugin {
         console.log('RAG context error:', error)
       }
     }
+    await statusCallback?.('Loading workspace context...')
     const files = !this.workspaceAgent.ctxFiles ? await this.workspaceAgent.getCurrentWorkspaceFiles() : this.workspaceAgent.ctxFiles
     userPrompt = "Using the following workspace context: ```\n" + files + "```\n\n" + userPrompt
 
-    let result
-    if (this.isOnDesktop && !this.useRemoteInferencer) {
-      result = await this.call(this.remixDesktopPluginName, 'generateWorkspace', userPrompt, params)
-    } else {
-      result = await this.remoteInferencer.generateWorkspace(userPrompt, params)
-    }
-    return (result !== undefined) ? this.workspaceAgent.writeGenerationResults(result) : "### No Changes applied!"
+    await statusCallback?.('Generating workspace updates with AI...')
+    const result = await this.remoteInferencer.generateWorkspace(userPrompt, params)
+
+    await statusCallback?.('Applying changes to workspace...')
+    return (result !== undefined) ? this.workspaceAgent.writeGenerationResults(result, statusCallback) : "### No Changes applied!"
   }
 
-  async fixWorspaceErrors(continueGeneration=false): Promise<any> {
+  async fixWorspaceErrors(): Promise<any> {
     try {
-      if (continueGeneration) {
-        return this.contractor.continueCompilation()
-      } else {
-        return this.contractor.fixWorkspaceCompilationErrors(this.workspaceAgent)
-      }
+      return this.contractor.fixWorkspaceCompilationErrors(this.workspaceAgent)
     } catch (error) {
     }
   }

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

@@ -1,4 +1,5 @@
 import lunr from 'lunr';
+import { extractImportsFromFile } from "../helpers/localImportsExtractor";
 
 interface Document {
   id: number;
@@ -80,23 +81,10 @@ export class CodeCompletionAgent {
 
   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;
+      // Use extractImportsFromFile to get all imports from the current file
+      const imports = extractImportsFromFile(fileContent);
+      const localImports = imports.filter(imp => imp.isLocal);
+      return localImports.map(imp => imp.importPath);
     } catch (error) {
       return [];
     }

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

@@ -1,7 +1,7 @@
 import { AssistantParams } from "../types/models";
 import { workspaceAgent } from "./workspaceAgent";
 import { CompilationResult } from "../types/types";
-import { compilecontracts } from "../helpers/compile";
+import { compilecontracts, compilationParams } from "../helpers/compile";
 
 const COMPILATION_WARNING_MESSAGE = '⚠️**Warning**: The compilation failed. Please check the compilation errors in the Solidity compiler plugin. Enter `/continue` or `/c` if you want Remix AI to try again until a compilable solution is generated?'
 
@@ -12,11 +12,9 @@ export class ContractAgent {
   generationThreadID: string= ''
   workspaceName: string = ''
   contracts: any = {}
-  performCompile: boolean = false
-  overrideWorkspace: boolean = false
+  mainPrompt: string = ''
   static instance
   oldPayload: any = undefined
-  mainPrompt: string
 
   private constructor(props) {
     this.plugin = props;
@@ -33,23 +31,29 @@ export class ContractAgent {
   * Write the result of the generation to the workspace. Compiles the contracts one time and creates a new workspace.
   * @param payload - The payload containing the generated files
   * @param userPrompt - The user prompt used to generate the files
+  * @param statusCallback - Optional callback for status updates in chat window
   */
-  async writeContracts(payload, userPrompt) {
+  async writeContracts(payload, userPrompt, statusCallback?: (status: string) => Promise<void>) {
+    await statusCallback?.('Getting current workspace info...')
     const currentWorkspace = await this.plugin.call('filePanel', 'getCurrentWorkspace')
 
     const writeAIResults = async (parsedResults) => {
       if (this.plugin.isOnDesktop) {
+        await statusCallback?.('Preparing files for desktop...')
         const files = parsedResults.files.reduce((acc, file) => {
           acc[file.fileName] = file.content
           return acc
         }, {})
+        await statusCallback?.('Opening in new window...')
         await this.plugin.call('electronTemplates', 'loadTemplateInNewWindow', files)
         //return "Feature not only available in the browser version of Remix IDE. Please use the browser version to generate secure code."
         return "## New workspace created!  \nNavigate to the new window!"
       }
 
+      await statusCallback?.('Creating new workspace...')
       await this.createWorkspace(this.workspaceName)
-      if (!this.overrideWorkspace) await this.plugin.call('filePanel', 'switchToWorkspace', { name: this.workspaceName, isLocalHost: false })
+
+      await statusCallback?.('Writing files to workspace...')
       const dirCreated = []
       for (const file of parsedResults.files) {
         const dir = file.fileName.split('/').slice(0, -1).join('/')
@@ -66,44 +70,70 @@ export class ContractAgent {
 
     try {
       if (payload === undefined) {
+        await statusCallback?.('No payload received, retrying...')
         this.nAttempts += 1
         if (this.nAttempts > this.generationAttempts) {
-          this.performCompile = false
           if (this.oldPayload) {
             return await writeAIResults(this.oldPayload)
           }
           return "Max attempts reached! Please try again with a different prompt."
         }
-        return "No payload, try again while considering changing the assistant provider with the command `/setAssistant <openai|anthorpic|mistralai>`"
+        return "No payload, try again while considering changing the assistant provider to one of these choices `<openai|anthropic|mistralai>`"
       }
+
+      await statusCallback?.('Processing generated files...')
       this.contracts = {}
       const parsedFiles = payload
       this.oldPayload = payload
       this.generationThreadID = parsedFiles['threadID']
       this.workspaceName = parsedFiles['projectName']
 
       this.nAttempts += 1
-      if (this.nAttempts === 1) this.mainPrompt = userPrompt
+      if (this.nAttempts === 1) this.mainPrompt=userPrompt
 
       if (this.nAttempts > this.generationAttempts) {
         return await writeAIResults(parsedFiles)
       }
 
+      await statusCallback?.('Processing Solidity contracts...')
+      const genContrats = []
       for (const file of parsedFiles.files) {
         if (file.fileName.endsWith('.sol')) {
           this.contracts[file.fileName] = { content: file.content }
+          genContrats.push({ fileName: file.fileName, content: file.content })
         }
       }
 
+      await statusCallback?.('Compiling contracts...')
       const result:CompilationResult = await compilecontracts(this.contracts, this.plugin)
-      if (!result.compilationSucceeded && this.performCompile) {
+      if (!result.compilationSucceeded) {
+        await statusCallback?.('Compilation failed, fixing errors...')
         // console.log('Compilation failed, trying again recursively ...')
-        const newPrompt = `Payload:\n${JSON.stringify(result.errfiles)}}\n\nWhile considering this compilation error: Here is the error message\n. Try this again:${this.mainPrompt}\n `
-        return await this.plugin.generate(newPrompt, AssistantParams, this.generationThreadID); // reuse the same thread
+        const generatedContracts = genContrats.map(contract =>
+          `File: ${contract.fileName}\n${contract.content}`
+        ).join('\n\n');
+
+        // Format error files properly according to the schema
+        const formattedErrorFiles = Object.entries(result.errfiles).map(([fileName, fileData]: [string, any]) => {
+          const errors = fileData.errors.map((err: any) =>
+            `Error at ${err.errorStart}-${err.errorEnd}: ${err.errorMessage}`
+          ).join('\n  ');
+          return `File: ${fileName}\n  ${errors}`;
+        }).join('\n\n');
+
+        const newPrompt = `
+              Compilation parameters:\n${JSON.stringify(compilationParams)}\n\n
+              Compilation errors:\n${formattedErrorFiles}\n\n
+              Generated contracts:\n${generatedContracts}\n\nConsider other possible solutions and retry this main prompt again: \n${this.mainPrompt}\n `
+
+        await statusCallback?.('Regenerating workspace with error fixes...')
+        return await this.plugin.generate(newPrompt, AssistantParams, this.generationThreadID, false, statusCallback); // reuse the same thread
       }
 
+      await statusCallback?.('Finalizing workspace creation...')
       return result.compilationSucceeded ? await writeAIResults(parsedFiles) : await writeAIResults(parsedFiles) + "\n\n" + COMPILATION_WARNING_MESSAGE
     } catch (error) {
+      await statusCallback?.('Error occurred, cleaning up...')
       this.deleteWorkspace(this.workspaceName )
       this.nAttempts = 0
       await this.plugin.call('filePanel', 'switchToWorkspace', currentWorkspace)
@@ -148,21 +178,4 @@ export class ContractAgent {
     }
   }
 
-  async continueCompilation(){
-    try {
-      if (this.oldPayload === undefined) {
-        return "No payload, try again while considering changing the assistant provider with the command `/setAssistant <openai|anthorpic|mistralai>`"
-      }
-
-      this.performCompile = true
-      this.overrideWorkspace = true
-      return await this.writeContracts(this.oldPayload, this.mainPrompt)
-    } catch (error) {
-      return "Error during continue compilation. Please try again."
-    } finally {
-      this.performCompile = false
-      this.overrideWorkspace = false
-    }
-  }
-
 }