coflows - v1 (issues w/ state management here)

ChromaDBFlow.py

@@ -6,7 +6,7 @@ from copy import deepcopy
 from langchain.embeddings import OpenAIEmbeddings
 
 from chromadb import Client as ChromaClient
-
+from aiflows.messages import FlowMessage
 from aiflows.base_flows import AtomicFlow
 
 import hydra
@@ -96,14 +96,13 @@ class ChromaDBFlow(AtomicFlow):
         """
         return self.flow_config["output_keys"]
 
-    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+    def run(self, input_message: FlowMessage):
         """ This method runs the flow. It runs the ChromaDBFlow. It either writes or reads memories from the database.
         
-        :param input_data: The input data of the flow.
-        :type input_data: Dict[str, Any]
-        :return: The output data of the flow.
-        :rtype: Dict[str, Any]
+        :param input_message: The input message of the flow.
+        :type input_message: FlowMessage
         """
+        input_data = input_message.data
         api_information = self.backend.get_key()
 
         if api_information.backend_used == "openai":
@@ -144,4 +143,8 @@ class ChromaDBFlow(AtomicFlow):
                 )
             response["retrieved"] = ""
 
-        return response
+        reply = self._package_output_message(
+            input_message = input_message,
+            response = response
+        )
+        self.reply_to_message(reply = reply, to = input_message)

VectorStoreFlow.py

@@ -8,7 +8,7 @@ from langchain.embeddings import OpenAIEmbeddings
 from langchain.schema import Document
 from langchain.vectorstores import Chroma, FAISS
 from langchain.vectorstores.base import VectorStoreRetriever
-
+from aiflows.messages import FlowMessage
 from aiflows.base_flows import AtomicFlow
 import hydra
 
@@ -141,16 +141,14 @@ class VectorStoreFlow(AtomicFlow):
         # TODO(yeeef): support metadata
         return [Document(page_content=doc, metadata={"": ""}) for doc in documents]
 
-    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+    def run(self, input_message: FlowMessage):
         """ This method runs the flow. It either writes or reads memories from the database.
         
