Update CodeFileEditAtomicFlow.py

CodeFileEditAtomicFlow.py

@@ -48,17 +48,21 @@ class CodeFileEditAtomicFlow(AtomicFlow):
         language_of_code = input_data["language_of_code"]
         assert language_of_code.lower() == 'python', "sorry!! only writing python code is supported."
 
+    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"]
+        content_to_write = self._generate_content(code_lib_location, code_str)
+        file_location_to_write = self._generate_temp_file_location(code_lib_location)
+        return content_to_write, file_location_to_write
+
     def run(
             self,
             input_data: Dict[str, Any]
     ):
         self._check_input(input_data)
 
-        # ~~~ Generating temp file location and content to write ~~~
-        code_str = input_data['code']
-        code_lib_location = input_data["memory_files"]["code_library"]
-        content_to_write = self._generate_content(code_lib_location, code_str)
-        file_location_to_write = self._generate_temp_file_location(code_lib_location)
+        # ~~~ Getting input data to the file editor ~~~
+        content_to_write, file_location_to_write = self._generate_input_to_writer(input_data)
 
         # ~~~ Calling the writer function ~~~
         result, code_editor_output, temp_file_location, file_written_timestamp = self._write_code_content_to_file(