Clean Versions

ChatWithDemonstrationsFlow.py

@@ -4,7 +4,7 @@ from aiflows.base_flows import CompositeFlow
 from aiflows.utils import logging
 from aiflows.interfaces import KeyInterface
 logging.set_verbosity_debug()
-
+from aiflows.messages import FlowMessage
 log = logging.get_logger(__name__)
 
 
@@ -50,11 +50,16 @@ class ChatWithDemonstrationsFlow(CompositeFlow):
         )
 
     def set_up_flow_state(self):
+        """ Set up the flow state. It sets the last state of the flow to None.
+        """
         super().set_up_flow_state()
-        self.flow_state["last_flow_called"] = None
+        self.flow_state["last_state"] = None
             
     def run(self,input_message):
-        
+        """ Runs the flow.
+        :param input_message: The input message of the flow. 
+        :type input_message: FlowMessage
+        """
         # #~~~~~~~~~~~Solution 1 - Blocking ~~~~~~~
         # future = self.subflows["demonstration_flow"].get_reply_future(input_message)
         
@@ -65,23 +70,23 @@ class ChatWithDemonstrationsFlow(CompositeFlow):
         # self.send_message(reply)
         
         # #~~~~~~~~~~~Solution 2 - Non-Blocking ~~~~~~~
-        if self.flow_state["last_flow_called"] is None:
+        if self.flow_state["last_state"] is None:
             self.flow_state["input_message"] = input_message
             self.subflows["demonstration_flow"].get_reply(
                 input_message,
                 self.get_instance_id(),
             )
-            self.flow_state["last_flow_called"] = "demonstration_flow"
+            self.flow_state["last_state"] = "demonstration_flow"
         
-        elif self.flow_state["last_flow_called"] == "demonstration_flow":
+        elif self.flow_state["last_state"] == "demonstration_flow":
             self.subflows["chat_flow"].get_reply(
                 input_message,
                 self.get_instance_id(),
             )
-            self.flow_state["last_flow_called"] = "chat_flow"
+            self.flow_state["last_state"] = "chat_flow"
 
         else:
-            self.flow_state["last_flow_called"] = None
+            self.flow_state["last_state"] = None
             
             reply = self.package_output_message(
                 self.flow_state["input_message"],

ChatWithDemonstrationsFlow.yaml

@@ -5,9 +5,15 @@ description: "A sequential flow that answers questions with demonstrations"
 subflows_config:
   demonstration_flow:
     _target_: flow_modules.aiflows.ChatWithDemonstrationsFlowModule.DemonstrationsAtomicFlow.instantiate_from_default_config
+    flow_class_name: flow_modules.aiflows.ChatWithDemonstrationsFlowModule.DemonstrationsAtomicFlow
+    flow_endpoint: DemonstrationsAtomicFlow
+    user_id: local
     name: "DemonstrationsAtomicFlow"
     description: "A flow that answers questions with demonstrations"
   chat_flow:
     _target_: flow_modules.aiflows.ChatFlowModule.ChatAtomicFlow.instantiate_from_default_config
     name: "ChatAtomicFlow"
     description: "A flow that answers questions"
+    flow_class_name: flow_modules.aiflows.ChatFlowModule.ChatAtomicFlow
+    flow_endpoint: ChatAtomicFlow
+    user_id: local

DemonstrationsAtomicFlow.py

@@ -5,6 +5,7 @@ from aiflows.utils import general_helpers
 from typing import Dict,Any,Optional,List
 from aiflows.prompt_template import JinjaPrompt
 from copy import deepcopy
+from aiflows.messages import FlowMessage
 import os
 import hydra
 log = logging.get_logger(__name__)
@@ -41,7 +42,7 @@ class DemonstrationsAtomicFlow(AtomicFlow):
         
     *Output Interface*:
 
-    - The input interface expected by its successor flow (e.g. typically ChatAtomicFlow so the input interface expected by ChatAtomicFlow))
+    - Whichever data that was passed in the input_message (e.g. typically ChatAtomicFlow so the input interface expected by ChatAtomicFlow))
     - `demonstrations` (List[Dict[str, Any]]): A list of demonstrations. Each demonstration is a dictionary with the following keys:
                                                     - idx (int): The index of the demonstration
                                                     - query (str): The query of the demonstration