-        :param input_data: The input data of the flow.
-        :type input_data: Dict[str, Any]
-        :return: The output data of the flow.
-        :rtype: Dict[str, Any]
+        :param input_message: The input data of the flow.
+        :type input_message: FlowMessage
         """
         response = {}
-
+        input_data = input_message.data
         operation = input_data["operation"]
         assert operation in ["write", "read"], f"Operation '{operation}' not supported"
 
@@ -169,4 +167,8 @@ class VectorStoreFlow(AtomicFlow):
             self.vector_db.add_documents(documents)
             response["retrieved"] = ""
 
-        return response
+        reply = self._package_output_message(
+            input_message = input_message,
+            response = response
+        )
+        self.reply_to_message(reply = reply, to = input_message)

demo.yaml

@@ -1,85 +1,18 @@
 chroma_demo_flow:
-  input_interface:
-    - "operation"
-    - "content"
-  output_interface:
-    - "retrieved"
-  _target_: aiflows.base_flows.SequentialFlow.instantiate_from_default_config
-  name: "demoChromaDBFlow"
-  description: "An example flow of how to read and writed in a ChromaDBFlowModule."
-  subflows_config:
-    chroma_db:
-      input_interface:
-        _target_: aiflows.interfaces.KeyInterface
-        keys_to_select: ["operation","content"]
-      _target_: flow_modules.aiflows.VectorStoreFlowModule.ChromaDBFlow.instantiate_from_default_config
-      
-      backend:
-        _target_: aiflows.backends.llm_lite.LiteLLMBackend
-        api_infos: ???
-        model_name: "" #Not used in current implementation
-      n_results: 1 # number of results to retrieve when query
-  topology:
-    - goal: Write content to the ChromaDB
-      input_interface:
-        _target_: aiflows.interfaces.KeyInterface
-        keys_to_select: ["operation","content"]
-      flow: chroma_db
-      output_interface:
-        _target_: aiflows.interfaces.KeyInterface
-        keys_to_set:
-            operation: "read"
-        keys_to_rename:
-            retrieved: content
-        keys_to_select: ["operation","content"]
+  _target_: flow_modules.aiflows.VectorStoreFlowModule.ChromaDBFlow.instantiate_from_default_config
+  
+  backend:
+    _target_: aiflows.backends.llm_lite.LiteLLMBackend
+    api_infos: ???
+    model_name: "" #Not used in current implementation
+  n_results: 1 # number of results to retrieve when query
 
-    - goal: Read content from the ChromaDB
-      input_interface:
-        _target_: aiflows.interfaces.KeyInterface
-        keys_to_select: ["operation","content"]
-      flow: chroma_db
-      output_interface:
-        _target_: aiflows.interfaces.KeyInterface
-        keys_to_select: ["retrieved"]
 
 vector_store_demo_flow:
-  input_interface:
-      - "operation"
-      - "content"
-  output_interface:
-    - "retrieved"
-  name: "demoVectorStoreFlow"
-  description: "An example flow of how to read and write in a VectorStoreFlowModule."
-  _target_: aiflows.base_flows.SequentialFlow.instantiate_from_default_config
-  subflows_config:
-    
-    vs_db:
-      _target_: flow_modules.aiflows.VectorStoreFlowModule.VectorStoreFlow.instantiate_from_default_config
-      backend:
-        _target_: aiflows.backends.llm_lite.LiteLLMBackend
-        api_infos: ???
-        model_name: "" #Not used in current implementation
-
-  topology:
-    - goal: Write content to the VectorStore
-      input_interface:
-        _target_: aiflows.interfaces.KeyInterface
-        keys_to_select: ["operation","content"]
-      flow: vs_db
-      output_interface:
-        _target_: aiflows.interfaces.KeyInterface
-        keys_to_set:
-            operation: "read"
-        keys_to_rename:
-            retrieved: content
-        keys_to_select: ["operation","content"]
 
-    - goal: Read content from the VectorStore
-      input_interface:
-        _target_: aiflows.interfaces.KeyInterface
-        keys_to_select: ["operation","content"]
-      flow: vs_db
-      output_interface:
-        _target_: aiflows.interfaces.KeyInterface
-        keys_to_select: ["retrieved"]
+  _target_: flow_modules.aiflows.VectorStoreFlowModule.VectorStoreFlow.instantiate_from_default_config
+  backend:
+    _target_: aiflows.backends.llm_lite.LiteLLMBackend
+    api_infos: ???
+    model_name: "" #Not used in current implementation
   

run.py

@@ -1,3 +1,4 @@
+"""A simple script to run a Flow that can be used for development and debugging."""
 
 import os
 
@@ -5,104 +6,153 @@ import hydra
 
 import aiflows
 from aiflows.flow_launchers import FlowLauncher
-from aiflows.utils.general_helpers import read_yaml_file
 from aiflows.backends.api_info import ApiInfo
+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/VectorStoreFlowModule", "revision": os.getcwd()},
+    {"url": "aiflows/VectorStoreFlowModule", "revision": os.getcwd()}
 ]
+
 from aiflows import flow_verse
 flow_verse.sync_dependencies(dependencies)
-
 if __name__ == "__main__":
     
+    #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
+    # # 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)
     
-    cfg["vector_store_demo_flow"]["subflows_config"]["vs_db"]["backend"]["api_infos"] = api_information
-    cfg["chroma_demo_flow"]["subflows_config"]["chroma_db"]["backend"]["api_infos"] = api_information
-    
-
-    # ~~~ Get the data ~~~
-    # This can be a list of samples
-    data = {"id": 0, "operation": "write", "content": "demo of writing"}  # Add your data here
-
-    # ~~~ Run inference ~~~
-    path_to_output_file = None
-    # path_to_output_file = "output.jsonl"  # Uncomment this line to save the output to disk
-    
-    #### CHROMA DEMO ####
-    ### DUMBY DEMO OF WRITING  "demo of writing" AND READIN ""  (Nothing)###
-    print("DEMO: ChromaDBFlow")
-
-    flow_with_interfaces_chroma = {
-        "flow": hydra.utils.instantiate(cfg['chroma_demo_flow'], _recursive_=False, _convert_="partial"),
-        "input_interface": (
-            None
-            if getattr(cfg, "input_interface", None) is None
-            else hydra.utils.instantiate(cfg['input_interface'], _recursive_=False)
-        ),
-        "output_interface": (
-            None
-            if getattr(cfg, "output_interface", None) is None
-            else hydra.utils.instantiate(cfg['output_interface'], _recursive_=False)
-        ),
-    }
-    
-    _, outputs = FlowLauncher.launch(
-        flow_with_interfaces=flow_with_interfaces_chroma,
-        data=data,
-        path_to_output_file=path_to_output_file,
+    #3. ~~~~ Serve The Flow ~~~~
+    serve_utils.recursive_serve_flow(
+        cl = cl,
+        flow_type="ChromaDBFlowModule",
+        default_config=cfg["chroma_demo_flow"],
+        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)
 
-    # ~~~ Print the output ~~~
-    flow_output_data = outputs[0]
-    print(flow_output_data)
-    
-    #### END CHROM DEMO ####
-    
-    #### VECTOR STORE DEMO ####
-    
-    print("DEMO: VECTOR STORE DEMO")
-     
-    flow_with_interfaces_vstore = {
-        "flow": hydra.utils.instantiate(cfg['vector_store_demo_flow'], _recursive_=False, _convert_="partial"),
-        "input_interface": (
-            None
-            if getattr(cfg, "input_interface", None) is None
-            else hydra.utils.instantiate(cfg['input_interface'], _recursive_=False)
-        ),
-        "output_interface": (
-            None
-            if getattr(cfg, "output_interface", None) is None
-            else hydra.utils.instantiate(cfg['output_interface'], _recursive_=False)
-        ),
-    }
-    
-    
-    _, outputs = FlowLauncher.launch(
-        flow_with_interfaces=flow_with_interfaces_vstore,
-        data=data,
-        path_to_output_file=path_to_output_file,
+    #5 ~~~~~Mount the flow and get its proxy~~~~~~
+    proxy_flow_cdb = serve_utils.recursive_mount(
+        cl=cl,
+        client_id="local",
+        flow_type="ChromaDBFlowModule",
+        config_overrides=None,
+        initial_state=None,
+        dispatch_point_override=None,
+    ) 
+    
+    #3.(2) ~~~~ Serve The Flow ~~~~
+    serve_utils.recursive_serve_flow(
+        cl = cl,
+        flow_type="VectoreStoreFlowModule",
+        default_config=cfg["vector_store_demo_flow"],
+        default_state=None,
+        default_dispatch_point="coflows_dispatch"
+    )
+    
+    #4.(2) ~~~~~Start A Worker Thread~~~~~
+    run_dispatch_worker_thread(cl, dispatch_point="coflows_dispatch", flow_modules_base_path=FLOW_MODULES_PATH)
+
+    #5.(2) ~~~~~Mount the flow and get its proxy~~~~~~
+    proxy_flow_vs = serve_utils.recursive_mount(
+        cl=cl,
+        client_id="local",
+        flow_type="VectoreStoreFlowModule",
+        config_overrides=None,
+        initial_state=None,
+        dispatch_point_override=None,
+    )   
+    
+    #6. ~~~ Get the data ~~~
+    data_write = {"id": 0, "operation": "write", "content": "The capital of Switzerland is Bern"}  # Add your data here
+    data_read = {"id": 1, "operation": "read", "content": "Capital of Switzerland"}  # Add your data here
+    #option1: use the FlowMessage class
+    input_message_write = FlowMessage(
+        data=data_write,
     )
+    
+    input_message_read = FlowMessage(
+        data=data_read
+    )
+
+    #option2: use the proxy_flow
+    #input_message = proxy_flow._package_input_message(data = data)
+    
+    #7. ~~~ Run inference ~~~
+    print("##########CHROMA DB DEMO###############")
+    #write to DB
+    proxy_flow_cdb.send_message_async(input_message_write)
+    #read from DB
+    future = proxy_flow_cdb.send_message_blocking(input_message_read)
+    
+    #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)
+    
+    
+    print("##########VECTOR STORE DEMO###############")
+    #write to DB
+    proxy_flow_vs.send_message_async(input_message_write)
+    #read from DB
+    future = proxy_flow_vs.send_message_blocking(input_message_read)
 
+    #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, "FlowModule")