KeyArgo
/
smolagents
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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:
Taha Yassine 2025-02-05 09:00:19 +01:00 committed by GitHub
parent 8b1dd44010
commit 26c733c745
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/smolagents/agents.py
View File

@ -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)) self.memory.steps.append(TaskStep(task=self.task, task_images=images))
if single_step: if single_step:
self.step_number = 1
step_start_time = time.time() step_start_time = time.time()
memory_step = ActionStep(start_time=step_start_time, observations_images=images) memory_step = ActionStep(start_time=step_start_time, observations_images=images)
memory_step.end_time = time.time() memory_step.end_time = time.time()

7
tests/test_agents.py
View File

@ -655,6 +655,13 @@ nested_answer()
class TestMultiStepAgent: 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): def test_planning_step_first_step(self):
fake_model = MagicMock() fake_model = MagicMock()
agent = MultiStepAgent( agent = MultiStepAgent(