merge with coflows

ChatWithDemonstrationsFlow.py

@@ -1,14 +1,14 @@
 
 
-from aiflows.base_flows import SequentialFlow
+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__)
 
 
-class ChatWithDemonstrationsFlow(SequentialFlow):
+class ChatWithDemonstrationsFlow(CompositeFlow):
     """ A Chat with Demonstrations Flow. It is a flow that consists of multiple sub-flows that are executed sequentially.
     It's parent class is SequentialFlow.
     
@@ -44,4 +44,68 @@ class ChatWithDemonstrationsFlow(SequentialFlow):
     :param \**kwargs: Arguments to be passed to the parent class SequentialFlow constructor.
     """
     def __init__(self,**kwargs):
-        super().__init__(**kwargs)
+        super().__init__(**kwargs)
+        self.output_interface = KeyInterface(
+            keys_to_rename={"api_output": "answer"}
+        )
+
+    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_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)
+        
+        # answer = self.subflows["chat_flow"].get_reply_future(future.get_message())
+        
+        #reply =  self.package_output_message(input_message, self.output_interface(answer.get_data()))
+        
+        # self.send_message(reply)
+        
+        # #~~~~~~~~~~~Solution 2 - Non-Blocking ~~~~~~~
+        if self.flow_state["last_state"] is None:
+            
+            self.flow_state["initial_message"] = input_message
+            
+            msg = self.package_input_message(
+                input_message.data
+            )
+            
+            self.subflows["demonstration_flow"].get_reply(
+                msg,
+            )
+            self.flow_state["last_state"] = "demonstration_flow"
+        
+        elif self.flow_state["last_state"] == "demonstration_flow":
+            
+            msg = self.package_input_message(
+                input_message.data
+            )
+            
+            self.subflows["chat_flow"].get_reply(
+                msg,
+            )
+
+            self.flow_state["last_state"] = "chat_flow"
+
+        else:
+            self.flow_state["last_state"] = None
+
+            reply = self.package_output_message(
+                self.flow_state["initial_message"],
+                response = self.output_interface(input_message).data
+            )
+                
+            self.send_message(
+                reply
+            )
+
+        
+       

ChatWithDemonstrationsFlow.yaml

@@ -1,33 +1,17 @@
+
 name: "ChatAtomic_Flow_with_Demonstrations"
 description: "A sequential flow that answers questions with demonstrations"
-
+_target_: flow_modules.aiflows.ChatWithDemonstrationsFlowModule.ChatWithDemonstrationsFlow.instantiate_from_default_config
 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
-
-topology:
-  - goal: Get Demonstrations
-    input_interface:
-      _target_: aiflows.interfaces.KeyInterface
-    flow: demonstration_flow
-    output_interface:
-      _target_: aiflows.interfaces.KeyInterface
-
-  - goal: Answer the question
-    input_interface:
-      _target_: aiflows.interfaces.KeyInterface
-    flow: chat_flow
-    output_interface:
-      _target_: aiflows.interfaces.KeyInterface
-      keys_to_rename:
-        api_output: answer  # Rename the api_output to answer
-
-
-input_interface_initialized:
-  - "query"
-
-output_interface:
-  - "answer"
+    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,12 +183,15 @@ 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]:
-        """ 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
         """
-        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.send_message(reply)

DemonstrationsAtomicFlow.yaml

@@ -1,9 +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: ???

__init__.py

@@ -1,8 +1,11 @@
 # ~~~ Specify the dependencies ~~
+
+
 dependencies = [
     {"url": "aiflows/ChatFlowModule", "revision": "main"},
 ]
 from aiflows import flow_verse
 flow_verse.sync_dependencies(dependencies)
+
 from .ChatWithDemonstrationsFlow import ChatWithDemonstrationsFlow
-from .DemonstrationsAtomicFlow import DemonstrationsAtomicFlow
+from .DemonstrationsAtomicFlow import DemonstrationsAtomicFlow

demo.yaml