@@ -182,13 +183,11 @@ class DemonstrationsAtomicFlow(AtomicFlow):
         log.info("Loaded the demonstrations for %d datapoints from %s", len(self.data), self.params["data_dir"])
 
     def run(self,
-            input_message):
-        """ This method runs the flow. It returns the input data of the flow with the demonstrations added to it.
+            input_message: FlowMessage):
+        """ This method runs the flow. It returns the data of the input_message with the demonstrations added to it.
         
-        :param input_data: The input data of the flow.
-        :type input_data: Dict[str, Any]
-        :return: The input data of the flow with the demonstrations added to it.
-        :rtype: Dict[str, Any]
+        :param input_message: The input message of the flow.
+        :type input_message: FlowMessage
         """
         input_data = input_message.data
         reply = self.package_output_message(

DemonstrationsAtomicFlow.yaml

@@ -1,6 +1,6 @@
 name: "DemonstrationAtomicFlow"
 description: "A flow which returns Demonstrations"
-
+_target_: flow_modules.aiflows.ChatWithDemonstrationsFlowModule.DemonstrationsAtomicFlow.instantiate_from_default_config
 data: ??? #e.g. [{"query_data": {"query": "What is the capital of France?"}, "response_data": {"response": "Paris, my sir."}}]
 params:
   data_dir: ???

demo.yaml

@@ -9,7 +9,11 @@ output_interface: # Connector between the Flow's output and the caller
 
 subflows_config:
   demonstration_flow:
-    _target_: flow_modules.aiflows.ChatWithDemonstrationsFlowModule.DemonstrationsAtomicFlow.instantiate_from_default_config
+    _target_:  flow_modules.aiflows.ChatWithDemonstrationsFlowModule.DemonstrationsAtomicFlow.instantiate_from_default_config
+    name: "proxy of DemonstrationsAtomicFlow"
+    flow_endpoint: DemonstrationsAtomicFlow
+    user_id: local
+    description: "A flow that answers questions with demonstrations"
     data: 
       - query_data:
           query: "What is the capital of Turkey?"
@@ -32,7 +36,10 @@ subflows_config:
 
   chat_flow:
     _target_: flow_modules.aiflows.ChatFlowModule.ChatAtomicFlow.instantiate_from_default_config
-    name: "SimpleQA_Flow"
+    name: "proxy SimpleQA_Flow"
+    description: "A flow that answers questions"
+    flow_endpoint: ChatAtomicFlow
+    user_id: local
     # ~~~ Input interface specification ~~~
     input_interface_non_initialized:
       - "question"

run.py

@@ -41,17 +41,6 @@ if __name__ == "__main__":
     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") )
-    
-
     #2.1 ~~~ Set the API information ~~~
     # OpenAI backend
     api_information = [ApiInfo(backend_used="openai",
@@ -66,26 +55,25 @@ if __name__ == "__main__":
 
     
     #3. ~~~~ Serve The Flow ~~~~
+    
     serve_utils.recursive_serve_flow(
         cl = cl,
-        flow_type="ChatWithDemonstrationFlowModule",
-        default_config=cfg,
-        default_state=None,
-        default_dispatch_point="coflows_dispatch"
+        flow_class_name="flow_modules.aiflows.ChatWithDemonstrationsFlowModule.ChatWithDemonstrationsFlow",
+        flow_endpoint="ChatWithDemonstrationsFlow",
     )
+
     
     #4. ~~~~~Start A Worker Thread~~~~~
-    run_dispatch_worker_thread(cl, dispatch_point="coflows_dispatch", flow_modules_base_path=FLOW_MODULES_PATH)
+    run_dispatch_worker_thread(cl)
 
     #5. ~~~~~Mount the flow and get its proxy~~~~~~
-    proxy_flow = serve_utils.recursive_mount(
+    proxy_flow= serve_utils.get_flow_instance(
         cl=cl,
-        client_id="local",
-        flow_type="ChatWithDemonstrationFlowModule",
-        config_overrides=None,
-        initial_state=None,
-        dispatch_point_override=None,
-    )   
+        flow_endpoint="ChatWithDemonstrationsFlow",
+        user_id="local",
+        config_overrides = cfg
+    )
+    
     
     #6. ~~~ Get the data ~~~
     data = {"id": 0, "question": "What is the capital of France?"}  # This can be a list of samples