coflow compatible

ChatHumanFlowModule.py

@@ -1,10 +1,12 @@
-from aiflows.base_flows import CircularFlow
+from aiflows.base_flows import CompositeFlow
 from aiflows.utils import logging
-
+from aiflows.messages import FlowMessage
+from aiflows.interfaces import KeyInterface
+from aiflows.data_transformations import RegexFirstOccurrenceExtractor,EndOfInteraction
 log = logging.get_logger(f"aiflows.{__name__}")
 
 
-class ChatHumanFlowModule(CircularFlow):
+class ChatHumanFlowModule(CompositeFlow):
     """ This class implements a Chat Human Flow Module. It is a flow that consists of two sub-flows that are executed circularly. It Contains the following subflows:
     
     - A User Flow: A flow makes queries to the Assistant Flow. E.g. The user asks the assistant (LLM) a question.
@@ -48,8 +50,104 @@ class ChatHumanFlowModule(CircularFlow):
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
+        
+        
+        self.regex_extractor = RegexFirstOccurrenceExtractor(**self.flow_config["regex_first_occurrence_extractor"])
+        
+        
+        self.end_of_interaction = EndOfInteraction(**self.flow_config["end_of_interaction"])
+        
+        self.input_interface_assistant = KeyInterface(
+            keys_to_rename = {"human_input": "query"},
+            additional_transformations = [self.regex_extractor, self.end_of_interaction]
+        )
+        
+    def set_up_flow_state(self):
+        """ This method sets up the flow state. It is called when the flow is executed."""
+        super().set_up_flow_state()
+        self.flow_state["last_flow_called"] = None
+        self.flow_state["current_round"] = 0
+        self.flow_state["user_inputs"] = []
+        self.flow_state["assistant_outputs"] = []
+        self.flow_state["input_message"] = None
+        self.flow_state["end_of_interaction"] = False
 
     @classmethod
     def type(cls):
         """ This method returns the type of the flow."""
         return "OpenAIChatHumanFlowModule"
+    
+    def max_rounds_reached(self):
+        return self.flow_config["max_rounds"] is not None and self.flow_state["current_round"] >= self.flow_config["max_rounds"]
+    
+    def generate_reply(self):
+        
+        reply = self._package_output_message(
+            input_message = self.flow_state["input_message"],
+            response = {
+                "user_inputs": self.flow_state["user_inputs"],
+                "assistant_outputs":  self.flow_state["assistant_outputs"],
+                "end_of_interaction": self.flow_state["end_of_interaction"]
+            },
+        )
+        
+        self.reply_to_message(
+            reply = reply,
+            to = self.flow_state["input_message"]
+        )
+        
+            
+    
+    def call_to_user(self,input_message):
+        
+        self.flow_state["assistant_outputs"].append(input_message.data["api_output"])
+
+        if self.max_rounds_reached():
+            self.generate_reply()
+        else:
+            self.subflows["User"].send_message_async(input_message,pipe_to=self.flow_config["flow_ref"])
+            self.flow_state["last_flow_called"] = "User"
+        
+        self.flow_state["current_round"] += 1
+        
+        
+            
+    def call_to_assistant(self,input_message):
+        message = self.input_interface_assistant(input_message)
+        
+        if  self.flow_state["last_flow_called"] is None:
+            self.flow_state["input_message"] = input_message
+        
+        else:
+            self.flow_state["user_inputs"].append(input_message.data["query"])
+        
+        if message.data["end_of_interaction"]:
+            self.flow_state["end_of_interaction"] = True
+            self.generate_reply()
+            
+        else:
+            self.subflows["Assistant"].send_message_async(message,pipe_to=self.flow_config["flow_ref"])
+            self.flow_state["last_flow_called"] = "Assistant"    
+    
+    def run(self,input_message: FlowMessage):
+        """ This method runs the flow. It is the main method of the flow and it is called when the flow is executed.
+        
+        :param input_message: The input message to the flow.
+        :type input_message: FlowMessage
+        """
+        last_flow_called = self.flow_state["last_flow_called"]
+        
+        if last_flow_called is None or last_flow_called == "User":
+            self.call_to_assistant(input_message=input_message)
+            
+            
+        elif last_flow_called == "Assistant":
+            self.call_to_user(input_message=input_message)
+            
+            
+            
+            
+            
+              
+        
+        

ChatHumanFlowModule.yaml

@@ -4,11 +4,24 @@ description: "Flow that enables chatting between a ChatAtomicFlow and a user pro
 max_rounds: null # Run until early exit is detected
 
 input_interface:
-  _target_: aiflows.interfaces.KeyInterface
+  - "query"
   
 output_interface:
   - "end_of_interaction"
-  - "answer"
+  - "user_inputs"
+  - "assistant_outputs"
+
+regex_first_occurrence_extractor:
+  regex: "(?<=```answer)([\\s\\S]*?)(?=```)"
+  input_key: "query"
+  output_key: "answer"
+  strip: true
+  assert_unique: true
+
+end_of_interaction:
+  end_of_interaction_string: "<END>"
+  input_key: "query"
+  output_key: "end_of_interaction"
 
 subflows_config:
   Assistant:
@@ -16,52 +29,5 @@ subflows_config:
   User:
     _target_: flow_modules.aiflows.HumanStandardInputFlowModule.HumanStandardInputFlow.instantiate_from_default_config
 
-topology:
-  - goal: "Query the assistant"
-
-    ### Input Interface
-    input_interface:
-      _target_: aiflows.interfaces.KeyInterface
-      additional_transformations:
-        - _target_: aiflows.data_transformations.KeyMatchInput
-
-    ### Flow Specification
-    flow: Assistant
-    reset_every_round: false
-
-    ### Output Interface
-    output_interface:
-      _target_: aiflows.interfaces.KeyInterface
-      additional_transformations:
-        - _target_: aiflows.data_transformations.PrintPreviousMessages
-
-  - goal: "Ask the user for input"
-
-    ### Input Interface
-    input_interface:
-      _target_: aiflows.interfaces.KeyInterface
-      additional_transformations:
-        - _target_: aiflows.data_transformations.KeyMatchInput
-
-    ### Flow Specification
-    flow: User
-    reset_every_round: true
 
-    ### Output Interface
-    output_interface:
-      _target_: aiflows.interfaces.KeyInterface
-      keys_to_rename:
-        human_input: query
-      additional_transformations:
-        - _target_: aiflows.data_transformations.RegexFirstOccurrenceExtractor
-          regex: '(?<=```answer)([\s\S]*?)(?=```)'
-          input_key: "query"
-          output_key: "answer"
-          strip: True
-          assert_unique: True
-        - _target_: aiflows.data_transformations.EndOfInteraction
-          end_of_interaction_string: "<END>"
-          input_key: "query"
-          output_key: "end_of_interaction"
 
-early_exit_key: "end_of_interaction"

__init__.py

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

demo.yaml

@@ -1,31 +1,29 @@
 
-
-flow:
-  max_rounds: 5
-  _target_: flow_modules.aiflows.ChatInteractiveFlowModule.ChatHumanFlowModule.instantiate_from_default_config
-  subflows_config:
-    Assistant:
-      _target_: flow_modules.aiflows.ChatFlowModule.ChatAtomicFlow.instantiate_from_default_config
-      backend:
-        _target_: aiflows.backends.llm_lite.LiteLLMBackend
-        api_infos: ???
-        model_name: 
-          openai: "gpt-4"
-          azure: "azure/gpt-4"
-      input_interface_non_initialized: []
-      
-    User:
-      _target_: flow_modules.aiflows.HumanStandardInputFlowModule.HumanStandardInputFlow.instantiate_from_default_config
-      request_multi_line_input_flag: False
-      query_message_prompt_template: 
-        _target_: aiflows.prompt_template.JinjaPrompt
-        template: |2-
-          {{api_output}}
-          
-          To end an Interaction, type <END> and press enter.
-          
-        input_variables: ["api_output"]
-      input_interface:
-        - "api_output"
-      output_interface:
-        - "human_input"
+max_rounds: 2
+_target_: flow_modules.aiflows.ChatInteractiveFlowModule.ChatHumanFlowModule.instantiate_from_default_config
+subflows_config:
+  Assistant:
+    _target_: flow_modules.aiflows.ChatFlowModule.ChatAtomicFlow.instantiate_from_default_config
+    backend:
+      _target_: aiflows.backends.llm_lite.LiteLLMBackend
+      api_infos: ???
+      model_name: 
+        openai: "gpt-4"
+        azure: "azure/gpt-4"
+    input_interface_non_initialized: []
+    
+  User:
+    _target_: flow_modules.aiflows.HumanStandardInputFlowModule.HumanStandardInputFlow.instantiate_from_default_config
+    request_multi_line_input_flag: False
+    query_message_prompt_template: 
+      _target_: aiflows.prompt_template.JinjaPrompt
+      template: |2-
+        {{api_output}}
+        
+        To end an Interaction, type <END> and press enter.
+        
+      input_variables: ["api_output"]
+    input_interface:
+      - "api_output"
+    output_interface:
+      - "human_input"

run.py

@@ -1,5 +1,3 @@
-"""A simple script to run a Flow that can be used for development and debugging."""
-
 import os
 
 import hydra
@@ -7,71 +5,110 @@ import hydra
 import aiflows
 from aiflows.flow_launchers import FlowLauncher
 from aiflows.backends.api_info import ApiInfo
-from aiflows.utils.general_helpers import read_yaml_file
+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
 
-CACHING_PARAMETERS.do_caching = False  # Set to True to enable caching
+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()
 
+
 dependencies = [
-    {"url": "aiflows/ChatInteractiveFlowModule", "revision": os.getcwd()},
+    {"url": "aiflows/ChatInteractiveFlowModule", "revision": os.getcwd()}
 ]
+
 from aiflows import flow_verse
 flow_verse.sync_dependencies(dependencies)
-
 if __name__ == "__main__":
-    # ~~~ Set the API information ~~~
-    # OpenAI backend
+    
+    #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 = "azure",
     #                           api_base = os.getenv("AZURE_API_BASE"),
     #                           api_key = os.getenv("AZURE_OPENAI_KEY"),
     #                           api_version =  os.getenv("AZURE_API_VERSION") )
-
-    root_dir = "."
-    cfg_path = os.path.join(root_dir, "demo.yaml")
-    cfg = read_yaml_file(cfg_path)
     
-    cfg["flow"]["subflows_config"]["Assistant"]["backend"]["api_infos"] = api_information
-    
-    # ~~~ Instantiate the Flow ~~~
-    flow_with_interfaces = {
-        "flow": hydra.utils.instantiate(cfg['flow'], _recursive_=False, _convert_="partial"),
-        "input_interface": (
-            None
-            if cfg.get( "input_interface", None) is None
-            else hydra.utils.instantiate(cfg['input_interface'], _recursive_=False)
-        ),
-        "output_interface": (
-            None
-            if cfg.get( "output_interface", None) is None
-            else hydra.utils.instantiate(cfg['output_interface'], _recursive_=False)
-        ),
-    }
-    # ~~~ Get the data ~~~
-    data = {"id": 0}  # This can be a list of samples
-    # data = {"id": 0, "question": "Who was the NBA champion in 2023?"}  # This can be a list of samples
+    
+    quick_load_api_keys(cfg, api_information, key="api_infos")
 
-    # ~~~ Run inference ~~~
-    path_to_output_file = None
-    # path_to_output_file = "output.jsonl"  # Uncomment this line to save the output to disk
+    
+    #3. ~~~~ Serve The Flow ~~~~
+    serve_utils.recursive_serve_flow(
+        cl = cl,
+        flow_type="ChatInteractiveFlowModule",
+        default_config=cfg,
+        default_state=None,
+        default_dispatch_point="coflows_dispatch"
+    )
+    
+    #4. ~~~~~Start A Worker Thread~~~~~
+    run_dispatch_worker_thread(cl, dispatch_point="coflows_dispatch", flow_modules_base_path=FLOW_MODULES_PATH)
 
-    _, outputs = FlowLauncher.launch(
-        flow_with_interfaces=flow_with_interfaces,
+    #5. ~~~~~Mount the flow and get its proxy~~~~~~
+    proxy_flow = serve_utils.recursive_mount(
+        cl=cl,
+        client_id="local",
+        flow_type="ChatInteractiveFlowModule",
+        config_overrides=None,
+        initial_state=None,
+        dispatch_point_override=None,
+    )   
+    
+    #6. ~~~ Get the data ~~~
+    data = {"id": 0, "query": "I want to ask you a few questions"}  # This can be a list of samples
+    # data = {"id": 0, "question": "Who was the NBA champion in 2023?"}  # This can be a list of samples
+    
+    #option1: use the FlowMessage class
+    input_message = FlowMessage(
         data=data,
-        path_to_output_file=path_to_output_file
     )
 
+    #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 ~~~
-    flow_output_data = outputs[0]
-    print(flow_output_data)
+    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, "ChatWithDemonstrationFlowModule")    o_caching = False  # Set to True to enable caching