diff --git a/src/app/parser/typescript/cleanup.ts b/src/app/parser/typescript/cleanup.ts
index 180738a..90b6158 100644
--- a/src/app/parser/typescript/cleanup.ts
+++ b/src/app/parser/typescript/cleanup.ts
@@ -1,7 +1,9 @@
 import * as ts from "ts-morph";
 import * as types from "../../types";
 import * as context from "../../context";
+import * as logger from "../../../util/logger";
 import * as fs from "fs";
+import * as path from "path";
 
 export function fixImports(generatedCodeList: types.GeneratedCode[]) {
   let project: ts.Project;
@@ -23,12 +25,31 @@ export function fixImports(generatedCodeList: types.GeneratedCode[]) {
     );
   }
 
-  for (const sourceFile of project.getSourceFiles()) {
-    sourceFile.fixMissingImports().organizeImports();
+  try {
+    for (const sourceFile of project.getSourceFiles()) {
+      sourceFile.fixMissingImports().organizeImports();
 
-    const code = generatedCodeList.filter(
-      (item) => item.filePath === sourceFile.getFilePath(),
-    )[0]!;
-    code.fileContent = sourceFile.getFullText();
+      // find the source file in the generated code file list
+      let curSourceFile =
+        generatedCodeList.filter(
+          (item) => item.filePath === sourceFile.getFilePath(),
+        )[0] ??
+        // if the source file is not found by path comparison, try to find it by comparing the file name
+        generatedCodeList.filter(
+          (item) =>
+            path.basename(item.filePath) ===
+            path.basename(sourceFile.getFilePath()),
+        )[0];
+
+      if (curSourceFile) {
+        curSourceFile!.fileContent = sourceFile.getFullText();
+      } else {
+        logger.error(
+          `Error while fixing imports: Unable to find the source file for ${sourceFile.getFilePath()}\n\nSkipping import fixing and cleanup for this file`,
+        );
+      }
+    }
+  } catch (error) {
+    logger.error("Error while fixing imports. Skipping", error);
   }
 }