coflows compatible+

ControllerAtomicFlow.py

@@ -2,7 +2,7 @@ import json
 from copy import deepcopy
 from typing import Any, Dict, List
 from flow_modules.aiflows.ChatFlowModule import ChatAtomicFlow
-
+from aiflows.messages import FlowMessage
 from dataclasses import dataclass
 
 
@@ -117,14 +117,28 @@ class ControllerAtomicFlow(ChatAtomicFlow):
         # ~~~ Instantiate flow ~~~
         return cls(**kwargs)
 
-    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+    def run(self, input_message: FlowMessage):
         """ This method runs the flow. Note that the response of the LLM is in the JSON format, but it's not a hard constraint (it can hallucinate and return an invalid JSON)
         
-        :param input_data: The input data of the flow.
-        :type input_data: Dict[str, Any]
-        :return: The output data of the flow (thought, reasoning, criticism, command, command_args)
-        :rtype: Dict[str, Any]
+        :param input_message: The input data of the flow.
+        :type input_message: FlowMessage
         """
-        api_output = super().run(input_data)["api_output"].strip()
+        
+        input_data = input_message.data
+        
+        if "goal" in input_data:
+            self.flow_state["goal"] = input_data["goal"]
+        
+        else:
+            input_data["goal"] = self.flow_state["goal"]
+        
+        api_output = self.query_llm(input_data)
+        
         response = json.loads(api_output)
-        return response
+        
+        reply = self._package_output_message(
+            input_message=input_message,
+            response=response
+        )
+        
+        self.reply_to_message(reply = reply, to = input_message)

ControllerAtomicFlow.yaml

@@ -11,6 +11,7 @@ input_interface_non_initialized: # initial input keys
 
 input_interface_initialized: # input_keys
   - "observation"
+  - "goal"
 
 #######################################################
 # Output keys
@@ -67,10 +68,13 @@ system_message_prompt_template:
 human_message_prompt_template:
   _target_: aiflows.prompt_template.JinjaPrompt
   template: |2-
+    Here is the goal you need to achieve:
+    {{goal}}
     Here is the response to your last action:
     {{observation}}
   input_variables:
     - "observation"
+    - "goal"
 
 init_human_message_prompt_template:
   _target_: aiflows.prompt_template.JinjaPrompt

ControllerExecutorFlow.py

@@ -4,6 +4,7 @@ from aiflows.base_flows import CompositeFlow
 from aiflows.utils import logging
 
 from .ControllerAtomicFlow import ControllerAtomicFlow
+from aiflows.messages import FlowMessage
 from aiflows.interfaces import KeyInterface
 logging.set_verbosity_debug()
 log = logging.get_logger(__name__)
@@ -65,51 +66,133 @@ class ControllerExecutorFlow(CompositeFlow):
     :param subflows: A list of subflows. Required when instantiating the subflow programmatically (it replaces subflows_config from flow_config).                                  
     """
     
+    def __init__(self,**kwargs):
+        super().__init__(**kwargs)
+            
+        self.input_interface_controller = KeyInterface(
+            keys_to_select = ["goal","observation"],
+        )
+        self.input_interface_first_round_controller = KeyInterface(
+            keys_to_select = ["goal"],
+        )
+        
+        self.reply_interface = KeyInterface(
+            keys_to_select = ["answer","status", "EARLY_EXIT"],
+        )
+        
+        self.next_flow_to_call = {
+            None: "Controller",
+            "Controller": "Executor",
+            "Executor": "Controller" 
+        } 
+    
+    def generate_reply(self):
+        """ This method generates the reply of the flow. It's called when the flow is finished. """
+        
+        reply = self._package_output_message(
+            input_message = self.flow_state["input_message"],
+            response = self.reply_interface(self.flow_state)
+        )
+        self.reply_to_message(reply,to=self.flow_state["input_message"])
+    
+    def get_next_flow_to_call(self):
+        
+        if self.flow_config["max_rounds"] is not None and self.flow_state["current_round"] >= self.flow_config["max_rounds"]:
+            return None
+        
+        return self.next_flow_to_call[self.flow_state["last_called"]]
+    
     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({
+            "EARLY_EXIT": False,
             "answer": "The maximum amount of rounds was reached before the model found an answer.",
             "status": "unfinished"
         })
         
-    def _single_round_controller_executor(self, reply: Dict[str,Any]) -> Dict[str,Any]:
+    def set_up_flow_state(self):
+        super().set_up_flow_state()
+        self.flow_state["last_called"] = None
+        self.flow_state["current_round"] = 0
         
