Update run_planwriter.py

run_planwriter.py

@@ -2,12 +2,13 @@ import os
 
 import hydra
 
-from flows.backends.api_info import ApiInfo
-from flows.messages import InputMessage
-from flows.utils.general_helpers import read_yaml_file
+from aiflows.backends.api_info import ApiInfo
+from aiflows.messages import InputMessage
+from aiflows.utils.general_helpers import read_yaml_file
 
-from flows import logging
-from flows.flow_cache import CACHING_PARAMETERS, clear_cache
+from aiflows import logging
+from aiflows.flow_cache import CACHING_PARAMETERS, clear_cache
+from aiflows.utils.general_helpers import update_api_infos
 
 CACHING_PARAMETERS.do_caching = False  # Set to True in order to disable caching
 # clear_cache() # Uncomment this line to clear the cache
@@ -22,7 +23,7 @@ dependencies = [
     {"url": "aiflows/ChatFlowModule", "revision": "a749ad10ed39776ba6721c37d0dc22af49ca0f17"},
 ]
 
-from flows import flow_verse
+from aiflows import flow_verse
 
 flow_verse.sync_dependencies(dependencies)
 
@@ -32,15 +33,12 @@ if __name__ == "__main__":
     api_information = [ApiInfo(backend_used="openai", api_key=os.getenv("OPENAI_API_KEY"))]
     path_to_output_file = None
 
-    # ~~~ setting api information into config ~~~
     current_dir = os.getcwd()
     cfg_path = os.path.join(current_dir, "PlanWriterFlow.yaml")
     cfg = read_yaml_file(cfg_path)
-    cfg["subflows_config"]["Controller"]["backend"]["api_infos"] = api_information
-    cfg["subflows_config"]["Executor"]["subflows_config"]["write_plan"]["subflows_config"]["PlanGenerator"]["backend"]["api_infos"] = api_information
-
-    # # ~~~ setting memory file of planner (empty dict)~~~
-    # cfg["subflows_config"]["Executor"]["subflows_config"]["write_plan"]["memory_files"] = {}
+    
+    # ~~~ setting api information into config ~~~
+    update_api_infos(cfg, api_information)
 
     # ~~~ instantiating the flow and input data ~~~
     PlanWriterFlow = hydra.utils.instantiate(cfg, _recursive_=False, _convert_="partial")