From 26c733c745007e6f5e84b1fccc6f7be125e4ccc5 Mon Sep 17 00:00:00 2001
From: Taha Yassine <40228615+taha-yassine@users.noreply.github.com>
Date: Wed, 5 Feb 2025 09:00:19 +0100
Subject: [PATCH] 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>
---
 src/smolagents/agents.py | 1 +
 tests/test_agents.py     | 7 +++++++
 2 files changed, 8 insertions(+)

diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py
index 6477217..5db3d78 100644
--- a/src/smolagents/agents.py
+++ b/src/smolagents/agents.py
@@ -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()
diff --git a/tests/test_agents.py b/tests/test_agents.py
index c593551..9caa221 100644
--- a/tests/test_agents.py
+++ b/tests/test_agents.py
@@ -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(