-        in_data = reply["executor_reply"]
+    
+    def call_controller(self):
+        #first_round  
+        if self.flow_state["last_called"] is None:
+            input_interface = self.input_interface_first_round_controller
+        else:
+            input_interface = self.input_interface_controller
+     
+        message = self._package_input_message(
+            data = input_interface(self.flow_state),
+            dst_flow = "Controller"
+        )
         
-        controller_reply = self.ask_subflow("Controller", in_data).get_data()
+        self.subflows["Controller"].send_message_async(message, pipe_to= self.flow_config["flow_ref"])
         
-        if controller_reply["command"] == "finish":
-            return {
-                "EARLY_EXIT": True,
-                "answer": controller_reply["command_args"]["answer"],
-                "status": "finished"
-            }
+    def call_executor(self):
         
-        executor_reply = {
-                "observation":  self.ask_subflow(controller_reply["command"], controller_reply["command_args"]).get_data()
-        }
+        #detect and early exit
+        if self.flow_state["command"] == "finish":
             
-        return {
-            "controller_reply": controller_reply,
-            "executor_reply": executor_reply
-        }
-    
-    def run(self,input_data):
-
-        reply = {
-            "executor_reply": input_data,
-            "controller_reply": None
-        }
-        
-        for round in range(self.flow_config["max_rounds"]):
-            reply = self._single_round_controller_executor(reply)
+            self._state_update_dict(
+                {
+                    "EARLY_EXIT": True,
+                    "answer": self.flow_state["command_args"]["answer"],
+                    "status": "finished"
+                }
+            )
+            self.generate_reply()
             
-            if reply.get("EARLY_EXIT",False):
-                return reply
-                
-        self._on_reach_max_round()
-        return {
-            "EARLY_EXIT": False,
-            "answer": reply["executor_reply"]["observation"],
-            "status": "unfinished"
-        }
+        #call executor
+        else:
+            executor_branch_to_call = self.flow_state["command"]
+            message = self._package_input_message(
+                data = self.flow_state["command_args"],
+                dst_flow = executor_branch_to_call
+            )
+            
+            self.subflows[executor_branch_to_call].send_message_async(message, pipe_to= self.flow_config["flow_ref"])
+            
+            
+    def register_data_to_state(self, input_message):
+        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
+        
+        else:
+            self._state_update_dict(
+                {
+                    "command": input_message.data["command"],
+                    "command_args": input_message.data["command_args"]
+                }
+            )
+        
+    def run(self,input_message: FlowMessage):
+        """ Runs the WikiSearch Atomic Flow. It's used to execute a Wikipedia search and get page summaries.
+        
+        :param input_message: The input message
+        :type input_message: FlowMessage
+        """
+     
+        self.register_data_to_state(input_message)       
+        
+        flow_to_call = self.get_next_flow_to_call()
+        
+        if flow_to_call == "Controller":
+            self.flow_state["observation"] = input_message.data
+            self.call_controller()
+        elif flow_to_call == "Executor":
+            self.call_executor()
+            self.flow_state["current_round"] += 1
+        else:
+            self._on_reach_max_round()
+            self.generate_reply()
+        
+        self.flow_state["last_called"] = flow_to_call
+        

ControllerExecutorFlow.yaml

@@ -15,9 +15,10 @@ subflows_config:
     name: "ControllerAtomicFlow"
     description: "A flow that calls other flows to solve a problem."
     _target_: flow_modules.aiflows.ControllerAtomicFlow.instantiate_from_default_config
-    finish:
-      description: "Signal that the objective has been satisfied, and returns the answer to the user."
-      input_args: ["answer"]
+    commands:
+      finish:
+        description: "Signal that the objective has been satisfied, and returns the answer to the user."
+        input_args: ["answer"]
 # E.g.,
 #    commands:
 #      wiki_search:
@@ -29,6 +30,3 @@ subflows_config:
   # wiki_search:
   #   _target_: .WikiSearchAtomicFlow.instantiate_from_default_config
 
-
-early_exit_key: "EARLY_EXIT"
-

WikiSearchAtomicFlow.py

@@ -3,7 +3,7 @@ from copy import deepcopy
 from typing import List, Dict, Optional, Any
 
 from aiflows.base_flows import AtomicFlow
-
+from aiflows.messages import FlowMessage
 from aiflows.utils import logging
 from .wikipediaAPI import WikipediaAPIWrapper
 