@@ -1,76 +1,74 @@
-flow:  # Overrides the OpenAIChatAtomicFlow config
-  _target_: flow_modules.aiflows.ChatWithDemonstrationsFlowModule.ChatWithDemonstrationsFlow.instantiate_from_default_config
-  name: "SimpleQA_Flow_with_Demonstrations"
-  description: "A sequential flow that answers questions with demonstrations"
-  
-  input_interface:  # Connector between the "input data" and the Flow
-    - "questions"
-  output_interface:  # Connector between the Flow's output and the caller
-    - "answer"
-  
-  subflows_config:
-    demonstration_flow:
-      data: 
-        - query_data:
-            query: "What is the capital of Turkey?"
-          response_data:
-            response: "Istambul, my sir."
-        - query_data:
-            query: "what is the capital of Germany?"
-          response_data:
-            response: "Berlin, my sir."
-      params:
-        data_dir: null
-        demonstrations_id: my_sir_demo
-      query_prompt_template:
-        template: |2-
-          Answer the following question: {{query}}
-        input_variables:
-          - "query"
-      response_prompt_template:
-        template: |2-
-          {{response}}
-        input_variables:
-          - response
+_target_: flow_modules.aiflows.ChatWithDemonstrationsFlowModule.ChatWithDemonstrationsFlow.instantiate_from_default_config
+name: "SimpleQA_Flow_with_Demonstrations"
+description: "A sequential flow that answers questions with demonstrations"
 
-    chat_flow:
-      name: "SimpleQA_Flow"
-      # ~~~ Input interface specification ~~~
-      input_interface_non_initialized:
-        - "question"
+input_interface: # Connector between the "input data" and the Flow
+  - "questions"
+output_interface: # Connector between the Flow's output and the caller
+  - "answer"
 
-      # ~~~ 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
+subflows_config:
+  demonstration_flow:
+    name: "proxy of DemonstrationsAtomicFlow"
+    description: "A flow that answers questions with demonstrations"
+    data: 
+      - query_data:
+          query: "What is the capital of Turkey?"
+        response_data:
+          response: "Istambul, my sir."
+      - query_data:
+          query: "what is the capital of Germany?"
+        response_data:
+          response: "Berlin, my sir."
+    query_prompt_template:
+      template: |2-
+        Answer the following question: {{query}}
+      input_variables:
+        - "query"
+    response_prompt_template:
+      template: |2-
+        {{response}}
+      input_variables:
+        - response
 
-        top_p: 0.2
-        frequency_penalty: 0
-        presence_penalty: 0
+  chat_flow:
+    name: "proxy SimpleQA_Flow"
+    description: "A flow that answers questions"
+    # ~~~ Input interface specification ~~~
+    input_interface_non_initialized:
+      - "question"
 
-      n_api_retries: 6
-      wait_time_between_retries: 20
+    # ~~~ backend model parameters ~~
+    backend:
+      _target_: aiflows.backends.llm_lite.LiteLLMBackend
+      api_infos: ???
+      model_name:
+        openai: "gpt-3.5-turbo"
+        azure: "azure/gpt-4"
 
-      # ~~~ Prompt specification ~~~
-      system_message_prompt_template:
-        _target_: aiflows.prompt_template.JinjaPrompt
-        template: |2-
-          You are a helpful chatbot that truthfully answers questions. Answer in a similar way to your previous replies.
-        input_variables: []
-        partial_variables: {}
- 
+      # ~~~ generation_parameters ~~
+      n: 1
+      max_tokens: 3000
+      temperature: 0.3
 
-      init_human_message_prompt_template:
-        _target_: aiflows.prompt_template.JinjaPrompt
-        template: |2-
-          Answer the following question: {{question}}
-        input_variables: ["question"]
-        partial_variables: {}
+      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. Answer in a similar way to your previous replies.
+      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

@@ -4,72 +4,102 @@ import hydra
 
 import aiflows
 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
 
-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()},
+    {"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
+    
+    #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.recursive_serve_flow(
+        cl = cl,
+        flow_class_name="flow_modules.aiflows.ChatWithDemonstrationsFlowModule.ChatWithDemonstrationsFlow",
+        flow_endpoint="ChatWithDemonstrationsFlow",
+    )
 
-    cfg["flow"]["subflows_config"]["chat_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)
-        ),
-    }
-    # ~~~ Get the data ~~~
-    data = {"id": 0, "question": "What's 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
+    
+    #4. ~~~~~Start A Worker Thread~~~~~
+    run_dispatch_worker_thread(cl)
 
-    _, outputs = FlowLauncher.launch(
-        flow_with_interfaces=flow_with_interfaces,
-        data=data,
-        path_to_output_file=path_to_output_file
+    #5. ~~~~~Mount the flow and get its proxy~~~~~~
+    proxy_flow= serve_utils.get_flow_instance(
+        cl=cl,
+        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
+    # data = {"id": 0, "question": "Who was the NBA champion in 2023?"}  # This can be a list of samples
+    
 
+    input_message = proxy_flow.package_input_message(data = data)
+    
+    #7. ~~~ Run inference ~~~
+    future = proxy_flow.get_reply_future(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")