Update ControllerFlow_ExtLib.py

ControllerFlow_ExtLib.py

@@ -1,65 +1,127 @@
-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
+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