@@ -44,15 +44,13 @@ class WikiSearchAtomicFlow(AtomicFlow):
         super().__init__(**kwargs)
 
     def run(self,
-            input_data: Dict[str, Any]) -> Dict[str, Any]:
+            input_message: FlowMessage):
         """ Runs the WikiSearch Atomic Flow. It's used to execute a Wikipedia search and get page summaries.
         
-        :param input_data: The input data dictionary
-        :type input_data: Dict[str, Any]
-        :return: The output data dictionary
-        :rtype: Dict[str, Any]
+        :param input_message: The input message
+        :type input_message: FlowMessage
         """
-
+        input_data = input_message.data
         # ~~~ Process input ~~~
         term = input_data.get("search_term", None)
         api_wrapper = WikipediaAPIWrapper(
@@ -70,4 +68,10 @@ class WikiSearchAtomicFlow(AtomicFlow):
 
         # Log the update to the flow messages list
         observation = search_response["wiki_content"] if search_response["wiki_content"] else search_response["relevant_pages"]
-        return {"wiki_content": observation}
+        
+        reply = self._package_output_message(
+            input_message = input_message,
+            response = {"wiki_content": observation},
+        )
+        
+        self.reply_to_message(reply=reply, to=input_message)

__init__.py

@@ -1,6 +1,6 @@
 # ~~~ Specify the dependencies ~~~
 dependencies = [
-    {"url": "aiflows/ChatFlowModule", "revision": "main"},
+    {"url": "aiflows/ChatFlowModule", "revision": "coflows"},
 ]
 from aiflows import flow_verse
 

demo.yaml

@@ -9,12 +9,10 @@ subflows_config:
       wiki_search:
         description: "Performs a search on Wikipedia."
         input_args: ["search_term"]
-      ddg_search:
-        description: "Query the search engine DuckDuckGo."
-        input_args: ["query"]
-      finish:
-        description: "Signal that the objective has been satisfied, and returns the answer to the user."
-        input_args: ["answer"]
+      # ddg_search:
+      #   description: "Query the search engine DuckDuckGo."
+      #   input_args: ["query"]
+  
     backend:
       _target_: aiflows.backends.llm_lite.LiteLLMBackend
       api_infos: ???

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,87 +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/ControllerExecutorFlowModule", "revision":  os.getcwd()},
+    {"url": "aiflows/ControllerExecutorFlowModule", "revision": os.getcwd()}
 ]
 
+from aiflows import flow_verse
 flow_verse.sync_dependencies(dependencies)
-
 if __name__ == "__main__":
-    # ~~~ 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 = "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")
+    #1. ~~~~~ Set up a colink server ~~~~
+    FLOW_MODULES_PATH = "./"
     
-    cl = serve_utils.start_colink_component(
-        "Reverse Number Demo",
-        {"jwt": jwt, "addr": addr}
-    )
-    # path_to_output_file = "output.jsonl"  # Uncomment this line to save the output to disk
+    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)
-    # put the API information in the config
+    
+    #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 = "azure",
+    #                           api_base = os.getenv("AZURE_API_BASE"),
+    #                           api_key = os.getenv("AZURE_OPENAI_KEY"),
+    #                           api_version =  os.getenv("AZURE_API_VERSION") )
+    
+    
+    quick_load_api_keys(cfg, api_information, key="api_infos")
+
+    
+    #3. ~~~~ Serve The Flow ~~~~
     serve_utils.recursive_serve_flow(
         cl = cl,
-        flow_type="ReAct_served",
+        flow_type="ControllerExecutorFlowModule",
         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
-    #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")
+    #4. ~~~~~Start A Worker Thread~~~~~
+    run_dispatch_worker_thread(cl, dispatch_point="coflows_dispatch", flow_modules_base_path=FLOW_MODULES_PATH)
 
-    # ~~~ Get the data ~~~
-    # This can be a list of samples
-    # 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?"}
-    data = {
-        "id": 0,
-        "goal": "Answer the following question: What is the profession and date of birth of Michael Jordan?",
-    }
-    # ~~~ Run inference ~~~
+    #5. ~~~~~Mount the flow and get its proxy~~~~~~
     proxy_flow = serve_utils.recursive_mount(
         cl=cl,
         client_id="local",
-        flow_type="ReAct_served",
-        config_overrides=cfg,
+        flow_type="ControllerExecutorFlowModule",
+        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)
     
-    future = proxy_flow.ask(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(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))
+    
+    
+    #9. ~~~~~Optional: Unserve Flow~~~~~~
+    # serve_utils.delete_served_flow(cl, "FlowModule")
+