diff --git a/agentverse/agentverse.py b/agentverse/agentverse.py index a716b810c..8b1505075 100644 --- a/agentverse/agentverse.py +++ b/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 = [] diff --git a/agentverse/gui.py b/agentverse/gui.py index ca8738b3d..911207f48 100644 --- a/agentverse/gui.py +++ b/agentverse/gui.py @@ -30,7 +30,7 @@ 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 @@ -38,9 +38,9 @@ def __init__(self, task: str): 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 diff --git a/agentverse/initialization.py b/agentverse/initialization.py index cc6e3cd42..ed21733e0 100644 --- a/agentverse/initialization.py +++ b/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): diff --git a/agentverse/simulation.py b/agentverse/simulation.py index f4d3dc982..27a439e4b 100644 --- a/agentverse/simulation.py +++ b/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 = [] diff --git a/agentverse/tasks/simulation/sde_team/readme.md b/agentverse/tasks/simulation/sde_team/readme.md index 62faa3117..b022b820b 100644 --- a/agentverse/tasks/simulation/sde_team/readme.md +++ b/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() ``` diff --git a/agentverse/tasksolving.py b/agentverse/tasksolving.py index b90b81b0e..2bd2bd9b7 100644 --- a/agentverse/tasksolving.py +++ b/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"] diff --git a/benchmark.py b/benchmark.py index 304adc6e1..01d2cb232 100644 --- a/benchmark.py +++ b/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) diff --git a/main_simulation_cli.py b/main_simulation_cli.py index a6be9e66d..48f56e567 100644 --- a/main_simulation_cli.py +++ b/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() diff --git a/main_simulation_gui.py b/main_simulation_gui.py index a7c4ad42f..576b980c5 100644 --- a/main_simulation_gui.py +++ b/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()