From 08bb8f524c6e0ddc98e3dd040494442c89c6a991 Mon Sep 17 00:00:00 2001 From: "clementine@huggingface.co" Date: Tue, 21 Jan 2025 15:00:47 +0000 Subject: [PATCH 1/8] fuse stream and non stream calls --- src/smolagents/agents.py | 69 ++++++---------------------------------- 1 file changed, 10 insertions(+), 59 deletions(-) diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py index 22af24809..91b8ce151 100644 --- a/src/smolagents/agents.py +++ b/src/smolagents/agents.py @@ -497,12 +497,9 @@ def run( result = self.step(step_log) return result - if stream: - return self.stream_run(self.task) - else: - return self.direct_run(self.task) + return self._run(task = self.task, stream = stream) - def stream_run(self, task: str): + def _run(self, task: str, stream: bool): """ Runs the agent in streaming mode, yielding steps as they are executed: should be launched only in the `run` method. """ @@ -538,7 +535,8 @@ def stream_run(self, task: str): for callback in self.step_callbacks: callback(step_log) self.step_number += 1 - yield step_log + if stream: + yield step_log if final_answer is None and self.step_number == self.max_steps: error_message = "Reached max steps." @@ -551,61 +549,14 @@ def stream_run(self, task: str): final_step_log.duration = step_log.end_time - step_start_time for callback in self.step_callbacks: callback(final_step_log) - yield final_step_log - - yield handle_agent_output_types(final_answer) - - def direct_run(self, task: str): - """ - Runs the agent in direct mode, returning outputs only at the end: should be launched only in the `run` method. - """ - final_answer = None - self.step_number = 0 - while final_answer is None and self.step_number < self.max_steps: - step_start_time = time.time() - step_log = ActionStep(step=self.step_number, start_time=step_start_time) - try: - if self.planning_interval is not None and self.step_number % self.planning_interval == 0: - self.planning_step( - task, - is_first_step=(self.step_number == 0), - step=self.step_number, - ) - self.logger.log( - Rule( - f"[bold]Step {self.step_number}", - characters="━", - style=YELLOW_HEX, - ), - level=LogLevel.INFO, - ) - - # Run one step! - final_answer = self.step(step_log) - - except AgentError as e: - step_log.error = e - finally: - step_end_time = time.time() - step_log.end_time = step_end_time - step_log.duration = step_end_time - step_start_time - self.logs.append(step_log) - for callback in self.step_callbacks: - callback(step_log) - self.step_number += 1 + if stream: + yield final_step_log - if final_answer is None and self.step_number == self.max_steps: - error_message = "Reached max steps." - final_step_log = ActionStep(error=AgentMaxStepsError(error_message)) - self.logs.append(final_step_log) - final_answer = self.provide_final_answer(task) - self.logger.log(Text(f"Final answer: {final_answer}"), level=LogLevel.INFO) - final_step_log.action_output = final_answer - final_step_log.duration = 0 - for callback in self.step_callbacks: - callback(final_step_log) + if stream: + yield handle_agent_output_types(final_answer) + else: + return handle_agent_output_types(final_answer) - return handle_agent_output_types(final_answer) def planning_step(self, task, is_first_step: bool, step: int): """ From e2fe8a480be1c580393d97833256766955d6fb47 Mon Sep 17 00:00:00 2001 From: "clementine@huggingface.co" Date: Tue, 21 Jan 2025 15:06:19 +0000 Subject: [PATCH 2/8] style --- src/smolagents/agents.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py index 91b8ce151..3fab370c9 100644 --- a/src/smolagents/agents.py +++ b/src/smolagents/agents.py @@ -497,7 +497,7 @@ def run( result = self.step(step_log) return result - return self._run(task = self.task, stream = stream) + return self._run(task=self.task, stream=stream) def _run(self, task: str, stream: bool): """ @@ -557,7 +557,6 @@ def _run(self, task: str, stream: bool): else: return handle_agent_output_types(final_answer) - def planning_step(self, task, is_first_step: bool, step: int): """ Used periodically by the agent to plan the next steps to reach the objective. From a982c7ed7bf614fbf93d49e1c99959023c94f0b3 Mon Sep 17 00:00:00 2001 From: "clementine@huggingface.co" Date: Tue, 21 Jan 2025 15:16:40 +0000 Subject: [PATCH 3/8] update type hinting and doc --- src/smolagents/agents.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py index 3fab370c9..ede7d7242 100644 --- a/src/smolagents/agents.py +++ b/src/smolagents/agents.py @@ -17,7 +17,7 @@ import time from dataclasses import dataclass from enum import IntEnum -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union, Generator from rich import box from rich.console import Console, Group @@ -499,9 +499,15 @@ def run( return self._run(task=self.task, stream=stream) - def _run(self, task: str, stream: bool): + def _run(self, task: str, stream: bool) -> Union[Generator[str], str]: """ - Runs the agent in streaming mode, yielding steps as they are executed: should be launched only in the `run` method. + Runs the agent. Running can be done in direct or streaming mode. + Note: in all cases, this function is internal and should be used only in the `run` method. + + Args: + task (`str`): The task to perform. + stream (`bool`): If True, the steps are returned as they are executed through a generator to iterate on. + If False, outputs are returned only at the end. """ final_answer = None self.step_number = 0 @@ -557,7 +563,7 @@ def _run(self, task: str, stream: bool): else: return handle_agent_output_types(final_answer) - def planning_step(self, task, is_first_step: bool, step: int): + def planning_step(self, task, is_first_step: bool, step: int) -> None: """ Used periodically by the agent to plan the next steps to reach the objective. From f8447923b725a8dfee77d884d6c7daa9a6230868 Mon Sep 17 00:00:00 2001 From: "clementine@huggingface.co" Date: Tue, 21 Jan 2025 15:19:36 +0000 Subject: [PATCH 4/8] fix type hinting --- src/smolagents/agents.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py index ede7d7242..d75653143 100644 --- a/src/smolagents/agents.py +++ b/src/smolagents/agents.py @@ -499,7 +499,7 @@ def run( return self._run(task=self.task, stream=stream) - def _run(self, task: str, stream: bool) -> Union[Generator[str], str]: + def _run(self, task: str, stream: bool) -> Union[Generator[str, None, None], str]: """ Runs the agent. Running can be done in direct or streaming mode. Note: in all cases, this function is internal and should be used only in the `run` method. From b3edee6a82707bea2c3af15812fb3e19dbb3c19c Mon Sep 17 00:00:00 2001 From: "clementine@huggingface.co" Date: Tue, 21 Jan 2025 16:35:52 +0000 Subject: [PATCH 5/8] fix style --- src/smolagents/agents.py | 27 +++++++++++++-------------- src/smolagents/types.py | 3 ++- tests/test_python_interpreter.py | 4 +--- 3 files changed, 16 insertions(+), 18 deletions(-) diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py index d75653143..92465f245 100644 --- a/src/smolagents/agents.py +++ b/src/smolagents/agents.py @@ -15,9 +15,10 @@ # See the License for the specific language governing permissions and # limitations under the License. import time +from collections import deque from dataclasses import dataclass from enum import IntEnum -from typing import Any, Callable, Dict, List, Optional, Tuple, Union, Generator +from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union from rich import box from rich.console import Console, Group @@ -497,17 +498,20 @@ def run( result = self.step(step_log) return result - return self._run(task=self.task, stream=stream) + if stream: # We want all the steps + return self._run(task=self.task) + # We only want the last step and want to go through the generator efficiently + return deque(self._run(task=self.task), maxlen=1)[0] - def _run(self, task: str, stream: bool) -> Union[Generator[str, None, None], str]: + def _run(self, task: str) -> Union[str, Generator[str, None, None]]: """ - Runs the agent. Running can be done in direct or streaming mode. + Runs the agent. Running can be done in direct or streaming mode. Note: in all cases, this function is internal and should be used only in the `run` method. Args: task (`str`): The task to perform. - stream (`bool`): If True, the steps are returned as they are executed through a generator to iterate on. - If False, outputs are returned only at the end. + stream (`bool`): If True, the steps are returned as they are executed through a generator to iterate on. + If False, outputs are returned only at the end as a string. """ final_answer = None self.step_number = 0 @@ -541,8 +545,7 @@ def _run(self, task: str, stream: bool) -> Union[Generator[str, None, None], str for callback in self.step_callbacks: callback(step_log) self.step_number += 1 - if stream: - yield step_log + yield step_log if final_answer is None and self.step_number == self.max_steps: error_message = "Reached max steps." @@ -555,13 +558,9 @@ def _run(self, task: str, stream: bool) -> Union[Generator[str, None, None], str final_step_log.duration = step_log.end_time - step_start_time for callback in self.step_callbacks: callback(final_step_log) - if stream: - yield final_step_log + yield final_step_log - if stream: - yield handle_agent_output_types(final_answer) - else: - return handle_agent_output_types(final_answer) + yield handle_agent_output_types(final_answer) def planning_step(self, task, is_first_step: bool, step: int) -> None: """ diff --git a/src/smolagents/types.py b/src/smolagents/types.py index 7077daa59..4ad78c0cc 100644 --- a/src/smolagents/types.py +++ b/src/smolagents/types.py @@ -18,6 +18,7 @@ import tempfile import uuid from io import BytesIO +from typing import Union import numpy as np import requests @@ -252,7 +253,7 @@ def handle_agent_input_types(*args, **kwargs): return args, kwargs -def handle_agent_output_types(output, output_type=None): +def handle_agent_output_types(output, output_type=None) -> Union[str, AgentType]: if output_type in _AGENT_TYPE_MAPPING: # If the class has defined outputs, we can map directly according to the class definition decoded_outputs = _AGENT_TYPE_MAPPING[output_type](output) diff --git a/tests/test_python_interpreter.py b/tests/test_python_interpreter.py index 540720f4c..0817c533d 100644 --- a/tests/test_python_interpreter.py +++ b/tests/test_python_interpreter.py @@ -902,6 +902,4 @@ def test_close_matches_subscript(self): code = 'capitals = {"Czech Republic": "Prague", "Monaco": "Monaco", "Bhutan": "Thimphu"};capitals["Butan"]' with pytest.raises(Exception) as e: evaluate_python_code(code) - assert "Maybe you meant one of these indexes instead" in str( - e - ) and "['Bhutan']" in str(e).replace("\\", "") + assert "Maybe you meant one of these indexes instead" in str(e) and "['Bhutan']" in str(e).replace("\\", "") From 43e723a2382e7da43e47688493117b2e55f09638 Mon Sep 17 00:00:00 2001 From: "clementine@huggingface.co" Date: Tue, 21 Jan 2025 16:49:38 +0000 Subject: [PATCH 6/8] correct type hinting since fn is now a generator --- src/smolagents/agents.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py index 7c326351e..46ff959f7 100644 --- a/src/smolagents/agents.py +++ b/src/smolagents/agents.py @@ -498,20 +498,18 @@ def run( result = self.step(step_log) return result - if stream: # We want all the steps + if stream: + # The steps are returned as they are executed through a generator to iterate on. return self._run(task=self.task) - # We only want the last step and want to go through the generator efficiently + # Outputs are returned only at the end as a string. We only look at the last step return deque(self._run(task=self.task), maxlen=1)[0] - def _run(self, task: str) -> Union[str, Generator[str, None, None]]: + def _run(self, task: str) -> Generator[str, None, None]: """ - Runs the agent. Running can be done in direct or streaming mode. - Note: in all cases, this function is internal and should be used only in the `run` method. + Runs the agent in streaming mode and returns a generator of all the steps. Args: task (`str`): The task to perform. - stream (`bool`): If True, the steps are returned as they are executed through a generator to iterate on. - If False, outputs are returned only at the end as a string. """ final_answer = None self.step_number = 0 From 612c5fbd45fa49df88f1cb73996dcc9e120ba572 Mon Sep 17 00:00:00 2001 From: "clementine@huggingface.co" Date: Tue, 21 Jan 2025 16:52:28 +0000 Subject: [PATCH 7/8] make style" --- src/smolagents/agents.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py index 46ff959f7..8c1378437 100644 --- a/src/smolagents/agents.py +++ b/src/smolagents/agents.py @@ -498,7 +498,7 @@ def run( result = self.step(step_log) return result - if stream: + if stream: # The steps are returned as they are executed through a generator to iterate on. return self._run(task=self.task) # Outputs are returned only at the end as a string. We only look at the last step From d290063fb856a351d697e9eeecff62f891bee5bd Mon Sep 17 00:00:00 2001 From: Aymeric Date: Wed, 22 Jan 2025 10:13:13 +0100 Subject: [PATCH 8/8] Fix type --- src/smolagents/types.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/smolagents/types.py b/src/smolagents/types.py index 4ad78c0cc..7077daa59 100644 --- a/src/smolagents/types.py +++ b/src/smolagents/types.py @@ -18,7 +18,6 @@ import tempfile import uuid from io import BytesIO -from typing import Union import numpy as np import requests @@ -253,7 +252,7 @@ def handle_agent_input_types(*args, **kwargs): return args, kwargs -def handle_agent_output_types(output, output_type=None) -> Union[str, AgentType]: +def handle_agent_output_types(output, output_type=None): if output_type in _AGENT_TYPE_MAPPING: # If the class has defined outputs, we can map directly according to the class definition decoded_outputs = _AGENT_TYPE_MAPPING[output_type](output)