[Feature] Add max_steps to LLMAgent.run and set handler result to MaxStepsReachedError (#91)

* add MaxStepsReachedError and unit tests

* changelog

* nit up date docstring

Author: nerdai

CHANGELOG.md

@@ -8,6 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 ## Unreleased
 
+### Added
+
+- Add `max_steps` to `LLMAgent.run` and set handler result to `MaxStepsReachedError` if reached (#91)
+
 ### Changed
 
 - Improve `NextStepDecision` to allow for only one next_step or task_result (#88)

src/llm_agents_from_scratch/agent/llm_agent.py

@@ -18,7 +18,11 @@
     TaskStepResult,
     ToolCallResult,
 )
-from llm_agents_from_scratch.errors import LLMAgentError, TaskHandlerError
+from llm_agents_from_scratch.errors import (
+    LLMAgentError,
+    MaxStepsReachedError,
+    TaskHandlerError,
+)
 from llm_agents_from_scratch.logger import get_logger
 
 from .templates import TaskHandlerTemplates, default_task_handler_templates
@@ -344,11 +348,13 @@ async def run_step(self, step: TaskStep) -> TaskStepResult:
                 content=final_content,
             )
 
-    def run(self, task: Task) -> TaskHandler:
+    def run(self, task: Task, max_steps: int | None = None) -> TaskHandler:
         """Agent's processing loop for executing tasks.
 
         Args:
             task (Task): the Task to perform.
+            max_steps (int | None): Maximum number of steps to run for task.
+                Defaults to None.
 
         Returns:
             TaskHandler: the TaskHandler object responsible for task execution.
@@ -364,8 +370,12 @@ async def _process_loop() -> None:
             """
             self.logger.info(f"🚀 Starting task: {task.instruction}")
             step_result = None
+            ix = 0
             while not task_handler.done():
                 try:
+                    if max_steps and ix == max_steps:
+                        raise MaxStepsReachedError("Max steps reached.")
+
                     next_step = await task_handler.get_next_step(step_result)
 
                     match next_step:
@@ -386,6 +396,8 @@ async def _process_loop() -> None:
 
                 except Exception as e:
                     task_handler.set_exception(e)
+                finally:
+                    ix += 1
 
         task_handler.background_task = asyncio.create_task(_process_loop())
 

src/llm_agents_from_scratch/errors/__init__.py

@@ -1,4 +1,4 @@
-from .agent import LLMAgentError
+from .agent import LLMAgentError, MaxStepsReachedError
 from .core import LLMAgentsFromScratchError, LLMAgentsFromScratchWarning
 from .task_handler import TaskHandlerError
 
@@ -8,6 +8,7 @@
     "LLMAgentsFromScratchWarning",
     # agent
     "LLMAgentError",
+    "MaxStepsReachedError",
     # task handler
     "TaskHandlerError",
 ]

src/llm_agents_from_scratch/errors/agent.py

@@ -7,3 +7,9 @@ class LLMAgentError(LLMAgentsFromScratchError):
     """Base error for all TaskHandler-related exceptions."""
 
     pass
+
+
+class MaxStepsReachedError(LLMAgentError):
+    """Raised if the maximum number of steps reached in a run() method call."""
+
+    pass

tests/agent/test_agent.py

@@ -12,7 +12,7 @@
     TaskResult,
     TaskStep,
 )
-from llm_agents_from_scratch.errors import LLMAgentError
+from llm_agents_from_scratch.errors import LLMAgentError, MaxStepsReachedError
 
 
 def test_init(mock_llm: BaseLLM) -> None:
@@ -113,3 +113,32 @@ async def test_run_exception(
     await asyncio.sleep(0.1)  # Let it run
 
     assert handler.exception() == err
+
+
+@pytest.mark.asyncio
+@patch.object(LLMAgent.TaskHandler, "get_next_step")
+async def test_run_max_steps_reached_error(
+    mock_get_next_step: AsyncMock,
+    mock_llm: BaseLLM,
+) -> None:
+    """Tests run method reaches max step."""
+
+    # arrange
+    task = Task(instruction="mock instruction")
+    mock_get_next_step.side_effect = [
+        TaskStep(
+            instruction="mock 1",
+            task_id=task.id_,
+        ),
+        TaskStep(
+            instruction="mock 2",
+            task_id=task.id_,
+        ),
+    ]
+    agent = LLMAgent(llm=mock_llm)
+
+    # act
+    handler = agent.run(task, max_steps=1)
+    await asyncio.sleep(0.1)  # Let it run
+
+    assert isinstance(handler.exception(), MaxStepsReachedError)