Update MemoryReadingAtomicFlow.py

MemoryReadingAtomicFlow.py

@@ -15,13 +15,15 @@ class MemoryReadingAtomicFlow(AtomicFlow):
     {"plan": "examples/JARVIS/plan.txt"}
     """
 
-    def __init__(self):
-        super().__init__()
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
         self.supported_mem_name = ["plan", "logs", "code_library"]
 
     def _check_input_data(self, input_data: Dict[str, Any]):
         """input data sanity check"""
-        for mem_name, mem_path in input_data.items():
+        assert "memory_files" in input_data, "memory_files not passed to MemoryReadingAtomicFlow"
+
+        for mem_name, mem_path in input_data["memory_files"].items():
             assert mem_name in self.supported_mem_name, (f"{mem_name} is not supported in MemoryReadingAtomicFlow, "
                                                          f"supported names are: {self.supported_mem_name}")
             assert os.path.exists(mem_path), f"{mem_path} does not exist."
@@ -50,7 +52,7 @@ class MemoryReadingAtomicFlow(AtomicFlow):
             input_data: Dict[str, Any]):
         self._check_input_data(input_data)
         response = {}
-        for mem_name, mem_path in input_data.items():
+        for mem_name, mem_path in input_data["memory_files"].items():
             if mem_name in ['plan', 'logs']:
                 response[mem_name] = self._read_text(mem_path)
             elif mem_name == 'code_library' and mem_path.endswith('.py'):