coflows update

AutoGPTFlow.py

@@ -76,128 +76,43 @@ class AutoGPTFlow(ControllerExecutorFlow):
         self.rename_human_output_interface = KeyInterface(
             keys_to_rename={"human_input": "human_feedback"}
         )
-    
-    def set_up_flow_state(self):
-        super().set_up_flow_state()
-        self.flow_state["early_exit_flag"] = True
-        self.flow_state["is_first_round"] = True
-        
-    def memory_read_step(self):
-        memory_read_input = self.prepare_memory_read_input()
-        output = self.ask_subflow("Memory", memory_read_input).get_data()
-        memory_read_output = self.prepare_memory_read_output(output)
-        
-        return memory_read_output
-    
-    def memory_write_step(self):
-        memory_write_input = self.prepare_memory_write_input()
-        self.tell_subflow("Memory", memory_write_input)
-    
-    def controller_executor_step(self, output_memory_retrieval):
-        
-        if self.flow_state["is_first_round"]:
-            additional_input_ctrl_ex = {
-                "goal": self.flow_state["goal"],
-            }
-        else:
-            additional_input_ctrl_ex = {
-                "observation": self.flow_state["observation"],
-                "human_feedback": self.flow_state["human_feedback"],
-            }
         
-        input_ctrl_ex =  {"executor_reply": {**output_memory_retrieval,**additional_input_ctrl_ex}}
-        
-        output_ctrl_ex = self._single_round_controller_executor(input_ctrl_ex)
-        
-        self.flow_state["early_exit_flag"] = output_ctrl_ex.get("EARLY_EXIT",False)
-        
-        if self.flow_state["early_exit_flag"]:
-            return output_ctrl_ex
-        
-        controller_reply = output_ctrl_ex["controller_reply"]
-        executor_reply = output_ctrl_ex["executor_reply"]
-        
-        self._state_update_dict(
-            {
-                "command": controller_reply["command"],
-                "command_args": controller_reply["command_args"],
-                "observation": executor_reply["observation"],
-            }
+        self.input_interface_controller = KeyInterface(
+            keys_to_select = ["goal","observation","human_feedback", "memory"],
         )
-        return output_ctrl_ex
-    
-    
-    def human_feedback_step(self):
-            
-        human_feedback_input_variables = {
-            "goal": self.flow_state["goal"],
-            "command": self.flow_state["command"],
-            "command_args": self.flow_state["command_args"],
-            "observation": self.flow_state["observation"],
-        }
         
-        human_feedback = self.rename_human_output_interface(
-            self.ask_subflow("HumanFeedback", human_feedback_input_variables).get_data()
+        self.input_interface_human_feedback = KeyInterface(
+            keys_to_select = ["goal","command","command_args","observation"],
         )
         
-        self.flow_state["human_feedback"] = human_feedback["human_feedback"]
-        
-        if human_feedback["human_feedback"].strip().lower() == "q":
-            self.flow_state["early_exit_flag"] = True
-            return {
-                "EARLY_EXIT": True,
-                "answer": "The user has chosen to exit before a final answer was generated.",
-                "status": "unfinished",
-            }
-        
-        return human_feedback
-        
-    
-    def _single_round_autogpt(self):
-        
-        
-        #1. Memory Retrieval
-        output_memory_retrieval = self.memory_read_step()
-        
-        
-        #2. ControllerExecutor
-        output_ctrl_ex = self.controller_executor_step(output_memory_retrieval)
-        
-        if self.flow_state["early_exit_flag"]:
-            return output_ctrl_ex
-        
-        #3. HumanFeedback
-        output_human_feedback = self.human_feedback_step()
-        
-        if self.flow_state["early_exit_flag"]:
-            return output_human_feedback
-        
-        #4. Memory Write
-        self.memory_write_step()
-        
-        return {** output_ctrl_ex, **output_human_feedback}
-        
-        
-        
-        
-    
-    def run(self,input_data):
-        self._state_update_dict({"goal": input_data["goal"]})
+        self.memory_read_ouput_interface = KeyInterface(
+            additional_transformations = [self.prepare_memory_read_output],
+            keys_to_select = ["memory"],
+        )
         
-        for round in range(self.flow_config["max_rounds"]):
-            output = self._single_round_autogpt()
+        self.human_feedback_ouput_interface = KeyInterface(
+            keys_to_rename={"human_input": "human_feedback"},
+            keys_to_select = ["human_feedback"],
+        )
             
-            self.flow_state["is_first_round"] = False
+        self.next_flow_to_call = {
+            None: "MemoryRead",
+            "MemoryRead": "Controller",
+            "Controller": "Executor",
+            "Executor": "HumanFeedback",
+            "HumanFeedback": "MemoryWrite",
+            "MemoryWrite": "MemoryRead",
+        }
             
-            if self.flow_state["early_exit_flag"]:
-                return output
+    def set_up_flow_state(self):
+        super().set_up_flow_state()
+        self.flow_state["early_exit_flag"] = False
         
-        return {
-            "EARLY_EXIT": False,
-            "answer": output,
-            "status": "unfinished"
-        }
+    def prepare_memory_read_output(self,data_dict: Dict[str, Any],**kwargs):
         
+        retrieved_memories = data_dict["retrieved"][0][1:]
+        return {"memory": "\n".join(retrieved_memories)}
+    
     def _get_memory_key(self):
         """ This method returns the memory key that is used to retrieve memories from the ChromaDB model.
         
@@ -246,11 +161,6 @@ class AutoGPTFlow(ControllerExecutorFlow):
             "content": query
         }
 
-    def prepare_memory_read_output(self, data: Dict[str, Any]):
-    
-        retrieved_memories = data["retrieved"][0][1:]
-        return {"memory": "\n".join(retrieved_memories)}
-
     def prepare_memory_write_input(self) -> Dict[str, Any]:
 
         query = self._get_memory_key()
@@ -259,3 +169,139 @@ class AutoGPTFlow(ControllerExecutorFlow):
             "operation": "write",
             "content": str(query)
         }
+
+    
+    def call_memory_read(self):
+        memory_read_input = self.prepare_memory_read_input()
+        
+        message = self._package_input_message(
+            data = memory_read_input,
+            dst_flow = "Memory",
+        )
+        
+        self.subflows["Memory"].send_message_async(
+            message,
+            pipe_to = self.flow_config["flow_ref"]
+        )
+
+    def call_memory_write(self):
+        memory_write_input = self.prepare_memory_write_input()
+        
+        message = self._package_input_message(
+            data = memory_write_input,
+            dst_flow = "Memory",
+        )
+        
+        self.subflows["Memory"].send_message_async(
+            message,
+            pipe_to = self.flow_config["flow_ref"]
+        )
+    
+    def call_human_feedback(self):
+            
+        message = self._package_input_message(
+            data = self.input_interface_human_feedback(self.flow_state),
+            dst_flow = "HumanFeedback",
+        )
+        
+        self.subflows["HumanFeedback"].send_message_async(
+            message,
+            pipe_to = self.flow_config["flow_ref"]
+        )
+             
+    def register_data_to_state(self, input_message):
+        
+        #Making this explicit so it's easier to understand
+        #I'm also showing different ways of writing to the state
+        # either programmatically or using the _state_update_dict and 
+        # input and ouput interface methods
+        
+        last_called = self.flow_state["last_called"]
+        
+        if last_called is None:            
+            self.flow_state["input_message"] = input_message
+            self.flow_state["goal"] = input_message.data["goal"]
+        
+        elif last_called == "Executor":
+            self.flow_state["observation"] = input_message.data
+        
+        elif last_called == "Controller":
+            self._state_update_dict(
+                {
+                    "command": input_message.data["command"],
+                    "command_args": input_message.data["command_args"]
+                }
+            )
+            
+            #detect and early exit
+            if self.flow_state["command"] == "finish":
+                
+                self._state_update_dict(
+                    {
+                        "EARLY_EXIT": True,
+                        "answer": self.flow_state["command_args"]["answer"],
+                        "status": "finished"
+                    }
+                )
+                self.flow_state["early_exit_flag"] = True
+                
+                
+        elif last_called == "MemoryRead":
+            self._state_update_dict(
+                self.memory_read_ouput_interface(input_message).data
+            )
+            
+        elif last_called == "HumanFeedback":
+            self._state_update_dict(
+                self.human_feedback_ouput_interface(input_message).data
+            )
+            
+            #detect early exit
+            if self.flow_state["human_feedback"].strip().lower() == "q":
+                
+                self._state_update_dict(
+                    {
+                        "EARLY_EXIT": True,
+                        "answer": "The user has chosen to exit before a final answer was generated.",
+                        "status": "unfinished",
+                    }
+                )
+                
+                self.flow_state["early_exit_flag"] = True
+        
+        
+            
+        
+    
+    def run(self,input_message):
+        
+        self.register_data_to_state(input_message)
+        
+        flow_to_call = self.get_next_flow_to_call()
+
+        if self.flow_state.get("early_exit_flag",False):
+            self.generate_reply()
+            
+        elif flow_to_call == "MemoryRead":
+            self.call_memory_read()
+        
+        elif flow_to_call == "Controller":
+            self.call_controller()
+        
+        elif flow_to_call == "Executor":
+            self.call_executor()
+        
+        elif flow_to_call == "HumanFeedback":
+            self.call_human_feedback()
+        
+        elif flow_to_call == "MemoryWrite":
+            self.call_memory_write()
+            self.flow_state["current_round"] += 1
+            
+        else:
+            self._on_reach_max_round()
+            self.generate_reply()
+            
+        self.flow_state["last_called"] = self.get_next_flow_to_call()
+        
+    

__init__.py

@@ -3,9 +3,9 @@ dependencies = [
     {"url": "aiflows/ControllerExecutorFlowModule",
      "revision": "coflows"},
     {"url": "aiflows/HumanStandardInputFlowModule",
-     "revision": "main"},
+     "revision": "coflows"},
     {"url": "aiflows/VectorStoreFlowModule",
-     "revision": "main"},
+     "revision": "coflows"},
 ]
 from aiflows import flow_verse
 

run.py

@@ -1,3 +1,5 @@
+"""A simple script to run a Flow that can be used for development and debugging."""
+
 import os
 
 import hydra
@@ -9,91 +11,109 @@ from aiflows.utils.general_helpers import read_yaml_file, quick_load_api_keys
 
 from aiflows import logging
 from aiflows.flow_cache import CACHING_PARAMETERS, clear_cache
+
 from aiflows.utils import serve_utils
 from aiflows.workers import run_dispatch_worker_thread
 from aiflows.messages import FlowMessage
 from aiflows.interfaces import KeyInterface
+from aiflows.utils.colink_utils import start_colink_server
+from aiflows.workers import run_dispatch_worker_thread
 
 CACHING_PARAMETERS.do_caching = False  # Set to True in order to disable caching
 # clear_cache() # Uncomment this line to clear the cache
 
 logging.set_verbosity_debug()
 
-from aiflows import flow_verse
-# ~~~ Load Flow dependecies from FlowVerse ~~~
+
 dependencies = [
     {"url": "aiflows/AutoGPTFlowModule", "revision": os.getcwd()},
-    {"url": "aiflows/LCToolFlowModule", "revision": "80c0c76181d90846ebff1057b8951d9689f93b62"},
+    {"url": "aiflows/LCToolFlowModule", "revision": "coflows"},
 ]
 
+from aiflows import flow_verse
 flow_verse.sync_dependencies(dependencies)
 if __name__ == "__main__":
-    # ~~~ Set the API information ~~~
+    
+    #1. ~~~~~ Set up a colink server ~~~~
+    FLOW_MODULES_PATH = "./"
+    
+    cl = start_colink_server()
+
+
+    #2. ~~~~~Load flow config~~~~~~
+    root_dir = "."
+    cfg_path = os.path.join(root_dir, "demo.yaml")
+    cfg = read_yaml_file(cfg_path)
+    
+    #2.1 ~~~ Set the API information ~~~
     # OpenAI backend
-    api_information = [ApiInfo(backend_used="openai", api_key=os.getenv("OPENAI_API_KEY"))]
-    # Azure backend
+    api_information = [ApiInfo(backend_used="openai",
+                              api_key = os.getenv("OPENAI_API_KEY"))]
+    # # Azure backend
     # api_information = ApiInfo(backend_used = "azure",
     #                           api_base = os.getenv("AZURE_API_BASE"),
     #                           api_key = os.getenv("AZURE_OPENAI_KEY"),
     #                           api_version =  os.getenv("AZURE_API_VERSION") )
     
-    FLOW_MODULES_PATH = "./"
-    
-    jwt = os.getenv("COLINK_JWT")
-    addr = os.getenv("LOCAL_COLINK_ADDRESS")
     
-    cl = serve_utils.start_colink_component(
-        "Reverse Number Demo",
-        {"jwt": jwt, "addr": addr}
-    )
+    quick_load_api_keys(cfg, api_information, key="api_infos")
 
-    root_dir = "."
-    cfg_path = os.path.join(root_dir, "demo.yaml")
-    cfg = read_yaml_file(cfg_path)
     
+    #3. ~~~~ Serve The Flow ~~~~
     serve_utils.recursive_serve_flow(
         cl = cl,
-        flow_type="demo_served",
+        flow_type="AutoGPTFlowModule",
         default_config=cfg,
         default_state=None,
-        default_dispatch_point="coflows_dispatch",
+        default_dispatch_point="coflows_dispatch"
     )
     
-    #in case you haven't started the dispatch worker thread, uncomment the 2 lines
-    #run_dispatch_worker_thread(cl, dispatch_point="coflows_dispatch", flow_modules_base_path=FLOW_MODULES_PATH)
-    #run_dispatch_worker_thread(cl, dispatch_point="coflows_dispatch", flow_modules_base_path=FLOW_MODULES_PATH)
+    #4. ~~~~~Start A Worker Thread~~~~~
+    run_dispatch_worker_thread(cl, dispatch_point="coflows_dispatch", flow_modules_base_path=FLOW_MODULES_PATH)
 
-    quick_load_api_keys(cfg, api_information, key="api_infos")
-
-
-    # ~~~ Get the data ~~~
-    # data = {"id": 0, "goal": "Answer the following question: What is the population of Canada?"}  # Uses wikipedia
-    # data = {"id": 0, "goal": "Answer the following question: Who was the NBA champion in 2023?"}  # Uses duckduckgo
-    data = {
-        "id": 0,
-        "goal": "Answer the following question: What is the profession and date of birth of Michael Jordan?",
-    }
-    # At first, we retrieve information about Michael Jordan the basketball player
-    # If we provide feedback, only in the first round, that we are not interested in the basketball player,
-    #   but the statistician, and skip the feedback in the next rounds, we get the correct answer
-
-    # ~~~ Run inference ~~~
+    #5. ~~~~~Mount the flow and get its proxy~~~~~~
     proxy_flow = serve_utils.recursive_mount(
         cl=cl,
         client_id="local",
-        flow_type="demo_served",
-        config_overrides=cfg,
+        flow_type="AutoGPTFlowModule",
+        config_overrides=None,
         initial_state=None,
         dispatch_point_override=None,
     )   
-    # ~~~ Print the output ~~~
+    
+    #6. ~~~ Get the data ~~~
+    data = {
+        "id": 0,
+        "goal": "Answer the following question: What is the profession and date of birth of Michael Jordan?",
+    }
+   
+    
+    #option1: use the FlowMessage class
     input_message = FlowMessage(
-        data= data,
-        src_flow="Coflows team",
-        dst_flow=proxy_flow,
-        is_input_msg=True
+        data=data,
     )
+
+    #option2: use the proxy_flow
+    #input_message = proxy_flow._package_input_message(data = data)
+    
+    #7. ~~~ Run inference ~~~
+    future = proxy_flow.send_message_blocking(input_message)
+    
+    #uncomment this line if you would like to get the full message back
+    #reply_message = future.get_message()
+    reply_data = future.get_data()
+    
+    # ~~~ Print the output ~~~
+    print("~~~~~~Reply~~~~~~")
+    print(reply_data)
+    
+    
+    #8. ~~~~ (Optional) apply output interface on reply ~~~~
+    # output_interface = KeyInterface(
+    #     keys_to_rename={"api_output": "answer"},
+    # )
+    # print("Output: ", output_interface(reply_data))
     
-    future = proxy_flow.ask(input_message)
     
-    print(future.get_data())
+    #9. ~~~~~Optional: Unserve Flow~~~~~~
+    # serve_utils.delete_served_flow(cl, "FlowModule")