| 62 | |
| 63 | |
| 64 | class Solution: |
| 65 | def __init__(self, config: SolutionConfig): |
| 66 | self.config = config |
| 67 | self.task: Task = ( |
| 68 | Task(config=TaskConfig(self.config.task)) if self.config.task else None |
| 69 | ) |
| 70 | self.agent_team: AgentTeam = AgentTeam( |
| 71 | config=AgentTeamConfig(self.config.agent_team) |
| 72 | ) |
| 73 | self.sop = SOP(config=SOPConfig(self.config.sop)) |
| 74 | self.sop.init_name_role_hash_for_nodes(agent_team=self.agent_team) |
| 75 | self.sop.init_node_prompts() |
| 76 | |
| 77 | def run(self, mode="test"): |
| 78 | if mode == "train": |
| 79 | trajectory = Trajectory([]) |
| 80 | |
| 81 | while not self.sop.finished: |
| 82 | current_node, current_agent_name = self.sop.next( |
| 83 | environment=self.agent_team.environment |
| 84 | ) |
| 85 | if current_node and current_agent_name: |
| 86 | action = self.agent_team.step(current_agent_name, current_node) |
| 87 | self.agent_team.execute(action) |
| 88 | if mode == "train": |
| 89 | trajectory.add_state( |
| 90 | State( |
| 91 | current_node, |
| 92 | self.agent_team.agents[current_agent_name], |
| 93 | action, |
| 94 | self.agent_team.environment, |
| 95 | ) |
| 96 | ) |
| 97 | else: |
| 98 | assert self.sop.finished == True |
| 99 | # TODO: Save environment shared short term memory |
| 100 | |
| 101 | if mode == "train": |
| 102 | return trajectory |
| 103 | else: |
| 104 | return None |
| 105 | |
| 106 | def update_prompt_template(self, node_name, agent_name, prompt_template): |
| 107 | """更新node中agent的prompt_template""" |
| 108 | pass |
| 109 | |
| 110 | def dump(self, base_path): |
| 111 | """Save the solution configuration to a file.""" |
| 112 | # get the base path, create the directory if it does not exist |
| 113 | base_path = str(base_path) |
| 114 | if base_path.endswith("solution.json"): |
| 115 | base_path = base_path[:-13] |
| 116 | os.makedirs(base_path, exist_ok=True) |
| 117 | |
| 118 | # get the path of the configuration files |
| 119 | solution_path = f"{base_path}/solution.json" |
| 120 | sop_path = f"{base_path}/sop.json" |
| 121 | task_path = f"{base_path}/task.json" |