Fix undefined `step_number` in single_step run (#495)
* Fix undefined step_number * Test step_number is defined after single_step run --------- Co-authored-by: Albert Villanova del Moral <8515462+albertvillanova@users.noreply.github.com>
This commit is contained in:
parent
8b1dd44010
commit
26c733c745
|
@ -423,6 +423,7 @@ You have been provided with these additional arguments, that you can access usin
|
|||
|
||||
self.memory.steps.append(TaskStep(task=self.task, task_images=images))
|
||||
if single_step:
|
||||
self.step_number = 1
|
||||
step_start_time = time.time()
|
||||
memory_step = ActionStep(start_time=step_start_time, observations_images=images)
|
||||
memory_step.end_time = time.time()
|
||||
|
|
|
@ -655,6 +655,13 @@ nested_answer()
|
|||
|
||||
|
||||
class TestMultiStepAgent:
|
||||
def test_step_number(self):
|
||||
fake_model = MagicMock()
|
||||
agent = MultiStepAgent(tools=[], model=fake_model)
|
||||
agent.run("Test task", single_step=True)
|
||||
assert hasattr(agent, "step_number"), "step_number attribute should be defined"
|
||||
assert agent.step_number == 1, "step_number should be set to 1 after run method is called"
|
||||
|
||||
def test_planning_step_first_step(self):
|
||||
fake_model = MagicMock()
|
||||
agent = MultiStepAgent(
|
||||
|
|
Loading…
Reference in New Issue