Update Controller_JarvisFlow.py

Controller_JarvisFlow.py

@@ -16,10 +16,41 @@ class Command:
 
 # TODO: controller should be generalized
 class Controller_JarvisFlow(ChatAtomicFlow):
+    """This class is a controller for JarvisFlow, it takes the plan generated by the planner, logs of previous executions,
+    depending on the initial goal or the subsequent feedback from the branching executors (and the human), to decide which
+    executor to call next (or to exit by calling finish).
+
+    *Configuration Parameters*:
+    - `commands` (dict): a dictionary of commands that the controller can call, each command has a name, a description, and a list of input arguments.
+        The commands will be injected into the system message prompt template.
+    - `system_message_prompt_template` (str): the template for the system message prompt, there are several components needs to be injected into the
+        template, including the commands, plan, plan_file_location, logs, and the goal. The injection of commands is done then initalizing the flow,
+        the rest of the components are injected at the beginning of each run.
+    - `previous_messages` (int): a sliding window of previous messages that will be passed to the model. This is the central part of short-term memory management.
+
+    *Input Interface Non Initialized*:
+    - `goal` (str): the initial goal of the conversation, this is the input to the model.
+    - `memory_files` (dict): a dictionary of file locations that contains the plan, logs.
+    - `plan` (str): the plan generated by the planner, the plan will change (marked as done, or re-plan) as execution preceeds.
+    - `logs` (str): the logs of previous executions, the logs will be appended as execution preceeds.
+
+    *Input Interface Initialized*:
+    - `result` (str): the result of the previous execution, this is the input to the model.
+    - `memory_files` (dict): a dictionary of file locations that contains the plan, logs.
+    - `plan` (str): the plan generated by the planner, the plan will change (marked as done, or re-plan) as execution preceeds.
+    - `logs` (str): the logs of previous executions, the logs will be appended as execution preceeds.
+    - `goal` (str): the initial goal, this is kept because the goal is also injected into the system prompts so that Jarvis does not
+        forget what the goal is, when the memory sliding window is implemented.
+
+    *Output Interface*:
+    - `command` (str): the command to be executed by the executor.
+    - `command_args` (dict): the arguments of the command to be executed by the executor.
+    """
     def __init__(
             self,
             commands: List[Command],
             **kwargs):
+        """Initialize the flow, inject the commands into the system message prompt template."""
         super().__init__(**kwargs)
         self.system_message_prompt_template = self.system_message_prompt_template.partial(
             commands=self._build_commands_manual(commands),
@@ -53,6 +84,7 @@ class Controller_JarvisFlow(ChatAtomicFlow):
         return content
     @staticmethod
     def _build_commands_manual(commands: List[Command]) -> str:
+        """Build the manual for the commands."""
         ret = ""
         for i, command in enumerate(commands):
             command_input_json_schema = json.dumps(
@@ -63,6 +95,7 @@ class Controller_JarvisFlow(ChatAtomicFlow):
 
     @classmethod
     def instantiate_from_config(cls, config):
+        """Setting up the flow from the config file. In particular, setting up the prompts, backend, and commands."""
         flow_config = deepcopy(config)
 
         kwargs = {"flow_config": flow_config}
@@ -85,6 +118,7 @@ class Controller_JarvisFlow(ChatAtomicFlow):
         return cls(**kwargs)
 
     def _update_prompts_and_input(self, input_data: Dict[str, Any]):
+        """Hinting the model to output in json format, updating the plan, logs to the system prompts."""
         if 'goal' in input_data:
             input_data['goal'] += self.hint_for_model
         if 'result' in input_data:
@@ -107,7 +141,7 @@ class Controller_JarvisFlow(ChatAtomicFlow):
             self._state_update_add_chat_message(content=updated_system_message_content,
                                                 role=self.flow_config["system_name"])
 
-
+        # ~~~run the model, special mechanism to deal with situations where the output is not in json format. ~~~
         while True:
             api_output = super().run(input_data)["api_output"].strip()
             try: