| from python.helpers.extension import Extension | |
| from agent import Agent, LoopData | |
| DATA_NAME_ITER_NO = "iteration_no" | |
| class IterationNo(Extension): | |
| async def execute(self, loop_data: LoopData = LoopData(), **kwargs): | |
| # total iteration number | |
| no = self.agent.get_data(DATA_NAME_ITER_NO) or 0 | |
| self.agent.set_data(DATA_NAME_ITER_NO, no + 1) | |
| def get_iter_no(agent: Agent) -> int: | |
| return agent.get_data(DATA_NAME_ITER_NO) or 0 |