Update ControllerFlow_ExtLib.py

ControllerFlow_ExtLib.py

@@ -1,127 +1,65 @@
-import json
-from copy import deepcopy
-from typing import Any, Dict, List
-
-from flow_modules.aiflows.ChatFlowModule import ChatAtomicFlow
-
-from dataclasses import dataclass
-
-
-@dataclass
-class Command:
-    name: str
-    description: str
-    input_args: List[str]
-
-class ControllerFlow_ExtLib(ChatAtomicFlow):
-    def __init__(
-            self,
-            commands: List[Command],
-            **kwargs):
-        super().__init__(**kwargs)
-        self.system_message_prompt_template = self.system_message_prompt_template.partial(
-            commands=self._build_commands_manual(commands),
-            plan_file_location="no location yet",
-            plan="no plan yet",
-            logs="no logs yet",
-        )
-        self.hint_for_model = """
-        Make sure your response is in the following format:
-              Response Format:
-              {
-              "command": "call one of the commands you have e.g. `write_code`",
-              "command_args": {
-                  "arg name": "value"
-                  }
-              }
-        """
-
-    @staticmethod
-    def _build_commands_manual(commands: List[Command]) -> str:
-        ret = ""
-        for i, command in enumerate(commands):
-            command_input_json_schema = json.dumps(
-                {input_arg: f"YOUR_{input_arg.upper()}" for input_arg in command.input_args})
-            ret += f"{i + 1}. {command.name}: {command.description} Input arguments (given in the JSON schema): {command_input_json_schema}\n"
-        return ret
-
-    def _get_plan_file_location(self, input_data: Dict[str, Any]):
-        assert "memory_files" in input_data, "memory_files not passed to Extlib/Controller"
-        assert "plan" in input_data["memory_files"], "plan not in memory files"
-        return input_data["memory_files"]["plan"]
-
-    def _get_plan_content(self, input_data: Dict[str, Any]):
-        assert "plan" in input_data, "plan not passed to Extlib/Controller"
-        plan_content = input_data["plan"]
-        if len(plan_content) == 0:
-            plan_content = 'No plan yet'
-        return plan_content
-
-    def _get_logs_content(self, input_data: Dict[str, Any]):
-        assert "logs" in input_data, "logs not passed to Extlib/Controller"
-        logs_content = input_data["logs"]
-        if len(logs_content) == 0:
-            logs_content = "No logs yet"
-        return logs_content
-
-    @classmethod
-    def instantiate_from_config(cls, config):
-        flow_config = deepcopy(config)
-
-        kwargs = {"flow_config": flow_config}
-
-        # ~~~ Set up prompts ~~~
-        kwargs.update(cls._set_up_prompts(flow_config))
-
-        # ~~~Set up backend ~~~
-        kwargs.update(cls._set_up_backend(flow_config))
-
-        # ~~~ Set up commands ~~~
-        commands = flow_config["commands"]
-        commands = [
-            Command(name, command_conf["description"], command_conf["input_args"]) for name, command_conf in
-            commands.items()
-        ]
-        kwargs.update({"commands": commands})
-
-        # ~~~ Instantiate flow ~~~
-        return cls(**kwargs)
-
-    def _update_prompts_and_input(self, input_data: Dict[str, Any]):
-        if 'goal' in input_data:
-            input_data['goal'] += self.hint_for_model
-        if 'result' in input_data:
-            input_data['result'] += self.hint_for_model
-        plan_file_location = self._get_plan_file_location(input_data)
-        plan_content = self._get_plan_content(input_data)
-        logs_content = self._get_logs_content(input_data)
-        self.system_message_prompt_template = self.system_message_prompt_template.partial(
-            plan_file_location=plan_file_location,
-            plan=plan_content,
-            logs=logs_content
-        )
-
-    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
-        self._update_prompts_and_input(input_data)
-
-        # ~~~when conversation is initialized, append the updated system prompts to the chat history ~~~
-        if self._is_conversation_initialized():
-            updated_system_message_content = self._get_message(self.system_message_prompt_template, input_data)
-            self._state_update_add_chat_message(content=updated_system_message_content,
-                                                role=self.flow_config["system_name"])
-            
-        while True:
-            api_output = super().run(input_data)["api_output"].strip()
-            try:
-                start = api_output.index("{")
-                end = api_output.rindex("}") + 1
-                json_str = api_output[start:end]
-                return json.loads(json_str)
-            except (json.decoder.JSONDecodeError, json.JSONDecodeError):
-                updated_system_message_content = self._get_message(self.system_message_prompt_template, input_data)
-                self._state_update_add_chat_message(content=updated_system_message_content,
-                                                    role=self.flow_config["system_name"])
-                new_goal = "The previous respond cannot be parsed with json.loads. Next time, do not provide any comments or code blocks. Make sure your next response is purely json parsable."
-                new_input_data = input_data.copy()
-                new_input_data['result'] = new_goal
-                input_data = new_input_data
+from typing import Dict, Any
+
+from flow_modules.Tachi67.AbstractBossFlowModule import CtrlExMemFlow
+from aiflows.base_flows import CircularFlow
+
+class CtrlExMem_ExtLib(CtrlExMemFlow):
+    """This class inherits from the CtrlExMemFlow class from AbstractBossFlowModule.
+    See: https://huggingface.co/Tachi67/AbstractBossFlowModule/blob/main/CtrlExMemFlow.py
+    *Input Interface*:
+    - `plan`
+    - `logs`
+    - `memory_files`
+    - `goal`
+    *Output Interface*
+    - `result` (str): The result of the flow, the result will be returned to the caller.
+    - `summary` (str): The summary of the flow, the summary will be logged into the logs of the caller flow.
+
+    """
+    def _on_reach_max_round(self):
+        self._state_update_dict({
+            "result": "the maximum amount of rounds was reached before the extend library flow has done the job",
+            "summary": "ExtendLibraryFlow: the maximum amount of rounds was reached before the flow has done the job",
+            "status": "unfinished"
+        })
+
+    @CircularFlow.output_msg_payload_processor
+    def detect_finish_or_continue(self, output_payload: Dict[str, Any], src_flow) -> Dict[str, Any]:
+        command = output_payload["command"]
+        if command == "finish":
+            return {
+                "EARLY_EXIT": True,
+                "result": output_payload["command_args"]["summary"],
+                "summary": "ExtendLibrary: " + output_payload["command_args"]["summary"],
+                "status": "finished"
+            }
+        elif command == "manual_finish":
+            # ~~~ return the manual quit status ~~~
+            return {
+                "EARLY_EXIT": True,
+                "result": "ExtendLibraryFlow was terminated explicitly by the user, process is unfinished",
+                "summary": "ExtendLibrary: process terminated by the user explicitly, nothing generated",
+                "status": "unfinished"
+            }
+        elif command == "save_code":
+            keys_to_fetch_from_state = ["code", "memory_files"]
+            fetched_state = self._fetch_state_attributes_by_keys(keys=keys_to_fetch_from_state)
+            output_payload["command_args"]["code"] = fetched_state["code"]
+            output_payload["command_args"]["memory_files"] = fetched_state["memory_files"]
+            return output_payload
+
+        elif command == "update_plan":
+            keys_to_fetch_from_state = ["memory_files"]
+            fetched_state = self._fetch_state_attributes_by_keys(keys=keys_to_fetch_from_state)
+            output_payload["command_args"]["memory_files"] = fetched_state["memory_files"]
+            return output_payload
+
+        elif command == "re_plan":
+            keys_to_fetch_from_state = ["plan", "memory_files"]
+            fetched_state = self._fetch_state_attributes_by_keys(keys=keys_to_fetch_from_state)
+            output_payload["command_args"]["plan_file_location"] = fetched_state["memory_files"]["plan"]
+            output_payload["command_args"]["plan"] = fetched_state["plan"]
+            return output_payload
+
+        else:
+            return output_payload