coflows update compatible

ChatAtomicFlow.py

@@ -14,6 +14,8 @@ from aiflows.prompt_template import JinjaPrompt
 
 from aiflows.backends.llm_lite import LiteLLMBackend
 
+from aiflows.messages import FlowMessage
+
 log = logging.get_logger(__name__)
 
 
@@ -356,15 +358,13 @@ class ChatAtomicFlow(AtomicFlow):
         self._state_update_add_chat_message(role=self.flow_config["user_name"],
                                             content=user_message_content)
 
-    def run(self,input_data: Dict[str, Any]):
+    def run(self,input_message: FlowMessage):
         """ This method runs the flow. It processes the input, calls the backend and updates the state of the flow.
         
-        :param input_data: The input data of the flow.
-        :type input_data: Dict[str, Any]
-        :return: The LLM's api output.
-        :rtype: Dict[str, Any]
+        :param input_message: The input data of the flow.
+        :type input_message: aiflows.messages.FlowMessage
         """
-        
+        input_data = input_message.data
         # ~~~ Process input ~~~
         self._process_input(input_data)
 
@@ -378,4 +378,7 @@ class ChatAtomicFlow(AtomicFlow):
                 content=answer
             )
         response = response if len(response) > 1 or len(response) == 0 else response[0]
-        return {"api_output": response}
+        
+        reply_message = self._package_output_message(input_message, response = {"api_output": response})
+        
+        self.reply_to_message(reply = reply_message, to = input_message)

demo.yaml

@@ -1,56 +1,46 @@
-input_interface:  # Connector between the "input data" and the Flow
-  _target_: aiflows.interfaces.KeyInterface
-  additional_transformations:
-    - _target_: aiflows.data_transformations.KeyMatchInput  # Pass the input parameters specified by the flow
-
-output_interface:  # Connector between the Flow's output and the caller
-  _target_: aiflows.interfaces.KeyInterface
-  keys_to_rename:
-    api_output: answer  # Rename the api_output to answer
-
-flow:  # Overrides the ChatAtomicFlow config
-  _target_: flow_modules.aiflows.ChatFlowModule.ChatAtomicFlow.instantiate_from_default_config
-
-  name: "SimpleQA_Flow"
-  description: "A flow that answers questions."
-
-  # ~~~ Input interface specification ~~~
-  input_interface_non_initialized:
-    - "question"
-
-  # ~~~ backend model parameters ~~
-  backend:
-    _target_: aiflows.backends.llm_lite.LiteLLMBackend
-    api_infos: ???
-    model_name: 
-      openai: "gpt-3.5-turbo"
-      azure: "azure/gpt-4"
-    
-    # ~~~ generation_parameters ~~
-    n: 1
-    max_tokens: 3000
-    temperature: 0.3
-
-    top_p: 0.2
-    frequency_penalty: 0
-    presence_penalty: 0
-
-  n_api_retries: 6
-  wait_time_between_retries: 20
-
-  # ~~~ Prompt specification ~~~
-  system_message_prompt_template:
-    _target_: aiflows.prompt_template.JinjaPrompt
-    template: |2-
-      You are a helpful chatbot that truthfully answers questions.
-    input_variables: []
-    partial_variables: {}
- 
-
-  init_human_message_prompt_template:
-    _target_: aiflows.prompt_template.JinjaPrompt
-    template: |2-
-      Answer the following question: {{question}}
-    input_variables: ["question"]
-    partial_variables: {}
+
+_target_: flow_modules.aiflows.ChatFlowModule.ChatAtomicFlow.instantiate_from_default_config
+
+name: "SimpleQA_Flow"
+description: "A flow that answers questions."
+
+# ~~~ Input interface specification ~~~
+input_interface_non_initialized:
+  - "question"
+
+# ~~~ backend model parameters ~~
+backend:
+  _target_: aiflows.backends.llm_lite.LiteLLMBackend
+  api_infos: ???
+  model_name: 
+    openai: "gpt-3.5-turbo"
+    azure: "azure/gpt-4"
+  
+  # ~~~ generation_parameters ~~
+  n: 1
+  max_tokens: 3000
+  temperature: 0.3
+
+  top_p: 0.2
+  frequency_penalty: 0
+  presence_penalty: 0
+
+n_api_retries: 6
+wait_time_between_retries: 20
+
+# ~~~ Prompt specification ~~~
+system_message_prompt_template:
+  _target_: aiflows.prompt_template.JinjaPrompt
+  template: |2-
+    You are a helpful chatbot that truthfully answers questions.
+  input_variables: []
+  partial_variables: {}
+
+
+init_human_message_prompt_template:
+  _target_: aiflows.prompt_template.JinjaPrompt
+  template: |2-
+    Answer the following question: {{question}}
+  input_variables: ["question"]
+  partial_variables: {}
 

run.py

@@ -5,11 +5,18 @@ 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
 
+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
 
@@ -22,52 +29,86 @@ 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") )
+    
+    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.serve_flow(
+        cl = cl,
+        flow_type="ChatFlowModule",
+        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)
 
-    cfg["flow"]["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)
-        ),
-    }
+    #5. ~~~~~Mount the flow and get its proxy~~~~~~
+    proxy_flow = serve_utils.recursive_mount(
+        cl=cl,
+        client_id="local",
+        flow_type="ChatFlowModule",
+        config_overrides=None,
+        initial_state=None,
+        dispatch_point_override=None,
+    )   
+    
 
-    # ~~~ Get the data ~~~
+    #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 ~~~
-    path_to_output_file = None
-    # path_to_output_file = "output.jsonl"  # Uncomment this line to save the output to disk
-
-    _, outputs = FlowLauncher.launch(
-        flow_with_interfaces=flow_with_interfaces,
+    
+    #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, "ReverseNumberAtomicFlow_served")
+    
+