Update RunCodeFileEditAtomicFlow.py

RunCodeFileEditAtomicFlow.py

@@ -54,6 +54,15 @@ class RunCodeFileEditAtomicFlow(CodeFileEditAtomicFlow):
     def _generate_input_to_writer(self, input_data: Dict[str, Any]):
         code_str = input_data['code']
         code_lib_location = input_data["memory_files"]["code_library"]
+        
+        # if we are running python code and if the code is importing the code library after one modification, we need
+        # to re-import the library.
+        code_lang = input_data["language"].lower()
+        code_lib_name = os.path.splitext(os.path.basename(code_lib_location))[0]
+        if code_lang == 'python' and code_lib_name in code_str:
+            prepend_code = f"import importlib\nimport {code_lib_name}\nimportlib.reload({code_lib_name})\n"
+            code_str = prepend_code + code_str
+        
         lang_details = self._generate_file_extension_and_comment_sign(input_data["language"])
         content_to_write = self._generate_content(code_str, lang_details["comment"])
         file_location_to_write = self._generate_temp_file_location(code_lib_location, lang_details["extension"])