From 882b134f351ee1381b9a253a4b23a8d334c7f295 Mon Sep 17 00:00:00 2001 From: Abubakar Abid Date: Wed, 29 Jan 2025 06:35:32 -0800 Subject: [PATCH] Fix return type of generator function in `Agent._run()` (#409) * Fix return type of generator function in `Agent._run()` --- src/smolagents/agents.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py index c9decd9..12a8569 100644 --- a/src/smolagents/agents.py +++ b/src/smolagents/agents.py @@ -57,7 +57,7 @@ from .tools import ( Tool, get_tool_description_with_args, ) -from .types import AgentAudio, AgentImage, handle_agent_output_types +from .types import AgentAudio, AgentImage, AgentType, handle_agent_output_types from .utils import ( AgentError, AgentExecutionError, @@ -564,7 +564,7 @@ You have been provided with these additional arguments, that you can access usin # Outputs are returned only at the end as a string. We only look at the last step return deque(self._run(task=self.task, images=images), maxlen=1)[0] - def _run(self, task: str, images: List[str] | None = None) -> Generator[str, None, None]: + def _run(self, task: str, images: List[str] | None = None) -> Generator[ActionStep | AgentType, None, None]: """ Run the agent in streaming mode and returns a generator of all the steps.