cleaned up run

ControllerExecutorFlow.py

@@ -64,33 +64,44 @@ class ControllerExecutorFlow(CompositeFlow):
     :param flow_config: The configuration of the flow (see Configuration Parameters).
     :param subflows: A list of subflows. Required when instantiating the subflow programmatically (it replaces subflows_config from flow_config).                                  
     """
-    def __init__(self, flow_config: Dict[str, Any], subflows: Dict[str, Any] = None):
-        super().__init__(flow_config, subflows)
+    
+    def _on_reach_max_round(self):
+        """ This method is called when the flow reaches the maximum amount of rounds. It updates the state of the flow and starts the process of terminating the flow."""
+        self._state_update_dict({
+            "answer": "The maximum amount of rounds was reached before the model found an answer.",
+            "status": "unfinished"
+        })
         
-    def set_up_flow_state(self):
-        super().set_up_flow_state()
+    def _single_round_controller_executor(self, executor_reply: Dict[str,Any]) -> Dict[str,Any]:
         
+        controller_reply = self.ask_subflow("Controller", executor_reply).get_data()
         
+        if controller_reply["command"] == "finish":
+            return {
+                "EARLY_EXIT": True,
+                "answer": controller_reply["command_args"]["answer"],
+                "status": "finished"
+            }
+        
+        executor_reply = {
+                "observation":  self.ask_subflow(controller_reply["command"], controller_reply["command_args"]).get_data()
+        }
+            
+        return executor_reply
+    
     def run(self,input_data):
 
         executor_reply = input_data
         
         for round in range(self.flow_config["max_rounds"]):
-            controller_reply = self.ask_subflow("Controller", executor_reply).get_data()
-            if controller_reply["command"] == "finish":
-                return {
-                    "EARLY_EXIT": True,
-                    "answer": controller_reply["command_args"]["answer"],
-                    "status": "finished"
-                }
-            
-            
+            executor_reply = self._single_round_controller_executor(executor_reply)
             
-            executor_reply = {
-                    "observation":  self.ask_subflow(controller_reply["command"], controller_reply["command_args"]).get_data()
-            }
-           
-        return  {
-            "answer": "The maximum amount of rounds was reached before the model found an answer.",
+            if executor_reply.get("EARLY_EXIT",False):
+                return executor_reply
+                
+        self._on_reach_max_round()
+        return {
+            "EARLY_EXIT": False,
+            "answer": executor_reply["observation"],
             "status": "unfinished"
-        }  
+        }