fix: allows users to customize config.yaml of the task

agentverse/agentverse.py

@@ -23,13 +23,13 @@ def __init__(self, agents: List[BaseAgent], environment: BaseEnvironment):
         self.environment = environment
 
     @classmethod
-    def from_task(cls, task: str):
+    def from_task(cls, task: str, tasks_dir: str):
         """Build an AgentVerse from a task name.
         The task name should correspond to a directory in `tasks` directory.
         Then this method will load the configuration from the yaml file in that directory.
         """
         # Prepare the config of the task
-        task_config = prepare_task_config(task)
+        task_config = prepare_task_config(task, tasks_dir)
 
         # Build the agents
         agents = []

agentverse/gui.py

@@ -30,17 +30,17 @@ class GUI:
     the UI of frontend
     """
 
-    def __init__(self, task: str):
+    def __init__(self, task: str, tasks_dir: str):
         """
         init a UI.
         default number of students is 0
         """
         self.messages = []
         self.task = task
         if task == "pipeline_brainstorming":
-            self.backend = TaskSolving.from_task(task)
+            self.backend = TaskSolving.from_task(task, tasks_dir)
         else:
-            self.backend = Simulation.from_task(task)
+            self.backend = Simulation.from_task(task, tasks_dir)
         self.turns_remain = 0
         self.agent_id = {
             self.backend.agents[idx].name: idx

agentverse/initialization.py

@@ -66,9 +66,9 @@ def load_agent(agent_config: Dict) -> BaseAgent:
     return agent
 
 
-def prepare_task_config(task):
+def prepare_task_config(task, tasks_dir):
     """Read the yaml config of the given task in `tasks` directory."""
-    all_task_dir = os.path.join(os.path.dirname(__file__), "tasks")
+    all_task_dir = tasks_dir
     task_path = os.path.join(all_task_dir, task)
     config_path = os.path.join(task_path, "config.yaml")
     if not os.path.exists(task_path):

agentverse/simulation.py

@@ -17,13 +17,13 @@ def __init__(self, agents: List[BaseAgent], environment: BaseEnvironment):
         self.environment = environment
 
     @classmethod
-    def from_task(cls, task: str):
+    def from_task(cls, task: str, tasks_dir: str):
         """Build an AgentVerse from a task name.
         The task name should correspond to a directory in `tasks` directory.
         Then this method will load the configuration from the yaml file in that directory.
         """
         # Prepare the config of the task
-        task_config = prepare_task_config(task)
+        task_config = prepare_task_config(task, tasks_dir)
 
         # Build the agents
         agents = []

agentverse/tasks/simulation/sde_team/readme.md

@@ -72,14 +72,17 @@ python build_config.py
 After generating `config.yaml`, run the `main.py` to start the task.
 
 ```python
+import os
 from agentverse.agentverse import AgentVerse
 from argparse import ArgumentParser
 
 parser = ArgumentParser()
 parser.add_argument("--task", type=str, default="sde_team/sde_team_2players")
+parser.add_argument("--tasks_dir", type=str, default=os.path.join(
+    os.path.dirname(__file__), "agentverse", "tasks"))
 
 args = parser.parse_args()
-agentverse = AgentVerse.from_task(args.task)
+agentverse = AgentVerse.from_task(args.task, args.tasks_dir)
 agentverse.run()
 ```
 

agentverse/tasksolving.py

@@ -23,13 +23,13 @@ def __init__(self, environment: BasicEnvironment, task: str = ""):
         self.task = task
 
     @classmethod
-    def from_task(cls, task: str):
+    def from_task(cls, task: str, tasks_dir: str):
         """Build an AgentVerse from a task name.
         The task name should correspond to a directory in `tasks` directory.
         Then this method will load the configuration from the yaml file in that directory.
         """
         # Prepare the config of the task
-        task_config = prepare_task_config(task)
+        task_config = prepare_task_config(task, tasks_dir)
 
         # Build the environment
         env_config = task_config["environment"]

benchmark.py

@@ -12,7 +12,9 @@
 
 parser = ArgumentParser()
 
-parser.add_argument("--task", type=str, default="responsegen")
+parser.add_argument("--task", type=str, default="tasksolving/responsegen")
+parser.add_argument("--tasks_dir", type=str, default=os.path.join(
+    os.path.dirname(__file__), "agentverse", "tasks"))
 parser.add_argument("--dataset_path", type=str, required=True)
 parser.add_argument("--output_path", type=str, default=None)
 parser.add_argument("--has_tools", action="store_true")
@@ -38,7 +40,7 @@ def get_dataloader(task, dataset_path):
     else:
         os.makedirs(args.output_path, exist_ok=True)
     shutil.copyfile(
-        f"./agentverse/tasks/{args.task}/config.yaml",
+        f"{args.tasks_dir}/{args.task}/config.yaml",
         f"{args.output_path}/config.yaml",
     )
 
@@ -57,7 +59,7 @@ def get_dataloader(task, dataset_path):
             assert args.tool_tmp_path is not None
             with open(args.tool_tmp_path, "w") as f:
                 f.write(json.dumps(example["tools"]))
-        agentversepipeline = TaskSolving.from_task(args.task)
+        agentversepipeline = TaskSolving.from_task(args.task, args.tasks_dir)
         agentversepipeline.environment.set_task_description(example["input"])
         # print(args.single_agent)
         # print(args.discussion_mode)

main_simulation_cli.py

@@ -1,3 +1,4 @@
+import os
 import logging
 from argparse import ArgumentParser
 
@@ -6,10 +7,12 @@
 
 parser = ArgumentParser()
 parser.add_argument("--task", type=str, default="simulation/prisoner_dilemma")
+parser.add_argument("--tasks_dir", type=str, default=os.path.join(
+    os.path.dirname(__file__), "agentverse", "tasks"))
 parser.add_argument("--debug", action="store_true")
 args = parser.parse_args()
 
 logger.set_level(logging.DEBUG if args.debug else logging.INFO)
 
-agentverse = Simulation.from_task(args.task)
+agentverse = Simulation.from_task(args.task, args.tasks_dir)
 agentverse.run()

main_simulation_gui.py

@@ -1,9 +1,12 @@
+import os
 from agentverse.gui import GUI
 from argparse import ArgumentParser
 
 parser = ArgumentParser()
 parser.add_argument("--task", type=str, default="simulation/nlp_classroom_9players")
+parser.add_argument("--tasks_dir", type=str, default=os.path.join(
+    os.path.dirname(__file__), "agentverse", "tasks"))
 args = parser.parse_args()
 
-ui = GUI(args.task)
+ui = GUI(args.task, args.tasks_dir)
 ui.launch()