From 8f4ad2398be5de3aaad55b12fba7857bce244fc8 Mon Sep 17 00:00:00 2001
From: STetsing <41009393+STetsing@users.noreply.github.com>
Date: Wed, 12 Feb 2025 14:48:53 +0100
Subject: [PATCH] added local import files to context

---
 .../src/agents/completionAgent.ts             | 32 +++++++++++++++++++
 .../src/inferencers/remote/remoteInference.ts |  2 +-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/libs/remix-ai-core/src/agents/completionAgent.ts b/libs/remix-ai-core/src/agents/completionAgent.ts
index d0df5b1bcf..ae2866f91c 100644
--- a/libs/remix-ai-core/src/agents/completionAgent.ts
+++ b/libs/remix-ai-core/src/agents/completionAgent.ts
@@ -74,6 +74,30 @@ 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;
+    } catch (error) {
+      return [];
+    }
+  }
+
   indexWorkspace() {
     this.getDcocuments().then((documents) => {
       this.indexer =lunr(function () {
@@ -131,6 +155,14 @@ export class CodeCompletionAgent {
       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/inferencers/remote/remoteInference.ts b/libs/remix-ai-core/src/inferencers/remote/remoteInference.ts
index a8fa702fc7..92eb0aea6a 100644
--- a/libs/remix-ai-core/src/inferencers/remote/remoteInference.ts
+++ b/libs/remix-ai-core/src/inferencers/remote/remoteInference.ts
@@ -12,7 +12,7 @@ export class RemoteInferencer implements ICompletions {
   max_history = 7
   model_op = RemoteBackendOPModel.CODELLAMA // default model operation change this to llama if necessary
   event: EventEmitter
-  test_env=false
+  test_env=true
   test_url="http://solcodertest.org"
 
   constructor(apiUrl?:string, completionUrl?:string) {