Update PlanFileEditAtomicFlow.py

PlanFileEditAtomicFlow.py

@@ -45,7 +45,7 @@ class PlanFileEditAtomicFlow(AtomicFlow):
         assert os.path.isfile(plan_file_loc), f"{plan_file_loc} is not a file"
 
     def _generate_input_to_writer(self, input_data: Dict[str, Any]):
-        plan_str = input_data['plan']
+        plan_str = input_data['plan'] if "plan" in input_data else input_data['new_plan']
         plan_file_location = input_data["plan_file_location"]
         content_to_write = self._generate_content(plan_file_location, plan_str)
         file_location_to_write = self._generate_temp_file_location(plan_file_location)