| 306 | |
| 307 | @dataclass |
| 308 | class CacheEnv(BaseEnv): |
| 309 | cache_file: str = '' |
| 310 | bad_count: int = 6 # type: ignore |
| 311 | |
| 312 | def __post_init__(self): |
| 313 | super().__post_init__() |
| 314 | self.cache = self.load_cache() |
| 315 | |
| 316 | def get_cost(self, usages): |
| 317 | """ |
| 318 | 1/100: scaling factor |
| 319 | * / 1000: total k token |
| 320 | """ |
| 321 | return 0.5 * ((usages[0]*1 + usages[1]*1) / 1000) |
| 322 | |
| 323 | def load_cache(self) -> Dict[str, Dict[str, Any]]: |
| 324 | """ |
| 325 | Read the task output of a single agent |
| 326 | {"task_id": str, "task": str, "output": "long output", "answer": "short answer", "usage": List[int], } |
| 327 | """ |
| 328 | |
| 329 | return json.load(open(self.cache_file, 'r')) |
| 330 | |
| 331 | def look_up_cache(self, task_id, agent_name) -> Response: |
| 332 | """ |
| 333 | Look up task output in cache |
| 334 | """ |
| 335 | model_name = self.map[agent_name] |
| 336 | responses = self.cache.get(task_id, {}).get(model_name, None) |
| 337 | idx = random.randint(0, 100000) % len(responses) |
| 338 | response = responses[idx] |
| 339 | if response is None: |
| 340 | return None |
| 341 | agent_resp = Response( |
| 342 | role=response['role'], |
| 343 | raw=response['raw'], |
| 344 | content=response['content'], |
| 345 | end=response['end'], |
| 346 | usage=response['usage'], |
| 347 | cost=0 |
| 348 | ) |
| 349 | return agent_resp |
| 350 | |
| 351 | def get_env_feedback(self, reward_fn, response, name, task) -> dict: |
| 352 | """ |
| 353 | Get environment feedback |
| 354 | """ |
| 355 | c = self.get_cost(response.usage) * self.agents[self.name2idx[name]].cost |
| 356 | r = reward_fn(task.type, response.content, task.ground_turth) if reward_fn else response.end |
| 357 | return {'reward': r, 'cost': c} |
| 358 | |
| 359 | |
| 360 | def run(self, |
| 361 | tasks: List[Task], |
| 362 | pre_reset=True, |
| 363 | reward_fn=None, |
| 364 | ): |
| 365 | if pre_reset: |