compatible to new version

ChatWithDemonstrationsFlow.py

@@ -53,26 +53,30 @@ class ChatWithDemonstrationsFlow(CompositeFlow):
         super().set_up_flow_state()
         self.flow_state["last_flow_called"] = None
             
-    def run(self,input_data):
+    def run(self,input_message):
         
-        #~~~~~~~~~~~Solution 1 - Blocking ~~~~~~~
-        future = self.ask_subflow("demonstration_flow",input_data)
+        # #~~~~~~~~~~~Solution 1 - Blocking ~~~~~~~
+        # future = self.subflows["demonstration_flow"].send_message_blocking(input_message)
         
-        answer = self.ask_subflow("chat_flow",future.get_data())
-        return self.output_interface(answer.get_data())
+        # answer = self.subflows["chat_flow"].send_message_blocking(future.get_message())
+        # self.reply_to_message(reply =self.output_interface(answer.get_message()), to=input_message)
         
         # #~~~~~~~~~~~Solution 2 - Non-Blocking ~~~~~~~
-        # if self.flow_state["last_flow_called"] is None:
-        #     self.ask_pipe_subflow("demonstration_flow",input_data)
-        #     self.flow_state["last_flow_called"] = "demonstration_flow"
-        #     return {"answer": "Just Called the Demonstration Flow"}
+        if self.flow_state["last_flow_called"] is None:
+            self.flow_state["input_message"] = input_message
+            self.subflows["demonstration_flow"].send_message_async(input_message,pipe_to=self.flow_config["flow_ref"])
+            self.flow_state["last_flow_called"] = "demonstration_flow"
         
-        # elif self.flow_state["last_flow_called"] == "demonstration_flow":
-        #     self.ask_pipe_subflow("chat_flow",input_data)
-        #     self.flow_state["last_flow_called"] = "chat_flow"
-        #     return {"answer": "Just Called the Demonstration Flow"}
-        
-        # self.flow_state["last_flow_called"] = None
-        return self.output_interface(input_data)
+        elif self.flow_state["last_flow_called"] == "demonstration_flow":
+            self.subflows["chat_flow"].send_message_async(input_message,pipe_to=self.flow_config["flow_ref"])
+            self.flow_state["last_flow_called"] = "chat_flow"
+
+        else:
+            self.flow_state["last_flow_called"] = None
+            self.reply_to_message(
+                reply =self.output_interface(input_message),
+                to=self.flow_state["input_message"]
+            )
+
         
        

DemonstrationsAtomicFlow.py

@@ -182,7 +182,7 @@ class DemonstrationsAtomicFlow(AtomicFlow):
         log.info("Loaded the demonstrations for %d datapoints from %s", len(self.data), self.params["data_dir"])
 
     def run(self,
-            input_data: Dict[str, Any]) -> Dict[str, Any]:
+            input_message):
         """ This method runs the flow. It returns the input data of the flow with the demonstrations added to it.
         
         :param input_data: The input data of the flow.
@@ -190,4 +190,9 @@ class DemonstrationsAtomicFlow(AtomicFlow):
         :return: The input data of the flow with the demonstrations added to it.
         :rtype: Dict[str, Any]
         """
-        return {**input_data,**{"demonstrations": self._get_io_pairs(input_data=input_data)}}
+        input_data = input_message.data
+        reply = self._package_output_message(
+            input_message=input_message,
+            response = {**{"demonstrations": self._get_io_pairs(input_data=input_data)},**input_data}
+        )
+        self.reply_to_message(reply = reply, to=input_message)

__init__.py

@@ -2,7 +2,7 @@
 
 
 dependencies = [
-    {"url": "aiflows/ChatFlowModule", "revision": "main"},
+    {"url": "aiflows/ChatFlowModule", "revision": "coflows"},
 ]
 from aiflows import flow_verse
 flow_verse.sync_dependencies(dependencies)

run.py

@@ -9,82 +9,115 @@ 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()  # Uncomment this line to see verbose logs
-
-from aiflows import flow_verse
+logging.set_verbosity_debug()
 
 
 dependencies = [
     {"url": "aiflows/ChatWithDemonstrationsFlowModule", "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"))]
+    
+    #1. ~~~~~ Set up a colink server ~~~~
+    FLOW_MODULES_PATH = "./"
+    
+    cl = start_colink_server()
 
 
-    FLOW_MODULES_PATH = "./"
+    #2. ~~~~~Load flow config~~~~~~
+    root_dir = "."
+    cfg_path = os.path.join(root_dir, "demo.yaml")
+    cfg = read_yaml_file(cfg_path)
     
-    jwt = os.getenv("COLINK_JWT")
-    addr = os.getenv("LOCAL_COLINK_ADDRESS")
+    #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") )
     
-    cl = serve_utils.start_colink_component(
-        "Reverse Number Demo",
-        {"jwt": jwt, "addr": addr}
-    )
 
+    #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")
 
-    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="simpleDemonstrationQA",
+        flow_type="ChatWithDemonstrationFlowModule",
         default_config=cfg,
         default_state=None,
-        default_dispatch_point="coflows_dispatch",
+        default_dispatch_point="coflows_dispatch"
     )
     
-    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)
+
+    #5. ~~~~~Mount the flow and get its proxy~~~~~~
     proxy_flow = serve_utils.recursive_mount(
         cl=cl,
         client_id="local",
-        flow_type="simpleDemonstrationQA",
-        config_overrides=cfg,
+        flow_type="ChatWithDemonstrationFlowModule",
+        config_overrides=None,
         initial_state=None,
         dispatch_point_override=None,
     )   
     
-    
-    # ~~~ Get the data ~~~
-    data = {"id": 0, "question": "What's the capital of France?"}  # This can be a list of samples
+    #6. ~~~ Get the data ~~~
+    data = {"id": 0, "question": "What is the capital of France?"}  # 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
-    # ~~~ Run inference ~~~
-   
+    
+    #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, "ChatWithDemonstrationFlowModule")