Ruff formatting

2024-12-11 16:16:18 +01:00 · 2024-12-11 16:16:18 +01:00 · 67deb6808f
parent 851e177e71
commit 67deb6808f
28 changed files with 1153 additions and 3460 deletions
--- a/agents/init.py
+++ b/agents/init.py
@ -24,10 +24,25 @@ from transformers.utils import (
 _import_structure = {
-    "agents": ["Agent", "CodeAgent", "ManagedAgent", "ReactAgent", "CodeAgent", "JsonAgent", "Toolbox"],
+    "agents": [
        "Agent",
        "CodeAgent",
        "ManagedAgent",
        "ReactAgent",
        "CodeAgent",
        "JsonAgent",
        "Toolbox",
    ],
    "llm_engine": ["HfApiEngine", "TransformersEngine"],
    "monitoring": ["stream_to_gradio"],
-    "tools": ["PipelineTool", "Tool", "ToolCollection", "launch_gradio_demo", "load_tool", "tool"],
+    "tools": [
        "PipelineTool",
        "Tool",
        "ToolCollection",
        "launch_gradio_demo",
        "load_tool",
        "tool",
    ],
 }
 try:
@ -45,10 +60,25 @@ else:
    _import_structure["translation"] = ["TranslationTool"]
 if TYPE_CHECKING:
-    from .agents import Agent, CodeAgent, ManagedAgent, ReactAgent, CodeAgent, JsonAgent, Toolbox
+    from .agents import (
        Agent,
        CodeAgent,
        ManagedAgent,
        ReactAgent,
        CodeAgent,
        JsonAgent,
        Toolbox,
    )
    from .llm_engine import HfApiEngine, TransformersEngine
    from .monitoring import stream_to_gradio
-    from .tools import PipelineTool, Tool, ToolCollection, launch_gradio_demo, load_tool, tool
+    from .tools import (
        PipelineTool,
        Tool,
        ToolCollection,
        launch_gradio_demo,
        load_tool,
        tool,
    )
    try:
        if not is_torch_available():
@ -66,4 +96,6 @@ if TYPE_CHECKING:
 else:
    import sys
-    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
+    sys.modules[__name__] = _LazyModule(
        __name__, globals()["__file__"], _import_structure, module_spec=__spec__
    )
--- a/agents/agent_types.py
+++ b/agents/agent_types.py
@ -19,7 +19,11 @@ import uuid
 import numpy as np
-from transformers.utils import is_soundfile_availble, is_torch_available, is_vision_available
+from transformers.utils import (
    is_soundfile_availble,
    is_torch_available,
    is_vision_available,
 )
 import logging
@ -108,7 +112,9 @@ class AgentImage(AgentType, ImageType):
        elif isinstance(value, np.ndarray):
            self._tensor = torch.from_numpy(value)
        else:
-            raise TypeError(f"Unsupported type for {self.__class__.__name__}: {type(value)}")
+            raise TypeError(
                f"Unsupported type for {self.__class__.__name__}: {type(value)}"
            )
    def _ipython_display_(self, include=None, exclude=None):
        """
@ -159,7 +165,7 @@ class AgentImage(AgentType, ImageType):
            return self._path
-    def save(self, output_bytes, format : str = None, **params):
+    def save(self, output_bytes, format: str = None, **params):
        """
        Saves the image to a file.
        Args:
@ -243,7 +249,9 @@ if is_torch_available():
 def handle_agent_inputs(*args, **kwargs):
    args = [(arg.to_raw() if isinstance(arg, AgentType) else arg) for arg in args]
-    kwargs = {k: (v.to_raw() if isinstance(v, AgentType) else v) for k, v in kwargs.items()}
+    kwargs = {
        k: (v.to_raw() if isinstance(v, AgentType) else v) for k, v in kwargs.items()
    }
    return args, kwargs
--- a/agents/agents.py
+++ b/agents/agents.py
@ -15,12 +15,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import time
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union, Tuple
 from dataclasses import dataclass
 from rich.syntax import Syntax
 from langfuse.decorators import langfuse_context, observe
 from transformers.utils import is_torch_available
 from .utils import console, parse_code_blob, parse_json_tool_call, truncate_content
@ -51,6 +49,7 @@ from .tools import (
 HUGGINGFACE_DEFAULT_TOOLS = {}
 class AgentError(Exception):
    """Base class for other agent-related exceptions"""
@ -60,7 +59,6 @@ class AgentError(Exception):
        console.print(f"[bold red]{message}[/bold red]")
 class AgentParsingError(AgentError):
    """Exception raised for errors in parsing in the agent"""
@ -84,12 +82,14 @@ class AgentGenerationError(AgentError):
    pass
 class AgentStep:
    pass
@dataclass
 class ActionStep(AgentStep):
-    tool_call: str | None = None
+    tool_call: Dict[str, str] | None = None
    start_time: float | None = None
    step_end_time: float | None = None
    iteration: int | None = None
@ -97,32 +97,43 @@ class ActionStep(AgentStep):
    error: AgentError | None = None
    step_duration: float | None = None
    llm_output: str | None = None
    observation: str | None = None
    agent_memory: List[Dict[str, str]] | None = None
    rationale: str | None = None
@dataclass
 class PlanningStep(AgentStep):
    plan: str
    facts: str
@dataclass
 class TaskStep(AgentStep):
    task: str
@dataclass
 class SystemPromptStep(AgentStep):
    system_prompt: str
-def format_prompt_with_tools(toolbox: Toolbox, prompt_template: str, tool_description_template: str) -> str:
+def format_prompt_with_tools(
    toolbox: Toolbox, prompt_template: str, tool_description_template: str
 ) -> str:
    tool_descriptions = toolbox.show_tool_descriptions(tool_description_template)
    prompt = prompt_template.replace("{{tool_descriptions}}", tool_descriptions)
    if "{{tool_names}}" in prompt:
-        prompt = prompt.replace("{{tool_names}}", ", ".join([f"'{tool_name}'" for tool_name in toolbox.tools.keys()]))
+        prompt = prompt.replace(
            "{{tool_names}}",
            ", ".join([f"'{tool_name}'" for tool_name in toolbox.tools.keys()]),
        )
    return prompt
-def show_agents_descriptions(managed_agents: list):
+def show_agents_descriptions(managed_agents: Dict):
    managed_agents_descriptions = """
 You can also give requests to team members.
 Calling a team member works the same as for calling a tool: simply, the only argument you can give in the call is 'request', a long string explaning your request.
@ -133,16 +144,24 @@ Here is a list of the team members that you can call:"""
    return managed_agents_descriptions
-def format_prompt_with_managed_agents_descriptions(prompt_template, managed_agents=None) -> str:
+def format_prompt_with_managed_agents_descriptions(
    prompt_template, managed_agents=None
 ) -> str:
    if managed_agents is not None:
-        return prompt_template.replace("<<managed_agents_descriptions>>", show_agents_descriptions(managed_agents))
+        return prompt_template.replace(
            "<<managed_agents_descriptions>>", show_agents_descriptions(managed_agents)
        )
    else:
        return prompt_template.replace("<<managed_agents_descriptions>>", "")
-def format_prompt_with_imports(prompt_template: str, authorized_imports: List[str]) -> str:
+def format_prompt_with_imports(
    prompt_template: str, authorized_imports: List[str]
 ) -> str:
    if "<<authorized_imports>>" not in prompt_template:
-        raise AgentError("Tag '<<authorized_imports>>' should be provided in the prompt.")
+        raise AgentError(
            "Tag '<<authorized_imports>>' should be provided in the prompt."
        )
    return prompt_template.replace("<<authorized_imports>>", str(authorized_imports))
@ -150,7 +169,7 @@ class BaseAgent:
    def __init__(
        self,
        tools: Union[List[Tool], Toolbox],
-        llm_engine: Callable = None,
+        llm_engine: Optional[Callable[[List[Dict[str, str]]], str]] = None,
        system_prompt: Optional[str] = None,
        tool_description_template: Optional[str] = None,
        additional_args: Dict = {},
@ -159,10 +178,12 @@ class BaseAgent:
        add_base_tools: bool = False,
        verbose: bool = False,
        grammar: Optional[Dict[str, str]] = None,
-        managed_agents: Optional[List] = None,
+        managed_agents: Optional[Dict] = None,
        step_callbacks: Optional[List[Callable]] = None,
        monitor_metrics: bool = True,
    ):
        if llm_engine is None:
            llm_engine = HfApiEngine()
        if system_prompt is None:
            system_prompt = CODE_SYSTEM_PROMPT
        if tool_parser is None:
@ -171,14 +192,16 @@ class BaseAgent:
        self.llm_engine = llm_engine
        self.system_prompt_template = system_prompt
        self.tool_description_template = (
-            tool_description_template if tool_description_template else DEFAULT_TOOL_DESCRIPTION_TEMPLATE
+            tool_description_template
            if tool_description_template
            else DEFAULT_TOOL_DESCRIPTION_TEMPLATE
        )
        self.additional_args = additional_args
        self.max_iterations = max_iterations
        self.tool_parser = tool_parser
        self.grammar = grammar
-        self.managed_agents = None
+        self.managed_agents = {}
        if managed_agents is not None:
            self.managed_agents = {agent.name: agent for agent in managed_agents}
@ -186,9 +209,13 @@ class BaseAgent:
            self._toolbox = tools
            if add_base_tools:
                if not is_torch_available():
-                    raise ImportError("Using the base tools requires torch to be installed.")
+                    raise ImportError(
                        "Using the base tools requires torch to be installed."
                    )
-                self._toolbox.add_base_tools(add_python_interpreter=(self.__class__ == JsonAgent))
+                self._toolbox.add_base_tools(
                    add_python_interpreter=(self.__class__ == JsonAgent)
                )
        else:
            self._toolbox = Toolbox(tools, add_base_tools=add_base_tools)
        self._toolbox.add_tool(FinalAnswerTool())
@ -196,7 +223,9 @@ class BaseAgent:
        self.system_prompt = format_prompt_with_tools(
            self._toolbox, self.system_prompt_template, self.tool_description_template
        )
-        self.system_prompt = format_prompt_with_managed_agents_descriptions(self.system_prompt, self.managed_agents)
+        self.system_prompt = format_prompt_with_managed_agents_descriptions(
            self.system_prompt, self.managed_agents
        )
        self.prompt_messages = None
        self.logs = []
        self.task = None
@ -222,15 +251,20 @@ class BaseAgent:
            self.system_prompt_template,
            self.tool_description_template,
        )
-        self.system_prompt = format_prompt_with_managed_agents_descriptions(self.system_prompt, self.managed_agents)
+        self.system_prompt = format_prompt_with_managed_agents_descriptions(
            self.system_prompt, self.managed_agents
        )
        if hasattr(self, "authorized_imports"):
            self.system_prompt = format_prompt_with_imports(
-                self.system_prompt, list(set(LIST_SAFE_MODULES) | set(self.authorized_imports))
+                self.system_prompt,
                list(set(LIST_SAFE_MODULES) | set(getattr(self, "authorized_imports"))),
            )
        return self.system_prompt
-    def write_inner_memory_from_logs(self, summary_mode: Optional[bool] = False) -> List[Dict[str, str]]:
+    def write_inner_memory_from_logs(
        self, summary_mode: Optional[bool] = False
    ) -> List[Dict[str, str]]:
        """
        Reads past llm_outputs, actions, and observations or errors from the logs into a series of messages
        that can be used as input to the LLM.
@ -253,7 +287,10 @@ class BaseAgent:
                memory.append(thought_message)
                if not summary_mode:
-                    thought_message = {"role": MessageRole.ASSISTANT, "content": "[PLAN]:\n" + step_log.plan.strip()}
+                    thought_message = {
                        "role": MessageRole.ASSISTANT,
                        "content": "[PLAN]:\n" + step_log.plan.strip(),
                    }
                    memory.append(thought_message)
            elif isinstance(step_log, TaskStep):
@ -265,13 +302,17 @@ class BaseAgent:
            elif isinstance(step_log, ActionStep):
                if step_log.llm_output is not None and not summary_mode:
-                    thought_message = {"role": MessageRole.ASSISTANT, "content": step_log.llm_output.strip()}
+                    thought_message = {
                        "role": MessageRole.ASSISTANT,
                        "content": step_log.llm_output.strip(),
                    }
                    memory.append(thought_message)
                if step_log.tool_call is not None and summary_mode:
                    tool_call_message = {
                        "role": MessageRole.ASSISTANT,
-                        "content": f"[STEP {i} TOOL CALL]: " + str(step_log.tool_call).strip(),
+                        "content": f"[STEP {i} TOOL CALL]: "
                        + str(step_log.tool_call).strip(),
                    }
                    memory.append(tool_call_message)
@ -284,15 +325,21 @@ class BaseAgent:
                        )
                    elif step_log.observation is not None:
                        message_content = f"[OUTPUT OF STEP {i}] -> Observation:\n{step_log.observation}"
-                    tool_response_message = {"role": MessageRole.TOOL_RESPONSE, "content": message_content}
+                    tool_response_message = {
                        "role": MessageRole.TOOL_RESPONSE,
                        "content": message_content,
                    }
                    memory.append(tool_response_message)
        return memory
    def get_succinct_logs(self):
-        return [{key: value for key, value in log.items() if key != "agent_memory"} for log in self.logs]
+        return [
            {key: value for key, value in log.items() if key != "agent_memory"}
            for log in self.logs
        ]
-    def extract_action(self, llm_output: str, split_token: str) -> str:
+    def extract_action(self, llm_output: str, split_token: str) -> Tuple[str, str]:
        """
        Parse action from the LLM output
@ -312,54 +359,6 @@ class BaseAgent:
            )
        return rationale.strip(), action.strip()
    def execute_tool_call(self, tool_name: str, arguments: Dict[str, str]) -> Any:
        """
        Execute tool with the provided input and returns the result.
        This method replaces arguments with the actual values from the state if they refer to state variables.
        Args:
            tool_name (`str`): Name of the Tool to execute (should be one from self.toolbox).
            arguments (Dict[str, str]): Arguments passed to the Tool.
        """
        available_tools = self.toolbox.tools
        if self.managed_agents is not None:
            available_tools = {**available_tools, **self.managed_agents}
        if tool_name not in available_tools:
            error_msg = f"Error: unknown tool {tool_name}, should be instead one of {list(available_tools.keys())}."
            console.print(f"[bold red]{error_msg}")
            raise AgentExecutionError(error_msg)
        try:
            if isinstance(arguments, str):
                observation = available_tools[tool_name](arguments)
            elif isinstance(arguments, dict):
                for key, value in arguments.items():
                    if isinstance(value, str) and value in self.state:
                        arguments[key] = self.state[value]
                observation = available_tools[tool_name](**arguments)
            else:
                error_msg = f"Arguments passed to tool should be a dict or string: got a {type(arguments)}."
                console.print(f"[bold red]{error_msg}")
                raise AgentExecutionError(error_msg)
            return observation
        except Exception as e:
            if tool_name in self.toolbox.tools:
                tool_description = get_tool_description_with_args(available_tools[tool_name])
                error_msg = (
                    f"Error in tool call execution: {e}\nYou should only use this tool with a correct input.\n"
                    f"As a reminder, this tool's description is the following:\n{tool_description}"
                )
                console.print(f"[bold red]{error_msg}")
                raise AgentExecutionError(error_msg)
            elif tool_name in self.managed_agents:
                error_msg = (
                    f"Error in calling team member: {e}\nYou should only ask this team member with a correct request.\n"
                    f"As a reminder, this team member's description is the following:\n{available_tools[tool_name]}"
                )
                console.print(f"[bold red]{error_msg}")
                raise AgentExecutionError(error_msg)
    def run(self, **kwargs):
        """To be implemented in the child class"""
        raise NotImplementedError
@ -382,8 +381,6 @@ class ReactAgent(BaseAgent):
        planning_interval: Optional[int] = None,
        **kwargs,
    ):
        if llm_engine is None:
            llm_engine = HfApiEngine()
        if system_prompt is None:
            system_prompt = CODE_SYSTEM_PROMPT
        if tool_description_template is None:
@ -423,8 +420,67 @@ class ReactAgent(BaseAgent):
            console.print(f"[bold red]{error_msg}[/bold red]")
            return error_msg
-    @observe
+    def execute_tool_call(self, tool_name: str, arguments: Dict[str, str]) -> Any:
-    def run(self, task: str, stream: bool = False, reset: bool = True, oneshot: bool = False, **kwargs):
+        """
        Execute tool with the provided input and returns the result.
        This method replaces arguments with the actual values from the state if they refer to state variables.
        Args:
            tool_name (`str`): Name of the Tool to execute (should be one from self.toolbox).
            arguments (Dict[str, str]): Arguments passed to the Tool.
        """
        available_tools = self.toolbox.tools
        if self.managed_agents is not None:
            available_tools = {**available_tools, **self.managed_agents}
        if tool_name not in available_tools:
            error_msg = f"Error: unknown tool {tool_name}, should be instead one of {list(available_tools.keys())}."
            console.print(f"[bold red]{error_msg}")
            raise AgentExecutionError(error_msg)
        try:
            if isinstance(arguments, str):
                observation = available_tools[tool_name](arguments)
            elif isinstance(arguments, dict):
                for key, value in arguments.items():
                    if isinstance(value, str) and value in self.state:
                        arguments[key] = self.state[value]
                observation = available_tools[tool_name](**arguments)
            else:
                error_msg = f"Arguments passed to tool should be a dict or string: got a {type(arguments)}."
                console.print(f"[bold red]{error_msg}")
                raise AgentExecutionError(error_msg)
            return observation
        except Exception as e:
            if tool_name in self.toolbox.tools:
                tool_description = get_tool_description_with_args(
                    available_tools[tool_name]
                )
                error_msg = (
                    f"Error in tool call execution: {e}\nYou should only use this tool with a correct input.\n"
                    f"As a reminder, this tool's description is the following:\n{tool_description}"
                )
                console.print(f"[bold red]{error_msg}")
                raise AgentExecutionError(error_msg)
            elif tool_name in self.managed_agents:
                error_msg = (
                    f"Error in calling team member: {e}\nYou should only ask this team member with a correct request.\n"
                    f"As a reminder, this team member's description is the following:\n{available_tools[tool_name]}"
                )
                console.print(f"[bold red]{error_msg}")
                raise AgentExecutionError(error_msg)
    def step(self, log_entry: ActionStep):
        """To be implemented in children classes"""
        pass
    def run(
        self,
        task: str,
        stream: bool = False,
        reset: bool = True,
        oneshot: bool = False,
        **kwargs,
    ):
        """
        Runs the agent for the given task.
@ -441,10 +497,11 @@ class ReactAgent(BaseAgent):
        agent.run("What is the result of 2 power 3.7384?")
        ```
        """
        print("LANGFUSE REF:", langfuse_context.get_current_trace_url())
        self.task = task
        if len(kwargs) > 0:
-            self.task += f"\nYou have been provided with these initial arguments: {str(kwargs)}."
+            self.task += (
                f"\nYou have been provided with these initial arguments: {str(kwargs)}."
            )
        self.state = kwargs.copy()
        self.initialize_system_prompt()
@ -460,7 +517,7 @@ class ReactAgent(BaseAgent):
            else:
                self.logs.append(system_prompt_step)
-        console.rule("[bold]New task", characters='=')
+        console.rule("[bold]New task", characters="=")
        console.print(self.task)
        self.logs.append(TaskStep(task=task))
@ -489,8 +546,13 @@ class ReactAgent(BaseAgent):
            step_start_time = time.time()
            step_log = ActionStep(iteration=iteration, start_time=step_start_time)
            try:
-                if self.planning_interval is not None and iteration % self.planning_interval == 0:
+                if (
-                    self.planning_step(task, is_first_step=(iteration == 0), iteration=iteration)
+                    self.planning_interval is not None
                    and iteration % self.planning_interval == 0
                ):
                    self.planning_step(
                        task, is_first_step=(iteration == 0), iteration=iteration
                    )
                console.rule("[bold]New step")
                self.step(step_log)
                if step_log.final_answer is not None:
@ -530,8 +592,13 @@ class ReactAgent(BaseAgent):
            step_start_time = time.time()
            step_log = ActionStep(iteration=iteration, start_time=step_start_time)
            try:
-                if self.planning_interval is not None and iteration % self.planning_interval == 0:
+                if (
-                    self.planning_step(task, is_first_step=(iteration == 0), iteration=iteration)
+                    self.planning_interval is not None
                    and iteration % self.planning_interval == 0
                ):
                    self.planning_step(
                        task, is_first_step=(iteration == 0), iteration=iteration
                    )
                console.rule("[bold]New step")
                self.step(step_log)
                if step_log.final_answer is not None:
@ -559,7 +626,7 @@ class ReactAgent(BaseAgent):
        return final_answer
-    def planning_step(self, task, is_first_step: bool = False, iteration: int = None):
+    def planning_step(self, task, is_first_step: bool, iteration: int):
        """
        Used periodically by the agent to plan the next steps to reach the objective.
@ -569,7 +636,10 @@ class ReactAgent(BaseAgent):
            iteration (`int`): The number of the current step, used as an indication for the LLM.
        """
        if is_first_step:
-            message_prompt_facts = {"role": MessageRole.SYSTEM, "content": SYSTEM_PROMPT_FACTS}
+            message_prompt_facts = {
                "role": MessageRole.SYSTEM,
                "content": SYSTEM_PROMPT_FACTS,
            }
            message_prompt_task = {
                "role": MessageRole.USER,
                "content": f"""Here is the task:
@ -589,15 +659,20 @@ Now begin!""",
                "role": MessageRole.USER,
                "content": USER_PROMPT_PLAN.format(
                    task=task,
-                    tool_descriptions=self._toolbox.show_tool_descriptions(self.tool_description_template),
+                    tool_descriptions=self._toolbox.show_tool_descriptions(
                        self.tool_description_template
                    ),
                    managed_agents_descriptions=(
-                        show_agents_descriptions(self.managed_agents) if self.managed_agents is not None else ""
+                        show_agents_descriptions(self.managed_agents)
                        if self.managed_agents is not None
                        else ""
                    ),
                    answer_facts=answer_facts,
                ),
            }
            answer_plan = self.llm_engine(
-                [message_system_prompt_plan, message_user_prompt_plan], stop_sequences=["<end_plan>"]
+                [message_system_prompt_plan, message_user_prompt_plan],
                stop_sequences=["<end_plan>"],
            )
            final_plan_redaction = f"""Here is the plan of action that I will follow to solve the task:
@ -608,7 +683,9 @@ Now begin!""",
 ```
 {answer_facts}
 ```""".strip()
-            self.logs.append(PlanningStep(plan=final_plan_redaction, facts=final_facts_redaction))
+            self.logs.append(
                PlanningStep(plan=final_plan_redaction, facts=final_facts_redaction)
            )
            console.rule("[orange]Initial plan")
            console.print(final_plan_redaction)
        else:  # update plan
@ -625,7 +702,9 @@ Now begin!""",
                "role": MessageRole.USER,
                "content": USER_PROMPT_FACTS_UPDATE,
            }
-            facts_update = self.llm_engine([facts_update_system_prompt] + agent_memory + [facts_update_message])
+            facts_update = self.llm_engine(
                [facts_update_system_prompt] + agent_memory + [facts_update_message]
            )
            # Redact updated plan
            plan_update_message = {
@ -636,25 +715,34 @@ Now begin!""",
                "role": MessageRole.USER,
                "content": USER_PROMPT_PLAN_UPDATE.format(
                    task=task,
-                    tool_descriptions=self._toolbox.show_tool_descriptions(self.tool_description_template),
+                    tool_descriptions=self._toolbox.show_tool_descriptions(
                        self.tool_description_template
                    ),
                    managed_agents_descriptions=(
-                        show_agents_descriptions(self.managed_agents) if self.managed_agents is not None else ""
+                        show_agents_descriptions(self.managed_agents)
                        if self.managed_agents is not None
                        else ""
                    ),
                    facts_update=facts_update,
                    remaining_steps=(self.max_iterations - iteration),
                ),
            }
            plan_update = self.llm_engine(
-                [plan_update_message] + agent_memory + [plan_update_message_user], stop_sequences=["<end_plan>"]
+                [plan_update_message] + agent_memory + [plan_update_message_user],
                stop_sequences=["<end_plan>"],
            )
            # Log final facts and plan
-            final_plan_redaction = PLAN_UPDATE_FINAL_PLAN_REDACTION.format(task=task, plan_update=plan_update)
+            final_plan_redaction = PLAN_UPDATE_FINAL_PLAN_REDACTION.format(
                task=task, plan_update=plan_update
            )
            final_facts_redaction = f"""Here is the updated list of the facts that I know:
 ```
 {facts_update}
 ```"""
-            self.logs.append(PlanningStep(plan=final_plan_redaction, facts=final_facts_redaction))
+            self.logs.append(
                PlanningStep(plan=final_plan_redaction, facts=final_facts_redaction)
            )
            console.rule("[orange]Updated plan")
            console.print(final_plan_redaction)
@ -705,14 +793,20 @@ class JsonAgent(ReactAgent):
        log_entry.agent_memory = agent_memory.copy()
        if self.verbose:
-            console.rule("[italic]Calling LLM engine with this last message:", align="left")
+            console.rule(
                "[italic]Calling LLM engine with this last message:", align="left"
            )
            console.print(self.prompt_messages[-1])
            console.rule()
        try:
-            additional_args = {"grammar": self.grammar} if self.grammar is not None else {}
+            additional_args = (
                {"grammar": self.grammar} if self.grammar is not None else {}
            )
            llm_output = self.llm_engine(
-                self.prompt_messages, stop_sequences=["<end_action>", "Observation:"], **additional_args
+                self.prompt_messages,
                stop_sequences=["<end_action>", "Observation:"],
                **additional_args,
            )
            log_entry.llm_output = llm_output
        except Exception as e:
@ -723,7 +817,9 @@ class JsonAgent(ReactAgent):
            console.print(llm_output)
        # Parse
-        rationale, action = self.extract_action(llm_output=llm_output, split_token="Action:")
+        rationale, action = self.extract_action(
            llm_output=llm_output, split_token="Action:"
        )
        try:
            tool_name, arguments = self.tool_parser(action)
@ -807,12 +903,18 @@ class CodeAgent(ReactAgent):
        )
        self.python_evaluator = evaluate_python_code
-        self.additional_authorized_imports = additional_authorized_imports if additional_authorized_imports else []
+        self.additional_authorized_imports = (
-        self.authorized_imports = list(set(LIST_SAFE_MODULES) | set(self.additional_authorized_imports))
+            additional_authorized_imports if additional_authorized_imports else []
-        self.system_prompt = self.system_prompt.replace("<<authorized_imports>>", str(self.authorized_imports))
+        )
        self.authorized_imports = list(
            set(LIST_SAFE_MODULES) | set(self.additional_authorized_imports)
        )
        self.system_prompt = self.system_prompt.replace(
            "<<authorized_imports>>", str(self.authorized_imports)
        )
        self.custom_tools = {}
-    def step(self, log_entry: Dict[str, Any]):
+    def step(self, log_entry: ActionStep):
        """
        Perform one step in the ReAct framework: the agent thinks, acts, and observes the result.
        The errors are raised here, they are caught and logged in the run() method.
@ -825,14 +927,20 @@ class CodeAgent(ReactAgent):
        log_entry.agent_memory = agent_memory.copy()
        if self.verbose:
-            console.rule("[italic]Calling LLM engine with these last messages:", align="left")
+            console.rule(
                "[italic]Calling LLM engine with these last messages:", align="left"
            )
            console.print(self.prompt_messages[-2:])
            console.rule()
        try:
-            additional_args = {"grammar": self.grammar} if self.grammar is not None else {}
+            additional_args = (
                {"grammar": self.grammar} if self.grammar is not None else {}
            )
            llm_output = self.llm_engine(
-                self.prompt_messages, stop_sequences=["<end_action>", "Observation:"], **additional_args
+                self.prompt_messages,
                stop_sequences=["<end_action>", "Observation:"],
                **additional_args,
            )
            log_entry.llm_output = llm_output
        except Exception as e:
@ -840,13 +948,19 @@ class CodeAgent(ReactAgent):
        if self.verbose:
            console.rule("[italic]Output message of the LLM:")
-            console.print(Syntax(llm_output, lexer='markdown', background_color='default'))
+            console.print(
                Syntax(llm_output, lexer="markdown", background_color="default")
            )
        # Parse
        try:
-            rationale, raw_code_action = self.extract_action(llm_output=llm_output, split_token="Code:")
+            rationale, raw_code_action = self.extract_action(
                llm_output=llm_output, split_token="Code:"
            )
        except Exception as e:
-            console.print(f"Error in extracting action, trying to parse the whole output. Error trace: {e}")
+            console.print(
                f"Error in extracting action, trying to parse the whole output. Error trace: {e}"
            )
            rationale, raw_code_action = llm_output, llm_output
        try:
@ -856,14 +970,17 @@ class CodeAgent(ReactAgent):
            raise AgentParsingError(error_msg)
        log_entry.rationale = rationale
-        log_entry.tool_call = {"tool_name": "code interpreter", "tool_arguments": code_action}
+        log_entry.tool_call = {
            "tool_name": "code interpreter",
            "tool_arguments": code_action,
        }
        # Execute
        if self.verbose:
            console.rule("[italic]Agent thoughts")
            console.print(rationale)
        console.rule("[bold]Agent is executing the code below:", align="left")
-        console.print(Syntax(code_action, lexer='python', background_color='default'))
+        console.print(Syntax(code_action, lexer="python", background_color="default"))
        console.rule("", align="left")
        try:
@ -886,7 +1003,9 @@ class CodeAgent(ReactAgent):
            if result is not None:
                console.rule("Last output from code snippet:", align="left")
                console.print(str(result))
-                observation += "Last output from code snippet:\n" + truncate_content(str(result))
+                observation += "Last output from code snippet:\n" + truncate_content(
                    str(result)
                )
            log_entry.observation = observation
        except Exception as e:
            error_msg = f"Code execution failed due to the following error:\n{str(e)}"
@ -902,7 +1021,14 @@ class CodeAgent(ReactAgent):
 class ManagedAgent:
-    def __init__(self, agent, name, description, additional_prompting=None, provide_run_summary=False):
+    def __init__(
        self,
        agent,
        name,
        description,
        additional_prompting=None,
        provide_run_summary=False,
    ):
        self.agent = agent
        self.name = name
        self.description = description
@ -925,18 +1051,22 @@ Your final_answer WILL HAVE to contain these parts:
 Put all these in your final_answer tool, everything that you do not pass as an argument to final_answer will be lost.
 And even if your task resolution is not successful, please return as much context as possible, so that your manager can act upon this feedback.
-<<additional_prompting>>"""
+{{additional_prompting}}"""
        if self.additional_prompting:
-            full_task = full_task.replace("\n<<additional_prompting>>", self.additional_prompting).strip()
+            full_task = full_task.replace(
                "\n{{additional_prompting}}", self.additional_prompting
            ).strip()
        else:
-            full_task = full_task.replace("\n<<additional_prompting>>", "").strip()
+            full_task = full_task.replace("\n{{additional_prompting}}", "").strip()
        return full_task
    def __call__(self, request, **kwargs):
        full_task = self.write_full_task(request)
        output = self.agent.run(full_task, **kwargs)
        if self.provide_run_summary:
-            answer = f"Here is the final answer from your managed agent '{self.name}':\n"
+            answer = (
                f"Here is the final answer from your managed agent '{self.name}':\n"
            )
            answer += str(output)
            answer += f"\n\nFor more detail, find below a summary of this agent's work:\nSUMMARY OF WORK FROM AGENT '{self.name}':\n"
            for message in self.agent.write_inner_memory_from_logs(summary_mode=True):
--- a/agents/default_tools.py
+++ b/agents/default_tools.py
@ -105,7 +105,9 @@ def get_remote_tools(logger, organization="huggingface-tools"):
    tools = {}
    for space_info in spaces:
        repo_id = space_info.id
-        resolved_config_file = hf_hub_download(repo_id, TOOL_CONFIG_FILE, repo_type="space")
+        resolved_config_file = hf_hub_download(
            repo_id, TOOL_CONFIG_FILE, repo_type="space"
        )
        with open(resolved_config_file, encoding="utf-8") as reader:
            config = json.load(reader)
        task = repo_id.split("/")[-1]
@ -131,7 +133,9 @@ class PythonInterpreterTool(Tool):
        if authorized_imports is None:
            self.authorized_imports = list(set(LIST_SAFE_MODULES))
        else:
-            self.authorized_imports = list(set(LIST_SAFE_MODULES) | set(authorized_imports))
+            self.authorized_imports = list(
                set(LIST_SAFE_MODULES) | set(authorized_imports)
            )
        self.inputs = {
            "code": {
                "type": "string",
@ -145,7 +149,11 @@ class PythonInterpreterTool(Tool):
    def forward(self, code):
        output = str(
-            evaluate_python_code(code, static_tools=BASE_PYTHON_TOOLS, authorized_imports=self.authorized_imports)
+            evaluate_python_code(
                code,
                static_tools=BASE_PYTHON_TOOLS,
                authorized_imports=self.authorized_imports,
            )
        )
        return output
@ -153,16 +161,21 @@ class PythonInterpreterTool(Tool):
 class FinalAnswerTool(Tool):
    name = "final_answer"
    description = "Provides a final answer to the given problem."
-    inputs = {"answer": {"type": "any", "description": "The final answer to the problem"}}
+    inputs = {
        "answer": {"type": "any", "description": "The final answer to the problem"}
    }
    output_type = "any"
    def forward(self, answer):
        return answer
 class UserInputTool(Tool):
    name = "user_input"
    description = "Asks for user's input on a specific question"
-    inputs = {"question": {"type": "string", "description": "The question to ask the user"}}
+    inputs = {
        "question": {"type": "string", "description": "The question to ask the user"}
    }
    output_type = "string"
    def forward(self, question):
--- a/agents/gradio_ui.py
+++ b/agents/gradio_ui.py
@ -18,6 +18,7 @@ from .agent_types import AgentAudio, AgentImage, AgentText
 from .agents import BaseAgent, AgentStep, ActionStep
 import gradio as gr
 def pull_messages_from_step(step_log: AgentStep, test_mode: bool = True):
    """Extract ChatMessage objects from agent steps"""
    if isinstance(step_log, ActionStep):
@ -33,7 +34,9 @@ def pull_messages_from_step(step_log: AgentStep, test_mode: bool = True):
                content=str(content),
            )
        if step_log.observation is not None:
-            yield gr.ChatMessage(role="assistant", content=f"```\n{step_log.observation}\n```")
+            yield gr.ChatMessage(
                role="assistant", content=f"```\n{step_log.observation}\n```"
            )
        if step_log.error is not None:
            yield gr.ChatMessage(
                role="assistant",
@ -42,7 +45,13 @@ def pull_messages_from_step(step_log: AgentStep, test_mode: bool = True):
            )
-def stream_to_gradio(agent, task: str, test_mode: bool = False, reset_agent_memory: bool=False, **kwargs):
+def stream_to_gradio(
    agent,
    task: str,
    test_mode: bool = False,
    reset_agent_memory: bool = False,
    **kwargs,
 ):
    """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""
    for step_log in agent.run(task, stream=True, reset=reset_agent_memory, **kwargs):
@ -52,7 +61,10 @@ def stream_to_gradio(agent, task: str, test_mode: bool = False, reset_agent_memo
    final_answer = step_log  # Last log is the run's final_answer
    if isinstance(final_answer, AgentText):
-        yield gr.ChatMessage(role="assistant", content=f"**Final answer:**\n```\n{final_answer.to_string()}\n```")
+        yield gr.ChatMessage(
            role="assistant",
            content=f"**Final answer:**\n```\n{final_answer.to_string()}\n```",
        )
    elif isinstance(final_answer, AgentImage):
        yield gr.ChatMessage(
            role="assistant",
@ -67,8 +79,9 @@ def stream_to_gradio(agent, task: str, test_mode: bool = False, reset_agent_memo
        yield gr.ChatMessage(role="assistant", content=str(final_answer))
-class GradioUI():
+class GradioUI:
    """A one-line interface to launch your agent in Gradio"""
    def __init__(self, agent: BaseAgent):
        self.agent = agent
@ -83,10 +96,17 @@ class GradioUI():
    def run(self):
        with gr.Blocks() as demo:
            stored_message = gr.State([])
-            chatbot = gr.Chatbot(label="Agent",
+            chatbot = gr.Chatbot(
                label="Agent",
                type="messages",
-                                avatar_images=(None, "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png"))
+                avatar_images=(
                    None,
                    "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png",
                ),
            )
            text_input = gr.Textbox(lines=1, label="Chat Message")
-            text_input.submit(lambda s: (s, ""), [text_input], [stored_message, text_input]).then(self.interact_with_agent, [stored_message, chatbot], [chatbot])
+            text_input.submit(
                lambda s: (s, ""), [text_input], [stored_message, text_input]
            ).then(self.interact_with_agent, [stored_message, chatbot], [chatbot])
        demo.launch()
--- a/agents/llm_engine.py
+++ b/agents/llm_engine.py
@ -39,7 +39,9 @@ class MessageRole(str, Enum):
        return [r.value for r in cls]
-def get_clean_message_list(message_list: List[Dict[str, str]], role_conversions: Dict[str, str] = {}):
+def get_clean_message_list(
    message_list: List[Dict[str, str]], role_conversions: Dict[str, str] = {}
 ):
    """
    Subsequent messages with the same role will be concatenated to a single message.
@ -54,12 +56,17 @@ def get_clean_message_list(message_list: List[Dict[str, str]], role_conversions:
        role = message["role"]
        if role not in MessageRole.roles():
-            raise ValueError(f"Incorrect role {role}, only {MessageRole.roles()} are supported for now.")
+            raise ValueError(
                f"Incorrect role {role}, only {MessageRole.roles()} are supported for now."
            )
        if role in role_conversions:
            message["role"] = role_conversions[role]
-        if len(final_message_list) > 0 and message["role"] == final_message_list[-1]["role"]:
+        if (
            len(final_message_list) > 0
            and message["role"] == final_message_list[-1]["role"]
        ):
            final_message_list[-1]["content"] += "\n=======\n" + message["content"]
        else:
            final_message_list.append(message)
@ -81,8 +88,12 @@ class HfEngine:
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        except Exception as e:
-            logger.warning(f"Failed to load tokenizer for model {model_id}: {e}. Loading default tokenizer instead.")
+            logger.warning(
-            self.tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-1.7B-Instruct")
+                f"Failed to load tokenizer for model {model_id}: {e}. Loading default tokenizer instead."
            )
            self.tokenizer = AutoTokenizer.from_pretrained(
                "HuggingFaceTB/SmolLM2-1.7B-Instruct"
            )
    def get_token_counts(self):
        return {
@ -91,12 +102,18 @@ class HfEngine:
        }
    def generate(
-        self, messages: List[Dict[str, str]], stop_sequences: Optional[List[str]] = None, grammar: Optional[str] = None
+        self,
        messages: List[Dict[str, str]],
        stop_sequences: Optional[List[str]] = None,
        grammar: Optional[str] = None,
    ):
        raise NotImplementedError
    def __call__(
-        self, messages: List[Dict[str, str]], stop_sequences: Optional[List[str]] = None, grammar: Optional[str] = None
+        self,
        messages: List[Dict[str, str]],
        stop_sequences: Optional[List[str]] = None,
        grammar: Optional[str] = None,
    ) -> str:
        """Process the input messages and return the model's response.
@ -127,11 +144,15 @@ class HfEngine:
            ```
        """
        if not isinstance(messages, List):
-            raise ValueError("Messages should be a list of dictionaries with 'role' and 'content' keys.")
+            raise ValueError(
                "Messages should be a list of dictionaries with 'role' and 'content' keys."
            )
        if stop_sequences is None:
            stop_sequences = []
        response = self.generate(messages, stop_sequences, grammar)
-        self.last_input_token_count = len(self.tokenizer.apply_chat_template(messages, tokenize=True))
+        self.last_input_token_count = len(
            self.tokenizer.apply_chat_template(messages, tokenize=True)
        )
        self.last_output_token_count = len(self.tokenizer.encode(response))
        # Remove stop sequences from LLM output
@ -175,18 +196,28 @@ class HfApiEngine(HfEngine):
        self.max_tokens = max_tokens
    def generate(
-        self, messages: List[Dict[str, str]], stop_sequences: Optional[List[str]] = None, grammar: Optional[str] = None
+        self,
        messages: List[Dict[str, str]],
        stop_sequences: Optional[List[str]] = None,
        grammar: Optional[str] = None,
    ) -> str:
        # Get clean message list
-        messages = get_clean_message_list(messages, role_conversions=llama_role_conversions)
+        messages = get_clean_message_list(
            messages, role_conversions=llama_role_conversions
        )
        # Send messages to the Hugging Face Inference API
        if grammar is not None:
            response = self.client.chat_completion(
-                messages, stop=stop_sequences, max_tokens=self.max_tokens, response_format=grammar
+                messages,
                stop=stop_sequences,
                max_tokens=self.max_tokens,
                response_format=grammar,
            )
        else:
-            response = self.client.chat_completion(messages, stop=stop_sequences, max_tokens=self.max_tokens)
+            response = self.client.chat_completion(
                messages, stop=stop_sequences, max_tokens=self.max_tokens
            )
        response = response.choices[0].message.content
        return response
@ -207,7 +238,9 @@ class TransformersEngine(HfEngine):
        max_length: int = 1500,
    ) -> str:
        # Get clean message list
-        messages = get_clean_message_list(messages, role_conversions=llama_role_conversions)
+        messages = get_clean_message_list(
            messages, role_conversions=llama_role_conversions
        )
        # Get LLM output
        if stop_sequences is not None and len(stop_sequences) > 0:
--- a/agents/monitoring.py
+++ b/agents/monitoring.py
@ -17,12 +17,14 @@
 from .utils import console
 class Monitor:
    def __init__(self, tracked_llm_engine):
        self.step_durations = []
        self.tracked_llm_engine = tracked_llm_engine
-        if getattr(self.tracked_llm_engine, "last_input_token_count", "Not found") != "Not found":
+        if (
            getattr(self.tracked_llm_engine, "last_input_token_count", "Not found")
            != "Not found"
        ):
            self.total_input_token_count = 0
            self.total_output_token_count = 0
@ -33,103 +35,11 @@ class Monitor:
        console.print(f"- Time taken: {step_duration:.2f} seconds")
        if getattr(self.tracked_llm_engine, "last_input_token_count", None) is not None:
-            self.total_input_token_count += self.tracked_llm_engine.last_input_token_count
+            self.total_input_token_count += (
-            self.total_output_token_count += self.tracked_llm_engine.last_output_token_count
+                self.tracked_llm_engine.last_input_token_count
            )
            self.total_output_token_count += (
                self.tracked_llm_engine.last_output_token_count
            )
            console.print(f"- Input tokens: {self.total_input_token_count:,}")
            console.print(f"- Output tokens: {self.total_output_token_count:,}")
 from typing import Optional, Union, List, Any
 import httpx
 import logging
 import os
 from langfuse.client import Langfuse, StatefulTraceClient, StatefulSpanClient, StateType
 class BaseTracker:
    def __init__(self):
        pass
    @classmethod
    def call(cls, *args, **kwargs):
        pass
 class LangfuseTracker(BaseTracker):
    log = logging.getLogger("langfuse")
    def __init__(self, *, public_key: Optional[str] = None, secret_key: Optional[str] = None,
                 host: Optional[str] = None, debug: bool = False, stateful_client: Optional[
                Union[StatefulTraceClient, StatefulSpanClient]
            ] = None, update_stateful_client: bool = False, version: Optional[str] = None,
                 session_id: Optional[str] = None, user_id: Optional[str] = None, trace_name: Optional[str] = None,
                 release: Optional[str] = None, metadata: Optional[Any] = None, tags: Optional[List[str]] = None,
                 threads: Optional[int] = None, flush_at: Optional[int] = None, flush_interval: Optional[int] = None,
                 max_retries: Optional[int] = None, timeout: Optional[int] = None, enabled: Optional[bool] = None,
                 httpx_client: Optional[httpx.Client] = None, sdk_integration: str = "default") -> None:
        super().__init__()
        self.version = version
        self.session_id = session_id
        self.user_id = user_id
        self.trace_name = trace_name
        self.release = release
        self.metadata = metadata
        self.tags = tags
        self.root_span = None
        self.update_stateful_client = update_stateful_client
        self.langfuse = None
        prio_public_key = public_key or os.environ.get("LANGFUSE_PUBLIC_KEY")
        prio_secret_key = secret_key or os.environ.get("LANGFUSE_SECRET_KEY")
        prio_host = host or os.environ.get(
            "LANGFUSE_HOST", "https://cloud.langfuse.com"
        )
        if stateful_client and isinstance(stateful_client, StatefulTraceClient):
            self.trace = stateful_client
            self._task_manager = stateful_client.task_manager
            return
        elif stateful_client and isinstance(stateful_client, StatefulSpanClient):
            self.root_span = stateful_client
            self.trace = StatefulTraceClient(
                stateful_client.client,
                stateful_client.trace_id,
                StateType.TRACE,
                stateful_client.trace_id,
                stateful_client.task_manager,
            )
            self._task_manager = stateful_client.task_manager
            return
        args = {
            "public_key": prio_public_key,
            "secret_key": prio_secret_key,
            "host": prio_host,
            "debug": debug,
        }
        if release is not None:
            args["release"] = release
        if threads is not None:
            args["threads"] = threads
        if flush_at is not None:
            args["flush_at"] = flush_at
        if flush_interval is not None:
            args["flush_interval"] = flush_interval
        if max_retries is not None:
            args["max_retries"] = max_retries
        if timeout is not None:
            args["timeout"] = timeout
        if enabled is not None:
            args["enabled"] = enabled
        if httpx_client is not None:
            args["httpx_client"] = httpx_client
        args["sdk_integration"] = sdk_integration
        self.langfuse = Langfuse(**args)
        self.trace: Optional[StatefulTraceClient] = None
        self._task_manager = self.langfuse.task_manager
    def call(self, i, o, name=None, **kwargs):
        self.langfuse.trace(input=i, output=o, name=name, metadata=kwargs)
--- a/agents/prompts.py
+++ b/agents/prompts.py
@ -42,7 +42,10 @@ def download_prompt(prompt_or_repo_id, agent_name, mode="run"):
        return prompt_or_repo_id
    prompt_file = cached_file(
-        prompt_or_repo_id, PROMPT_FILES[mode], repo_type="dataset", user_agent={"agent": agent_name}
+        prompt_or_repo_id,
        PROMPT_FILES[mode],
        repo_type="dataset",
        user_agent={"agent": agent_name},
    )
    with open(prompt_file, "r", encoding="utf-8") as f:
        return f.read()
--- a/agents/python_interpreter.py
+++ b/agents/python_interpreter.py
@ -26,6 +26,7 @@ import pandas as pd
 from .utils import truncate_content
 class InterpreterError(ValueError):
    """
    An error raised when the interpretor cannot evaluate a Python expression, due to syntax error or unsupported
@ -38,7 +39,8 @@ class InterpreterError(ValueError):
 ERRORS = {
    name: getattr(builtins, name)
    for name in dir(builtins)
-    if isinstance(getattr(builtins, name), type) and issubclass(getattr(builtins, name), BaseException)
+    if isinstance(getattr(builtins, name), type)
    and issubclass(getattr(builtins, name), BaseException)
 }
@ -92,7 +94,9 @@ def evaluate_unaryop(expression, state, static_tools, custom_tools):
    elif isinstance(expression.op, ast.Invert):
        return ~operand
    else:
-        raise InterpreterError(f"Unary operation {expression.op.__class__.__name__} is not supported.")
+        raise InterpreterError(
            f"Unary operation {expression.op.__class__.__name__} is not supported."
        )
 def evaluate_lambda(lambda_expression, state, static_tools, custom_tools):
@ -102,7 +106,9 @@ def evaluate_lambda(lambda_expression, state, static_tools, custom_tools):
        new_state = state.copy()
        for arg, value in zip(args, values):
            new_state[arg] = value
-        return evaluate_ast(lambda_expression.body, new_state, static_tools, custom_tools)
+        return evaluate_ast(
            lambda_expression.body, new_state, static_tools, custom_tools
        )
    return lambda_func
@ -120,7 +126,9 @@ def evaluate_while(while_loop, state, static_tools, custom_tools):
                break
        iterations += 1
        if iterations > max_iterations:
-            raise InterpreterError(f"Maximum number of {max_iterations} iterations in While loop exceeded")
+            raise InterpreterError(
                f"Maximum number of {max_iterations} iterations in While loop exceeded"
            )
    return None
@ -128,7 +136,10 @@ def create_function(func_def, state, static_tools, custom_tools):
    def new_func(*args, **kwargs):
        func_state = state.copy()
        arg_names = [arg.arg for arg in func_def.args.args]
-        default_values = [evaluate_ast(d, state, static_tools, custom_tools) for d in func_def.args.defaults]
+        default_values = [
            evaluate_ast(d, state, static_tools, custom_tools)
            for d in func_def.args.defaults
        ]
        # Apply default values
        defaults = dict(zip(arg_names[-len(default_values) :], default_values))
@ -180,26 +191,39 @@ def create_class(class_name, class_bases, class_body):
 def evaluate_function_def(func_def, state, static_tools, custom_tools):
-    custom_tools[func_def.name] = create_function(func_def, state, static_tools, custom_tools)
+    custom_tools[func_def.name] = create_function(
        func_def, state, static_tools, custom_tools
    )
    return custom_tools[func_def.name]
 def evaluate_class_def(class_def, state, static_tools, custom_tools):
    class_name = class_def.name
-    bases = [evaluate_ast(base, state, static_tools, custom_tools) for base in class_def.bases]
+    bases = [
        evaluate_ast(base, state, static_tools, custom_tools)
        for base in class_def.bases
    ]
    class_dict = {}
    for stmt in class_def.body:
        if isinstance(stmt, ast.FunctionDef):
-            class_dict[stmt.name] = evaluate_function_def(stmt, state, static_tools, custom_tools)
+            class_dict[stmt.name] = evaluate_function_def(
                stmt, state, static_tools, custom_tools
            )
        elif isinstance(stmt, ast.Assign):
            for target in stmt.targets:
                if isinstance(target, ast.Name):
-                    class_dict[target.id] = evaluate_ast(stmt.value, state, static_tools, custom_tools)
+                    class_dict[target.id] = evaluate_ast(
                        stmt.value, state, static_tools, custom_tools
                    )
                elif isinstance(target, ast.Attribute):
-                    class_dict[target.attr] = evaluate_ast(stmt.value, state, static_tools, custom_tools)
+                    class_dict[target.attr] = evaluate_ast(
                        stmt.value, state, static_tools, custom_tools
                    )
        else:
-            raise InterpreterError(f"Unsupported statement in class body: {stmt.__class__.__name__}")
+            raise InterpreterError(
                f"Unsupported statement in class body: {stmt.__class__.__name__}"
            )
    new_class = type(class_name, tuple(bases), class_dict)
    state[class_name] = new_class
@ -223,7 +247,9 @@ def evaluate_augassign(expression, state, static_tools, custom_tools):
        elif isinstance(target, ast.List):
            return [get_current_value(elt) for elt in target.elts]
        else:
-            raise InterpreterError("AugAssign not supported for {type(target)} targets.")
+            raise InterpreterError(
                "AugAssign not supported for {type(target)} targets."
            )
    current_value = get_current_value(expression.target)
    value_to_add = evaluate_ast(expression.value, state, static_tools, custom_tools)
@ -232,7 +258,9 @@ def evaluate_augassign(expression, state, static_tools, custom_tools):
    if isinstance(expression.op, ast.Add):
        if isinstance(current_value, list):
            if not isinstance(value_to_add, list):
-                raise InterpreterError(f"Cannot add non-list value {value_to_add} to a list.")
+                raise InterpreterError(
                    f"Cannot add non-list value {value_to_add} to a list."
                )
            updated_value = current_value + value_to_add
        else:
            updated_value = current_value + value_to_add
@ -259,7 +287,9 @@ def evaluate_augassign(expression, state, static_tools, custom_tools):
    elif isinstance(expression.op, ast.RShift):
        updated_value = current_value >> value_to_add
    else:
-        raise InterpreterError(f"Operation {type(expression.op).__name__} is not supported.")
+        raise InterpreterError(
            f"Operation {type(expression.op).__name__} is not supported."
        )
    # Update the state
    set_value(expression.target, updated_value, state, static_tools, custom_tools)
@ -311,7 +341,9 @@ def evaluate_binop(binop, state, static_tools, custom_tools):
    elif isinstance(binop.op, ast.RShift):
        return left_val >> right_val
    else:
-        raise NotImplementedError(f"Binary operation {type(binop.op).__name__} is not implemented.")
+        raise NotImplementedError(
            f"Binary operation {type(binop.op).__name__} is not implemented."
        )
 def evaluate_assign(assign, state, static_tools, custom_tools):
@ -321,7 +353,9 @@ def evaluate_assign(assign, state, static_tools, custom_tools):
        set_value(target, result, state, static_tools, custom_tools)
    else:
        if len(assign.targets) != len(result):
-            raise InterpreterError(f"Assign failed: expected {len(result)} values but got {len(assign.targets)}.")
+            raise InterpreterError(
                f"Assign failed: expected {len(result)} values but got {len(assign.targets)}."
            )
        expanded_values = []
        for tgt in assign.targets:
            if isinstance(tgt, ast.Starred):
@ -336,7 +370,9 @@ def evaluate_assign(assign, state, static_tools, custom_tools):
 def set_value(target, value, state, static_tools, custom_tools):
    if isinstance(target, ast.Name):
        if target.id in static_tools:
-            raise InterpreterError(f"Cannot assign to name '{target.id}': doing this would erase the existing tool!")
+            raise InterpreterError(
                f"Cannot assign to name '{target.id}': doing this would erase the existing tool!"
            )
        state[target.id] = value
    elif isinstance(target, ast.Tuple):
        if not isinstance(value, tuple):
@ -399,9 +435,14 @@ def evaluate_call(call, state, static_tools, custom_tools):
        else:
            args.append(evaluate_ast(arg, state, static_tools, custom_tools))
-    kwargs = {keyword.arg: evaluate_ast(keyword.value, state, static_tools, custom_tools) for keyword in call.keywords}
+    kwargs = {
        keyword.arg: evaluate_ast(keyword.value, state, static_tools, custom_tools)
        for keyword in call.keywords
    }
-    if isinstance(func, type) and len(func.__module__.split(".")) > 1:  # Check for user-defined classes
+    if (
        isinstance(func, type) and len(func.__module__.split(".")) > 1
    ):  # Check for user-defined classes
        # Instantiate the class using its constructor
        obj = func.__new__(func)  # Create a new instance of the class
        if hasattr(obj, "__init__"):  # Check if the class has an __init__ method
@ -441,7 +482,9 @@ def evaluate_subscript(subscript, state, static_tools, custom_tools):
    value = evaluate_ast(subscript.value, state, static_tools, custom_tools)
    if isinstance(value, str) and isinstance(index, str):
-        raise InterpreterError("You're trying to subscript a string with a string index, which is impossible")
+        raise InterpreterError(
            "You're trying to subscript a string with a string index, which is impossible"
        )
    if isinstance(value, pd.core.indexing._LocIndexer):
        parent_object = value.obj
        return parent_object.loc[index]
@ -453,11 +496,15 @@ def evaluate_subscript(subscript, state, static_tools, custom_tools):
        return value[index]
    elif isinstance(value, (list, tuple)):
        if not (-len(value) <= index < len(value)):
-            raise InterpreterError(f"Index {index} out of bounds for list of length {len(value)}")
+            raise InterpreterError(
                f"Index {index} out of bounds for list of length {len(value)}"
            )
        return value[int(index)]
    elif isinstance(value, str):
        if not (-len(value) <= index < len(value)):
-            raise InterpreterError(f"Index {index} out of bounds for string of length {len(value)}")
+            raise InterpreterError(
                f"Index {index} out of bounds for string of length {len(value)}"
            )
        return value[index]
    elif index in value:
        return value[index]
@ -483,7 +530,10 @@ def evaluate_name(name, state, static_tools, custom_tools):
 def evaluate_condition(condition, state, static_tools, custom_tools):
    left = evaluate_ast(condition.left, state, static_tools, custom_tools)
-    comparators = [evaluate_ast(c, state, static_tools, custom_tools) for c in condition.comparators]
+    comparators = [
        evaluate_ast(c, state, static_tools, custom_tools)
        for c in condition.comparators
    ]
    ops = [type(op) for op in condition.ops]
    result = True
@ -561,9 +611,13 @@ def evaluate_for(for_loop, state, static_tools, custom_tools):
 def evaluate_listcomp(listcomp, state, static_tools, custom_tools):
    def inner_evaluate(generators, index, current_state):
        if index >= len(generators):
-            return [evaluate_ast(listcomp.elt, current_state, static_tools, custom_tools)]
+            return [
                evaluate_ast(listcomp.elt, current_state, static_tools, custom_tools)
            ]
        generator = generators[index]
-        iter_value = evaluate_ast(generator.iter, current_state, static_tools, custom_tools)
+        iter_value = evaluate_ast(
            generator.iter, current_state, static_tools, custom_tools
        )
        result = []
        for value in iter_value:
            new_state = current_state.copy()
@ -572,7 +626,10 @@ def evaluate_listcomp(listcomp, state, static_tools, custom_tools):
                    new_state[elem.id] = value[idx]
            else:
                new_state[generator.target.id] = value
-            if all(evaluate_ast(if_clause, new_state, static_tools, custom_tools) for if_clause in generator.ifs):
+            if all(
                evaluate_ast(if_clause, new_state, static_tools, custom_tools)
                for if_clause in generator.ifs
            ):
                result.extend(inner_evaluate(generators, index + 1, new_state))
        return result
@ -586,7 +643,9 @@ def evaluate_try(try_node, state, static_tools, custom_tools):
    except Exception as e:
        matched = False
        for handler in try_node.handlers:
-            if handler.type is None or isinstance(e, evaluate_ast(handler.type, state, static_tools, custom_tools)):
+            if handler.type is None or isinstance(
                e, evaluate_ast(handler.type, state, static_tools, custom_tools)
            ):
                matched = True
                if handler.name:
                    state[handler.name] = e
@ -638,7 +697,9 @@ def evaluate_assert(assert_node, state, static_tools, custom_tools):
 def evaluate_with(with_node, state, static_tools, custom_tools):
    contexts = []
    for item in with_node.items:
-        context_expr = evaluate_ast(item.context_expr, state, static_tools, custom_tools)
+        context_expr = evaluate_ast(
            item.context_expr, state, static_tools, custom_tools
        )
        if item.optional_vars:
            state[item.optional_vars.id] = context_expr.__enter__()
            contexts.append(state[item.optional_vars.id])
@ -661,7 +722,9 @@ def evaluate_with(with_node, state, static_tools, custom_tools):
 def import_modules(expression, state, authorized_imports):
    def check_module_authorized(module_name):
        module_path = module_name.split(".")
-        module_subpaths = [".".join(module_path[:i]) for i in range(1, len(module_path) + 1)]
+        module_subpaths = [
            ".".join(module_path[:i]) for i in range(1, len(module_path) + 1)
        ]
        return any(subpath in authorized_imports for subpath in module_subpaths)
    if isinstance(expression, ast.Import):
@ -676,7 +739,9 @@ def import_modules(expression, state, authorized_imports):
        return None
    elif isinstance(expression, ast.ImportFrom):
        if check_module_authorized(expression.module):
-            module = __import__(expression.module, fromlist=[alias.name for alias in expression.names])
+            module = __import__(
                expression.module, fromlist=[alias.name for alias in expression.names]
            )
            for alias in expression.names:
                state[alias.asname or alias.name] = getattr(module, alias.name)
        else:
@ -691,9 +756,14 @@ def evaluate_dictcomp(dictcomp, state, static_tools, custom_tools):
        for value in iter_value:
            new_state = state.copy()
            set_value(gen.target, value, new_state, static_tools, custom_tools)
-            if all(evaluate_ast(if_clause, new_state, static_tools, custom_tools) for if_clause in gen.ifs):
+            if all(
                evaluate_ast(if_clause, new_state, static_tools, custom_tools)
                for if_clause in gen.ifs
            ):
                key = evaluate_ast(dictcomp.key, new_state, static_tools, custom_tools)
-                val = evaluate_ast(dictcomp.value, new_state, static_tools, custom_tools)
+                val = evaluate_ast(
                    dictcomp.value, new_state, static_tools, custom_tools
                )
                result[key] = val
    return result
@ -744,7 +814,10 @@ def evaluate_ast(
        # Constant -> just return the value
        return expression.value
    elif isinstance(expression, ast.Tuple):
-        return tuple(evaluate_ast(elt, state, static_tools, custom_tools) for elt in expression.elts)
+        return tuple(
            evaluate_ast(elt, state, static_tools, custom_tools)
            for elt in expression.elts
        )
    elif isinstance(expression, (ast.ListComp, ast.GeneratorExp)):
        return evaluate_listcomp(expression, state, static_tools, custom_tools)
    elif isinstance(expression, ast.UnaryOp):
@ -770,8 +843,13 @@ def evaluate_ast(
        return evaluate_function_def(expression, state, static_tools, custom_tools)
    elif isinstance(expression, ast.Dict):
        # Dict -> evaluate all keys and values
-        keys = [evaluate_ast(k, state, static_tools, custom_tools) for k in expression.keys]
+        keys = [
-        values = [evaluate_ast(v, state, static_tools, custom_tools) for v in expression.values]
+            evaluate_ast(k, state, static_tools, custom_tools) for k in expression.keys
        ]
        values = [
            evaluate_ast(v, state, static_tools, custom_tools)
            for v in expression.values
        ]
        return dict(zip(keys, values))
    elif isinstance(expression, ast.Expr):
        # Expression -> evaluate the content
@ -788,10 +866,18 @@ def evaluate_ast(
    elif hasattr(ast, "Index") and isinstance(expression, ast.Index):
        return evaluate_ast(expression.value, state, static_tools, custom_tools)
    elif isinstance(expression, ast.JoinedStr):
-        return "".join([str(evaluate_ast(v, state, static_tools, custom_tools)) for v in expression.values])
+        return "".join(
            [
                str(evaluate_ast(v, state, static_tools, custom_tools))
                for v in expression.values
            ]
        )
    elif isinstance(expression, ast.List):
        # List -> evaluate all elements
-        return [evaluate_ast(elt, state, static_tools, custom_tools) for elt in expression.elts]
+        return [
            evaluate_ast(elt, state, static_tools, custom_tools)
            for elt in expression.elts
        ]
    elif isinstance(expression, ast.Name):
        # Name -> pick up the value in the state
        return evaluate_name(expression, state, static_tools, custom_tools)
@ -815,7 +901,9 @@ def evaluate_ast(
            evaluate_ast(expression.upper, state, static_tools, custom_tools)
            if expression.upper is not None
            else None,
-            evaluate_ast(expression.step, state, static_tools, custom_tools) if expression.step is not None else None,
+            evaluate_ast(expression.step, state, static_tools, custom_tools)
            if expression.step is not None
            else None,
        )
    elif isinstance(expression, ast.DictComp):
        return evaluate_dictcomp(expression, state, static_tools, custom_tools)
@ -834,17 +922,24 @@ def evaluate_ast(
    elif isinstance(expression, ast.With):
        return evaluate_with(expression, state, static_tools, custom_tools)
    elif isinstance(expression, ast.Set):
-        return {evaluate_ast(elt, state, static_tools, custom_tools) for elt in expression.elts}
+        return {
            evaluate_ast(elt, state, static_tools, custom_tools)
            for elt in expression.elts
        }
    elif isinstance(expression, ast.Return):
        raise ReturnException(
-            evaluate_ast(expression.value, state, static_tools, custom_tools) if expression.value else None
+            evaluate_ast(expression.value, state, static_tools, custom_tools)
            if expression.value
            else None
        )
    else:
        # For now we refuse anything else. Let's add things as we need them.
        raise InterpreterError(f"{expression.__class__.__name__} is not supported.")
-def truncate_print_outputs(print_outputs: str, max_len_outputs: int = MAX_LEN_OUTPUT) -> str:
+def truncate_print_outputs(
    print_outputs: str, max_len_outputs: int = MAX_LEN_OUTPUT
 ) -> str:
    if len(print_outputs) < max_len_outputs:
        return print_outputs
    else:
@ -895,8 +990,12 @@ def evaluate_python_code(
    OPERATIONS_COUNT = 0
    try:
        for node in expression.body:
-            result = evaluate_ast(node, state, static_tools, custom_tools, authorized_imports)
+            result = evaluate_ast(
-        state["print_outputs"] = truncate_content(PRINT_OUTPUTS, max_length=MAX_LEN_OUTPUT)
+                node, state, static_tools, custom_tools, authorized_imports
            )
        state["print_outputs"] = truncate_content(
            PRINT_OUTPUTS, max_length=MAX_LEN_OUTPUT
        )
        return result
    except InterpreterError as e:
        msg = truncate_content(PRINT_OUTPUTS, max_length=MAX_LEN_OUTPUT)
--- a/agents/search.py
+++ b/agents/search.py
@ -26,7 +26,9 @@ class DuckDuckGoSearchTool(Tool):
    name = "web_search"
    description = """Perform a web search based on your query (think a Google search) then returns the top search results as a list of dict elements.
    Each result has keys 'title', 'href' and 'body'."""
-    inputs = {"query": {"type": "string", "description": "The search query to perform."}}
+    inputs = {
        "query": {"type": "string", "description": "The search query to perform."}
    }
    output_type = "any"
    def forward(self, query: str) -> str:
--- a/agents/tools.py
+++ b/agents/tools.py
@ -26,7 +26,13 @@ from functools import lru_cache, wraps
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Union
-from huggingface_hub import create_repo, get_collection, hf_hub_download, metadata_update, upload_folder
+from huggingface_hub import (
    create_repo,
    get_collection,
    hf_hub_download,
    metadata_update,
    upload_folder,
 )
 from huggingface_hub.utils import RepositoryNotFoundError, build_hf_headers, get_session
 from packaging import version
@ -73,7 +79,9 @@ def get_repo_type(repo_id, repo_type=None, **hub_kwargs):
            hf_hub_download(repo_id, TOOL_CONFIG_FILE, repo_type="model", **hub_kwargs)
            return "model"
        except RepositoryNotFoundError:
-            raise EnvironmentError(f"`{repo_id}` does not seem to be a valid repo identifier on the Hub.")
+            raise EnvironmentError(
                f"`{repo_id}` does not seem to be a valid repo identifier on the Hub."
            )
        except Exception:
            return "model"
    except Exception:
@ -158,7 +166,15 @@ class Tool:
            "inputs": dict,
            "output_type": str,
        }
-        authorized_types = ["string", "integer", "number", "image", "audio", "any", "boolean"]
+        authorized_types = [
            "string",
            "integer",
            "number",
            "image",
            "audio",
            "any",
            "boolean",
        ]
        for attr, expected_type in required_attributes.items():
            attr_value = getattr(self, attr, None)
@ -169,7 +185,9 @@ class Tool:
                    f"Attribute {attr} should have type {expected_type.__name__}, got {type(attr_value)} instead."
                )
        for input_name, input_content in self.inputs.items():
-            assert isinstance(input_content, dict), f"Input '{input_name}' should be a dictionary."
+            assert isinstance(
                input_content, dict
            ), f"Input '{input_name}' should be a dictionary."
            assert (
                "type" in input_content and "description" in input_content
            ), f"Input '{input_name}' should have keys 'type' and 'description', has only {list(input_content.keys())}."
@ -251,7 +269,11 @@ class Tool:
        # Save app file
        app_file = os.path.join(output_dir, "app.py")
        with open(app_file, "w", encoding="utf-8") as f:
-            f.write(APP_FILE_TEMPLATE.format(module_name=last_module, class_name=self.__class__.__name__))
+            f.write(
                APP_FILE_TEMPLATE.format(
                    module_name=last_module, class_name=self.__class__.__name__
                )
            )
        # Save requirements file
        requirements_file = os.path.join(output_dir, "requirements.txt")
@ -343,7 +365,9 @@ class Tool:
            custom_tool = config
        tool_class = custom_tool["tool_class"]
-        tool_class = get_class_from_dynamic_module(tool_class, repo_id, token=token, **hub_kwargs)
+        tool_class = get_class_from_dynamic_module(
            tool_class, repo_id, token=token, **hub_kwargs
        )
        if len(tool_class.name) == 0:
            tool_class.name = custom_tool["name"]
@ -420,7 +444,9 @@ class Tool:
        with tempfile.TemporaryDirectory() as work_dir:
            # Save all files.
            self.save(work_dir)
-            logger.info(f"Uploading the following files to {repo_id}: {','.join(os.listdir(work_dir))}")
+            logger.info(
                f"Uploading the following files to {repo_id}: {','.join(os.listdir(work_dir))}"
            )
            return upload_folder(
                repo_id=repo_id,
                commit_message=commit_message,
@ -432,7 +458,11 @@ class Tool:
    @staticmethod
    def from_space(
-        space_id: str, name: str, description: str, api_name: Optional[str] = None, token: Optional[str] = None
+        space_id: str,
        name: str,
        description: str,
        api_name: Optional[str] = None,
        token: Optional[str] = None,
    ):
        """
        Creates a [`Tool`] from a Space given its id on the Hub.
@ -485,7 +515,9 @@ class Tool:
                self.client = Client(space_id, hf_token=token)
                self.name = name
                self.description = description
-                space_description = self.client.view_api(return_format="dict", print_info=False)["named_endpoints"]
+                space_description = self.client.view_api(
                    return_format="dict", print_info=False
                )["named_endpoints"]
                # If api_name is not defined, take the first of the available APIs for this space
                if api_name is None:
@ -498,7 +530,9 @@ class Tool:
                try:
                    space_description_api = space_description[api_name]
                except KeyError:
-                    raise KeyError(f"Could not find specified {api_name=} among available api names.")
+                    raise KeyError(
                        f"Could not find specified {api_name=} among available api names."
                    )
                self.inputs = {}
                for parameter in space_description_api["parameters"]:
@ -523,9 +557,11 @@ class Tool:
                    temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
                    arg.save(temp_file.name)
                    arg = temp_file.name
-                if (isinstance(arg, (str, Path)) and Path(arg).exists() and Path(arg).is_file()) or is_http_url_like(
+                if (
-                    arg
+                    isinstance(arg, (str, Path))
-                ):
+                    and Path(arg).exists()
                    and Path(arg).is_file()
                ) or is_http_url_like(arg):
                    arg = handle_file(arg)
                return arg
@ -544,7 +580,9 @@ class Tool:
                    ]  # Sometime the space also returns the generation seed, in which case the result is at index 0
                return output
-        return SpaceToolWrapper(space_id, name, description, api_name=api_name, token=token)
+        return SpaceToolWrapper(
            space_id, name, description, api_name=api_name, token=token
        )
    @staticmethod
    def from_gradio(gradio_tool):
@ -561,7 +599,8 @@ class Tool:
                self._gradio_tool = _gradio_tool
                func_args = list(inspect.signature(_gradio_tool.run).parameters.items())
                self.inputs = {
-                    key: {"type": CONVERSION_DICT[value.annotation], "description": ""} for key, value in func_args
+                    key: {"type": CONVERSION_DICT[value.annotation], "description": ""}
                    for key, value in func_args
                }
                self.forward = self._gradio_tool.run
@ -603,7 +642,9 @@ DEFAULT_TOOL_DESCRIPTION_TEMPLATE = """
 """
-def get_tool_description_with_args(tool: Tool, description_template: str = DEFAULT_TOOL_DESCRIPTION_TEMPLATE) -> str:
+def get_tool_description_with_args(
    tool: Tool, description_template: str = DEFAULT_TOOL_DESCRIPTION_TEMPLATE
 ) -> str:
    compiled_template = compile_jinja_template(description_template)
    rendered = compiled_template.render(
        tool=tool,
@ -621,7 +662,10 @@ def compile_jinja_template(template):
        raise ImportError("template requires jinja2 to be installed.")
    if version.parse(jinja2.__version__) < version.parse("3.1.0"):
-        raise ImportError("template requires jinja2>=3.1.0 to be installed. Your version is " f"{jinja2.__version__}.")
+        raise ImportError(
            "template requires jinja2>=3.1.0 to be installed. Your version is "
            f"{jinja2.__version__}."
        )
    def raise_exception(message):
        raise TemplateError(message)
@ -697,7 +741,9 @@ class PipelineTool(Tool):
        if model is None:
            if self.default_checkpoint is None:
-                raise ValueError("This tool does not implement a default checkpoint, you need to pass one.")
+                raise ValueError(
                    "This tool does not implement a default checkpoint, you need to pass one."
                )
            model = self.default_checkpoint
        if pre_processor is None:
            pre_processor = model
@ -720,15 +766,21 @@ class PipelineTool(Tool):
        Instantiates the `pre_processor`, `model` and `post_processor` if necessary.
        """
        if isinstance(self.pre_processor, str):
-            self.pre_processor = self.pre_processor_class.from_pretrained(self.pre_processor, **self.hub_kwargs)
+            self.pre_processor = self.pre_processor_class.from_pretrained(
                self.pre_processor, **self.hub_kwargs
            )
        if isinstance(self.model, str):
-            self.model = self.model_class.from_pretrained(self.model, **self.model_kwargs, **self.hub_kwargs)
+            self.model = self.model_class.from_pretrained(
                self.model, **self.model_kwargs, **self.hub_kwargs
            )
        if self.post_processor is None:
            self.post_processor = self.pre_processor
        elif isinstance(self.post_processor, str):
-            self.post_processor = self.post_processor_class.from_pretrained(self.post_processor, **self.hub_kwargs)
+            self.post_processor = self.post_processor_class.from_pretrained(
                self.post_processor, **self.hub_kwargs
            )
        if self.device is None:
            if self.device_map is not None:
@ -768,8 +820,12 @@ class PipelineTool(Tool):
        encoded_inputs = self.encode(*args, **kwargs)
-        tensor_inputs = {k: v for k, v in encoded_inputs.items() if isinstance(v, torch.Tensor)}
+        tensor_inputs = {
-        non_tensor_inputs = {k: v for k, v in encoded_inputs.items() if not isinstance(v, torch.Tensor)}
+            k: v for k, v in encoded_inputs.items() if isinstance(v, torch.Tensor)
        }
        non_tensor_inputs = {
            k: v for k, v in encoded_inputs.items() if not isinstance(v, torch.Tensor)
        }
        encoded_inputs = send_to_device(tensor_inputs, self.device)
        outputs = self.forward({**encoded_inputs, **non_tensor_inputs})
@ -790,7 +846,9 @@ def launch_gradio_demo(tool_class: Tool):
    try:
        import gradio as gr
    except ImportError:
-        raise ImportError("Gradio should be installed in order to launch a gradio demo.")
+        raise ImportError(
            "Gradio should be installed in order to launch a gradio demo."
        )
    tool = tool_class()
@ -807,11 +865,15 @@ def launch_gradio_demo(tool_class: Tool):
    gradio_inputs = []
    for input_name, input_details in tool_class.inputs.items():
-        input_gradio_component_class = TYPE_TO_COMPONENT_CLASS_MAPPING[input_details["type"]]
+        input_gradio_component_class = TYPE_TO_COMPONENT_CLASS_MAPPING[
            input_details["type"]
        ]
        new_component = input_gradio_component_class(label=input_name)
        gradio_inputs.append(new_component)
-    output_gradio_componentclass = TYPE_TO_COMPONENT_CLASS_MAPPING[tool_class.output_type]
+    output_gradio_componentclass = TYPE_TO_COMPONENT_CLASS_MAPPING[
        tool_class.output_type
    ]
    gradio_output = output_gradio_componentclass(label=input_name)
    gr.Interface(
@ -875,7 +937,9 @@ def load_tool(task_or_repo_id, model_repo_id=None, token=None, **kwargs):
            f"the tools that you load. We recommend specifying a `revision` to ensure you're loading the "
            f"code that you have checked."
        )
-        return Tool.from_hub(task_or_repo_id, model_repo_id=model_repo_id, token=token, **kwargs)
+        return Tool.from_hub(
            task_or_repo_id, model_repo_id=model_repo_id, token=token, **kwargs
        )
 def add_description(description):
@ -935,7 +999,9 @@ class EndpointClient:
            payload["parameters"] = params
        # Make API call
-        response = get_session().post(self.endpoint_url, headers=self.headers, json=payload, data=data)
+        response = get_session().post(
            self.endpoint_url, headers=self.headers, json=payload, data=data
        )
        # By default, parse the response for the user.
        if output_image:
@ -972,7 +1038,9 @@ class ToolCollection:
    def __init__(self, collection_slug: str, token: Optional[str] = None):
        self._collection = get_collection(collection_slug, token=token)
-        self._hub_repo_ids = {item.item_id for item in self._collection.items if item.item_type == "space"}
+        self._hub_repo_ids = {
            item.item_id for item in self._collection.items if item.item_type == "space"
        }
        self.tools = {Tool.from_hub(repo_id) for repo_id in self._hub_repo_ids}
@ -986,7 +1054,9 @@ def tool(tool_function: Callable) -> Tool:
    """
    parameters = get_json_schema(tool_function)["function"]
    if "return" not in parameters:
-        raise TypeHintParsingException("Tool return type not found: make sure your function has a return type hint!")
+        raise TypeHintParsingException(
            "Tool return type not found: make sure your function has a return type hint!"
        )
    class_name = f"{parameters['name'].capitalize()}Tool"
    class SpecificTool(Tool):
@ -1000,9 +1070,9 @@ def tool(tool_function: Callable) -> Tool:
            return tool_function(*args, **kwargs)
    original_signature = inspect.signature(tool_function)
-    new_parameters = [inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD)] + list(
+    new_parameters = [
-        original_signature.parameters.values()
+        inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD)
-    )
+    ] + list(original_signature.parameters.values())
    new_signature = original_signature.replace(parameters=new_parameters)
    SpecificTool.forward.__signature__ = new_signature
@ -1049,7 +1119,10 @@ class Toolbox:
                The template to use to describe the tools. If not provided, the default template will be used.
        """
        return "\n".join(
-            [get_tool_description_with_args(tool, tool_description_template) for tool in self._tools.values()]
+            [
                get_tool_description_with_args(tool, tool_description_template)
                for tool in self._tools.values()
            ]
        )
    def add_tool(self, tool: Tool):
--- a/agents/utils.py
+++ b/agents/utils.py
@ -16,32 +16,29 @@
 # limitations under the License.
 import json
 import re
-from typing import Tuple, Dict
+from typing import Tuple, Dict, Union
 from transformers.utils.import_utils import _is_package_available
 _pygments_available = _is_package_available("pygments")
 def is_pygments_available():
    return _pygments_available
 from rich.console import Console
 console = Console()
 LENGTH_TRUNCATE_REPORTS = 10000
 def truncate_content(content: str, max_length: int = LENGTH_TRUNCATE_REPORTS):
    if len(content) < max_length:
        return content
    else:
        return content[:max_length//2] + "\n..._(Content was truncated because too long)_...\n---" + content[-max_length//2:]
 def parse_json_blob(json_blob: str) -> Dict[str, str]:
    try:
        first_accolade_index = json_blob.find("{")
        last_accolade_index = [a.start() for a in list(re.finditer("}", json_blob))][-1]
-        json_blob = json_blob[first_accolade_index : last_accolade_index + 1].replace('\\"', "'")
+        json_blob = json_blob[first_accolade_index : last_accolade_index + 1].replace(
            '\\"', "'"
        )
        json_data = json.loads(json_blob, strict=False)
        return json_data
    except json.JSONDecodeError as e:
@ -63,7 +60,12 @@ def parse_code_blob(code_blob: str) -> str:
    try:
        pattern = r"```(?:py|python)?\n(.*?)\n```"
        match = re.search(pattern, code_blob, re.DOTALL)
        if match is None:
            raise ValueError(
                f"No match ground for regex pattern {pattern} in {code_blob=}."
            )
        return match.group(1).strip()
    except Exception as e:
        raise ValueError(
            f"""
@ -77,7 +79,7 @@ Code:
        )
-def parse_json_tool_call(json_blob: str) -> Tuple[str, Dict[str, str]]:
+def parse_json_tool_call(json_blob: str) -> Tuple[str, Union[str, None]]:
    json_blob = json_blob.replace("```json", "").replace("```", "")
    tool_call = parse_json_blob(json_blob)
    if "action" in tool_call and "action_input" in tool_call:
@ -85,7 +87,25 @@ def parse_json_tool_call(json_blob: str) -> Tuple[str, Dict[str, str]]:
    elif "action" in tool_call:
        return tool_call["action"], None
    else:
-        missing_keys = [key for key in ['action', 'action_input'] if key not in tool_call]
+        missing_keys = [
            key for key in ["action", "action_input"] if key not in tool_call
        ]
        error_msg = f"Missing keys: {missing_keys} in blob {tool_call}"
        console.print(f"[bold red]{error_msg}[/bold red]")
        raise ValueError(error_msg)
 MAX_LENGTH_TRUNCATE_CONTENT = 20000
 def truncate_content(
    content: str, max_length: int = MAX_LENGTH_TRUNCATE_CONTENT
 ) -> str:
    if len(content) <= max_length:
        return content
    else:
        return (
            content[: MAX_LENGTH_TRUNCATE_CONTENT // 2]
            + f"\n..._This content has been truncated to stay below {max_length} characters_...\n"
            + content[-MAX_LENGTH_TRUNCATE_CONTENT // 2 :]
        )
--- a/poetry.lock
+++ b/poetry.lock
--- a/requirements.txt
+++ b/requirements.txt
@ -0,0 +1,271 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile requirements.in -o requirements.txt
 aiofiles==23.2.1
    # via gradio
 annotated-types==0.7.0
    # via pydantic
 anyio==4.7.0
    # via
    #   gradio
    #   httpx
    #   starlette
 appnope==0.1.4
    # via ipykernel
 asttokens==3.0.0
    # via stack-data
 beautifulsoup4==4.12.3
    # via markdownify
 certifi==2024.8.30
    # via
    #   httpcore
    #   httpx
    #   requests
 charset-normalizer==3.4.0
    # via requests
 click==8.1.7
    # via
    #   duckduckgo-search
    #   typer
    #   uvicorn
 comm==0.2.2
    # via ipykernel
 debugpy==1.8.10
    # via ipykernel
 decorator==5.1.1
    # via ipython
 diskcache==5.6.3
    # via llama-cpp-python
 duckduckgo-search==6.4.1
    # via -r requirements.in
 executing==2.1.0
    # via stack-data
 fastapi==0.115.6
    # via gradio
 ffmpy==0.4.0
    # via gradio
 filelock==3.16.1
    # via
    #   huggingface-hub
    #   transformers
 fsspec==2024.10.0
    # via
    #   gradio-client
    #   huggingface-hub
 gradio==5.8.0
    # via -r requirements.in
 gradio-client==1.5.1
    # via gradio
 h11==0.14.0
    # via
    #   httpcore
    #   uvicorn
 httpcore==1.0.7
    # via httpx
 httpx==0.28.1
    # via
    #   gradio
    #   gradio-client
    #   safehttpx
 huggingface-hub==0.26.5
    # via
    #   gradio
    #   gradio-client
    #   tokenizers
    #   transformers
 idna==3.10
    # via
    #   anyio
    #   httpx
    #   requests
 iniconfig==2.0.0
    # via pytest
 ipykernel==6.29.5
    # via -r requirements.in
 ipython==8.30.0
    # via ipykernel
 jedi==0.19.2
    # via ipython
 jinja2==3.1.4
    # via
    #   -r requirements.in
    #   gradio
    #   llama-cpp-python
 jupyter-client==8.6.3
    # via ipykernel
 jupyter-core==5.7.2
    # via
    #   ipykernel
    #   jupyter-client
 llama-cpp-python==0.3.5
    # via -r requirements.in
 markdown-it-py==3.0.0
    # via rich
 markdownify==0.14.1
    # via -r requirements.in
 markupsafe==2.1.5
    # via
    #   gradio
    #   jinja2
 matplotlib-inline==0.1.7
    # via
    #   ipykernel
    #   ipython
 mdurl==0.1.2
    # via markdown-it-py
 nest-asyncio==1.6.0
    # via ipykernel
 numpy==2.2.0
    # via
    #   gradio
    #   llama-cpp-python
    #   pandas
    #   transformers
 orjson==3.10.12
    # via gradio
 packaging==24.2
    # via
    #   gradio
    #   gradio-client
    #   huggingface-hub
    #   ipykernel
    #   pytest
    #   transformers
 pandas==2.2.3
    # via
    #   -r requirements.in
    #   gradio
 parso==0.8.4
    # via jedi
 pexpect==4.9.0
    # via ipython
 pillow==11.0.0
    # via
    #   -r requirements.in
    #   gradio
 platformdirs==4.3.6
    # via jupyter-core
 pluggy==1.5.0
    # via pytest
 primp==0.8.3
    # via duckduckgo-search
 prompt-toolkit==3.0.48
    # via ipython
 psutil==6.1.0
    # via ipykernel
 ptyprocess==0.7.0
    # via pexpect
 pure-eval==0.2.3
    # via stack-data
 pydantic==2.10.3
    # via
    #   fastapi
    #   gradio
 pydantic-core==2.27.1
    # via pydantic
 pydub==0.25.1
    # via gradio
 pygments==2.18.0
    # via
    #   ipython
    #   rich
 pytest==8.3.4
    # via -r requirements.in
 python-dateutil==2.9.0.post0
    # via
    #   jupyter-client
    #   pandas
 python-dotenv==1.0.1
    # via -r requirements.in
 python-multipart==0.0.19
    # via gradio
 pytz==2024.2
    # via pandas
 pyyaml==6.0.2
    # via
    #   gradio
    #   huggingface-hub
    #   transformers
 pyzmq==26.2.0
    # via
    #   ipykernel
    #   jupyter-client
 regex==2024.11.6
    # via transformers
 requests==2.32.3
    # via
    #   -r requirements.in
    #   huggingface-hub
    #   transformers
 rich==13.9.4
    # via
    #   -r requirements.in
    #   typer
 ruff==0.8.2
    # via gradio
 safehttpx==0.1.6
    # via gradio
 safetensors==0.4.5
    # via transformers
 semantic-version==2.10.0
    # via gradio
 shellingham==1.5.4
    # via typer
 six==1.17.0
    # via
    #   markdownify
    #   python-dateutil
 sniffio==1.3.1
    # via anyio
 soupsieve==2.6
    # via beautifulsoup4
 stack-data==0.6.3
    # via ipython
 starlette==0.41.3
    # via
    #   fastapi
    #   gradio
 tokenizers==0.21.0
    # via transformers
 tomlkit==0.13.2
    # via gradio
 tornado==6.4.2
    # via
    #   ipykernel
    #   jupyter-client
 tqdm==4.67.1
    # via
    #   huggingface-hub
    #   transformers
 traitlets==5.14.3
    # via
    #   comm
    #   ipykernel
    #   ipython
    #   jupyter-client
    #   jupyter-core
    #   matplotlib-inline
 transformers==4.47.0
    # via -r requirements.in
 typer==0.15.1
    # via gradio
 typing-extensions==4.12.2
    # via
    #   anyio
    #   fastapi
    #   gradio
    #   gradio-client
    #   huggingface-hub
    #   llama-cpp-python
    #   pydantic
    #   pydantic-core
    #   typer
 tzdata==2024.2
    # via pandas
 urllib3==2.2.3
    # via requests
 uvicorn==0.32.1
    # via gradio
 wcwidth==0.2.13
    # via prompt-toolkit
 websockets==14.1
    # via gradio-client
--- a/setup.py
+++ b/setup.py
@ -1,122 +0,0 @@
 # Copyright 2021 The HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from setuptools import find_packages, setup
 extras = {}
 extras["quality"] = [
    "black ~= 23.1",  # hf-doc-builder has a hidden dependency on `black`
    "hf-doc-builder >= 0.3.0",
    "ruff ~= 0.6.4",
 ]
 extras["docs"] = []
 extras["test_prod"] = ["pytest>=7.2.0,<=8.0.0", "pytest-xdist", "pytest-subtests", "parameterized"]
 extras["test_dev"] = [
    "datasets",
    "diffusers",
    "evaluate",
    "torchdata>=0.8.0",
    "torchpippy>=0.2.0",
    "transformers",
    "scipy",
    "scikit-learn",
    "tqdm",
    "bitsandbytes",
    "timm",
 ]
 extras["testing"] = extras["test_prod"] + extras["test_dev"]
 extras["deepspeed"] = ["deepspeed"]
 extras["rich"] = ["rich"]
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard", "dvclive"]
 extras["dev"] = extras["quality"] + extras["testing"] + extras["rich"]
 extras["sagemaker"] = [
    "sagemaker",  # boto3 is a required package in sagemaker
 ]
 setup(
    name="accelerate",
    version="1.2.0.dev0",
    description="Accelerate",
    long_description=open("README.md", encoding="utf-8").read(),
    long_description_content_type="text/markdown",
    keywords="deep learning",
    license="Apache",
    author="The HuggingFace team",
    author_email="zach.mueller@huggingface.co",
    url="https://github.com/huggingface/accelerate",
    package_dir={"": "src"},
    packages=find_packages("src"),
    entry_points={
        "console_scripts": [
            "accelerate=accelerate.commands.accelerate_cli:main",
            "accelerate-config=accelerate.commands.config:main",
            "accelerate-estimate-memory=accelerate.commands.estimate:main",
            "accelerate-launch=accelerate.commands.launch:main",
            "accelerate-merge-weights=accelerate.commands.merge:main",
        ]
    },
    python_requires=">=3.9.0",
    install_requires=[
        "numpy>=1.17,<3.0.0",
        "packaging>=20.0",
        "psutil",
        "pyyaml",
        "torch>=1.10.0",
        "huggingface_hub>=0.21.0",
        "safetensors>=0.4.3",
    ],
    extras_require=extras,
    classifiers=[
        "Development Status :: 5 - Production/Stable",
        "Intended Audience :: Developers",
        "Intended Audience :: Education",
        "Intended Audience :: Science/Research",
        "License :: OSI Approved :: Apache Software License",
        "Operating System :: OS Independent",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.8",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
 )
 # Release checklist
 # 1. Checkout the release branch (for a patch the current release branch, for a new minor version, create one):
 #      git checkout -b vXX.xx-release
 #    The -b is only necessary for creation (so remove it when doing a patch)
 # 2. Change the version in __init__.py and setup.py to the proper value.
 # 3. Commit these changes with the message: "Release: v<VERSION>"
 # 4. Add a tag in git to mark the release:
 #      git tag v<VERSION> -m 'Adds tag v<VERSION> for pypi'
 #    Push the tag and release commit to git: git push --tags origin vXX.xx-release
 # 5. Run the following commands in the top-level directory:
 #      rm -rf dist
 #      rm -rf build
 #      python setup.py bdist_wheel
 #      python setup.py sdist
 # 6. Upload the package to the pypi test server first:
 #      twine upload dist/* -r testpypi
 # 7. Check that you can install it in a virtualenv by running:
 #      pip install accelerate
 #      pip uninstall accelerate
 #      pip install -i https://testpypi.python.org/pypi accelerate
 #      accelerate env
 #      accelerate test
 # 8. Upload the final version to actual pypi:
 #      twine upload dist/* -r pypi
 # 9. Add release notes to the tag in github once everything is looking hunky-dory.
 # 10. Go back to the main branch and update the version in __init__.py, setup.py to the new version ".dev" and push to
 #     main.
--- a/tests/test_agent_types.py
+++ b/tests/test_agent_types.py
@ -19,19 +19,25 @@ import uuid
 from pathlib import Path
 from transformers.agents.agent_types import AgentAudio, AgentImage, AgentText
-from transformers.testing_utils import get_tests_dir, require_soundfile, require_torch, require_vision
+from transformers.testing_utils import (
-from transformers.utils import is_soundfile_availble, is_torch_available, is_vision_available
+    get_tests_dir,
    require_soundfile,
    require_torch,
    require_vision,
 )
 from transformers.utils import (
    is_soundfile_availble,
    is_torch_available,
    is_vision_available,
 )
 import torch
 from PIL import Image
 if is_torch_available():
    import torch
 if is_soundfile_availble():
    import soundfile as sf
 if is_vision_available():
    from PIL import Image
 def get_new_path(suffix="") -> str:
    directory = tempfile.mkdtemp()
--- a/tests/test_agents.py
+++ b/tests/test_agents.py
@ -19,17 +19,16 @@ import uuid
 import pytest
-from transformers.agents.agent_types import AgentText
+from agents.agent_types import AgentText
-from transformers.agents.agents import (
+from agents.agents import (
    AgentMaxIterationsError,
    CodeAgent,
    ManagedAgent,
-    ReactCodeAgent,
+    CodeAgent,
-    ReactJsonAgent,
+    JsonAgent,
    Toolbox,
 )
-from transformers.agents.default_tools import PythonInterpreterTool
+from agents.default_tools import PythonInterpreterTool
 from transformers.testing_utils import require_torch
 def get_new_path(suffix="") -> str:
@ -149,19 +148,26 @@ print(result)
 class AgentTests(unittest.TestCase):
    def test_fake_code_agent(self):
-        agent = CodeAgent(tools=[PythonInterpreterTool()], llm_engine=fake_code_llm_oneshot)
+        agent = CodeAgent(
            tools=[PythonInterpreterTool()], llm_engine=fake_code_llm_oneshot
        )
        output = agent.run("What is 2 multiplied by 3.6452?")
        assert isinstance(output, str)
        assert output == "7.2904"
    def test_fake_react_json_agent(self):
-        agent = ReactJsonAgent(tools=[PythonInterpreterTool()], llm_engine=fake_react_json_llm)
+        agent = JsonAgent(
            tools=[PythonInterpreterTool()], llm_engine=fake_react_json_llm
        )
        output = agent.run("What is 2 multiplied by 3.6452?")
        assert isinstance(output, str)
        assert output == "7.2904"
        assert agent.logs[0]["task"] == "What is 2 multiplied by 3.6452?"
        assert agent.logs[1]["observation"] == "7.2904"
-        assert agent.logs[1]["rationale"].strip() == "Thought: I should multiply 2 by 3.6452. special_marker"
+        assert (
            agent.logs[1]["rationale"].strip()
            == "Thought: I should multiply 2 by 3.6452. special_marker"
        )
        assert (
            agent.logs[2]["llm_output"]
            == """
@ -175,7 +181,9 @@ Action:
        )
    def test_fake_react_code_agent(self):
-        agent = ReactCodeAgent(tools=[PythonInterpreterTool()], llm_engine=fake_react_code_llm)
+        agent = CodeAgent(
            tools=[PythonInterpreterTool()], llm_engine=fake_react_code_llm
        )
        output = agent.run("What is 2 multiplied by 3.6452?")
        assert isinstance(output, float)
        assert output == 7.2904
@ -186,17 +194,19 @@ Action:
        }
    def test_react_code_agent_code_errors_show_offending_lines(self):
-        agent = ReactCodeAgent(tools=[PythonInterpreterTool()], llm_engine=fake_react_code_llm_error)
+        agent = CodeAgent(
            tools=[PythonInterpreterTool()], llm_engine=fake_react_code_llm_error
        )
        output = agent.run("What is 2 multiplied by 3.6452?")
        assert isinstance(output, AgentText)
        assert output == "got an error"
        assert "Evaluation stopped at line 'print = 2' because of" in str(agent.logs)
    def test_setup_agent_with_empty_toolbox(self):
-        ReactJsonAgent(llm_engine=fake_react_json_llm, tools=[])
+        JsonAgent(llm_engine=fake_react_json_llm, tools=[])
    def test_react_fails_max_iterations(self):
-        agent = ReactCodeAgent(
+        agent = CodeAgent(
            tools=[PythonInterpreterTool()],
            llm_engine=fake_code_llm_no_return,  # use this callable because it never ends
            max_iterations=5,
@ -208,51 +218,62 @@ Action:
    @require_torch
    def test_init_agent_with_different_toolsets(self):
        toolset_1 = []
-        agent = ReactCodeAgent(tools=toolset_1, llm_engine=fake_react_code_llm)
+        agent = CodeAgent(tools=toolset_1, llm_engine=fake_react_code_llm)
        assert (
            len(agent.toolbox.tools) == 1
        )  # when no tools are provided, only the final_answer tool is added by default
        toolset_2 = [PythonInterpreterTool(), PythonInterpreterTool()]
-        agent = ReactCodeAgent(tools=toolset_2, llm_engine=fake_react_code_llm)
+        agent = CodeAgent(tools=toolset_2, llm_engine=fake_react_code_llm)
        assert (
            len(agent.toolbox.tools) == 2
        )  # deduplication of tools, so only one python_interpreter tool is added in addition to final_answer
        toolset_3 = Toolbox(toolset_2)
-        agent = ReactCodeAgent(tools=toolset_3, llm_engine=fake_react_code_llm)
+        agent = CodeAgent(tools=toolset_3, llm_engine=fake_react_code_llm)
        assert (
            len(agent.toolbox.tools) == 2
        )  # same as previous one, where toolset_3 is an instantiation of previous one
        # check that add_base_tools will not interfere with existing tools
        with pytest.raises(KeyError) as e:
-            agent = ReactJsonAgent(tools=toolset_3, llm_engine=fake_react_json_llm, add_base_tools=True)
+            agent = JsonAgent(
                tools=toolset_3, llm_engine=fake_react_json_llm, add_base_tools=True
            )
        assert "already exists in the toolbox" in str(e)
        # check that python_interpreter base tool does not get added to code agents
-        agent = ReactCodeAgent(tools=[], llm_engine=fake_react_code_llm, add_base_tools=True)
+        agent = CodeAgent(tools=[], llm_engine=fake_react_code_llm, add_base_tools=True)
-        assert len(agent.toolbox.tools) == 7  # added final_answer tool + 6 base tools (excluding interpreter)
+        assert (
            len(agent.toolbox.tools) == 7
        )  # added final_answer tool + 6 base tools (excluding interpreter)
    def test_function_persistence_across_steps(self):
-        agent = ReactCodeAgent(
+        agent = CodeAgent(
-            tools=[], llm_engine=fake_react_code_functiondef, max_iterations=2, additional_authorized_imports=["numpy"]
+            tools=[],
            llm_engine=fake_react_code_functiondef,
            max_iterations=2,
            additional_authorized_imports=["numpy"],
        )
        res = agent.run("ok")
        assert res[0] == 0.5
    def test_init_managed_agent(self):
-        agent = ReactCodeAgent(tools=[], llm_engine=fake_react_code_functiondef)
+        agent = CodeAgent(tools=[], llm_engine=fake_react_code_functiondef)
        managed_agent = ManagedAgent(agent, name="managed_agent", description="Empty")
        assert managed_agent.name == "managed_agent"
        assert managed_agent.description == "Empty"
    def test_agent_description_gets_correctly_inserted_in_system_prompt(self):
-        agent = ReactCodeAgent(tools=[], llm_engine=fake_react_code_functiondef)
+        agent = CodeAgent(tools=[], llm_engine=fake_react_code_functiondef)
        managed_agent = ManagedAgent(agent, name="managed_agent", description="Empty")
-        manager_agent = ReactCodeAgent(
+        manager_agent = CodeAgent(
-            tools=[], llm_engine=fake_react_code_functiondef, managed_agents=[managed_agent]
+            tools=[],
            llm_engine=fake_react_code_functiondef,
            managed_agents=[managed_agent],
        )
        assert "You can also give requests to team members." not in agent.system_prompt
        assert "<<managed_agents_descriptions>>" not in agent.system_prompt
-        assert "You can also give requests to team members." in manager_agent.system_prompt
+        assert (
            "You can also give requests to team members." in manager_agent.system_prompt
        )
--- a/tests/test_document_question_answering.py
+++ b/tests/test_document_question_answering.py
@ -1,41 +0,0 @@
 # coding=utf-8
 # Copyright 2023 HuggingFace Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import unittest
 from datasets import load_dataset
 from transformers import load_tool
 from .test_tools_common import ToolTesterMixin
 class DocumentQuestionAnsweringToolTester(unittest.TestCase, ToolTesterMixin):
    def setUp(self):
        self.tool = load_tool("document_question_answering")
        self.tool.setup()
    def test_exact_match_arg(self):
        dataset = load_dataset("hf-internal-testing/example-documents", split="test")
        document = dataset[0]["image"]
        result = self.tool(document, "When is the coffee break?")
        self.assertEqual(result, "11-14 to 11:39 a.m.")
    def test_exact_match_kwarg(self):
        dataset = load_dataset("hf-internal-testing/example-documents", split="test")
        document = dataset[0]["image"]
        self.tool(document=document, question="When is the coffee break?")
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@ -87,7 +87,11 @@ class ExampleDifferenceTests(unittest.TestCase):
    examples_path = Path("examples").resolve()
    def one_complete_example(
-        self, complete_file_name: str, parser_only: bool, secondary_filename: str = None, special_strings: list = None
+        self,
        complete_file_name: str,
        parser_only: bool,
        secondary_filename: str = None,
        special_strings: list = None,
    ):
        """
        Tests a single `complete` example against all of the implemented `by_feature` scripts
@ -112,10 +116,15 @@ class ExampleDifferenceTests(unittest.TestCase):
                    with self.subTest(
                        tested_script=complete_file_name,
                        feature_script=item,
-                        tested_section="main()" if parser_only else "training_function()",
+                        tested_section="main()"
                        if parser_only
                        else "training_function()",
                    ):
                        diff = compare_against_test(
-                            self.examples_path / complete_file_name, item_path, parser_only, secondary_filename
+                            self.examples_path / complete_file_name,
                            item_path,
                            parser_only,
                            secondary_filename,
                        )
                        diff = "\n".join(diff)
                        if special_strings is not None:
@ -140,8 +149,12 @@ class ExampleDifferenceTests(unittest.TestCase):
            " " * 12,
            " " * 8 + "for step, batch in enumerate(active_dataloader):\n",
        ]
-        self.one_complete_example("complete_cv_example.py", True, cv_path, special_strings)
+        self.one_complete_example(
-        self.one_complete_example("complete_cv_example.py", False, cv_path, special_strings)
+            "complete_cv_example.py", True, cv_path, special_strings
        )
        self.one_complete_example(
            "complete_cv_example.py", False, cv_path, special_strings
        )
@mock.patch.dict(os.environ, {"TESTING_MOCKED_DATALOADERS": "1"})
--- a/tests/test_final_answer.py
+++ b/tests/test_final_answer.py
@ -47,9 +47,9 @@ class FinalAnswerToolTester(unittest.TestCase, ToolTesterMixin):
    def create_inputs(self):
        inputs_text = {"answer": "Text input"}
        inputs_image = {
-            "answer": Image.open(Path(get_tests_dir("fixtures/tests_samples/COCO")) / "000000039769.png").resize(
+            "answer": Image.open(
-                (512, 512)
+                Path(get_tests_dir("fixtures/tests_samples/COCO")) / "000000039769.png"
-            )
+            ).resize((512, 512))
        }
        inputs_audio = {"answer": torch.Tensor(np.ones(3000))}
        return {"string": inputs_text, "image": inputs_image, "audio": inputs_audio}
--- a/tests/test_image_question_answering.py
+++ b/tests/test_image_question_answering.py
@ -1,42 +0,0 @@
 # coding=utf-8
 # Copyright 2023 HuggingFace Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import unittest
 from pathlib import Path
 from transformers import is_vision_available, load_tool
 from transformers.testing_utils import get_tests_dir
 from .test_tools_common import ToolTesterMixin
 if is_vision_available():
    from PIL import Image
 class ImageQuestionAnsweringToolTester(unittest.TestCase, ToolTesterMixin):
    def setUp(self):
        self.tool = load_tool("image_question_answering")
        self.tool.setup()
    def test_exact_match_arg(self):
        image = Image.open(Path(get_tests_dir("fixtures/tests_samples/COCO")) / "000000039769.png")
        result = self.tool(image, "How many cats are sleeping on the couch?")
        self.assertEqual(result, "2")
    def test_exact_match_kwarg(self):
        image = Image.open(Path(get_tests_dir("fixtures/tests_samples/COCO")) / "000000039769.png")
        result = self.tool(image=image, question="How many cats are sleeping on the couch?")
        self.assertEqual(result, "2")
--- a/tests/test_monitoring.py
+++ b/tests/test_monitoring.py
@ -129,7 +129,9 @@ final_answer('This is the final answer.')
    def test_streaming_agent_image_output(self):
        def dummy_llm_engine(prompt, **kwargs):
-            return 'Action:{"action": "final_answer", "action_input": {"answer": "image"}}'
+            return (
                'Action:{"action": "final_answer", "action_input": {"answer": "image"}}'
            )
        agent = ReactJsonAgent(
            tools=[],
@ -138,7 +140,14 @@ final_answer('This is the final answer.')
        )
        # Use stream_to_gradio to capture the output
-        outputs = list(stream_to_gradio(agent, task="Test task", image=AgentImage(value="path.png"), test_mode=True))
+        outputs = list(
            stream_to_gradio(
                agent,
                task="Test task",
                image=AgentImage(value="path.png"),
                test_mode=True,
            )
        )
        self.assertEqual(len(outputs), 2)
        final_message = outputs[-1]
--- a/tests/test_python_interpreter.py
+++ b/tests/test_python_interpreter.py
@ -21,7 +21,10 @@ import pytest
 from transformers import load_tool
 from transformers.agents.agent_types import AGENT_TYPE_MAPPING
 from transformers.agents.default_tools import BASE_PYTHON_TOOLS
-from transformers.agents.python_interpreter import InterpreterError, evaluate_python_code
+from transformers.agents.python_interpreter import (
    InterpreterError,
    evaluate_python_code,
 )
 from .test_tools_common import ToolTesterMixin
@ -57,7 +60,12 @@ class PythonInterpreterToolTester(unittest.TestCase, ToolTesterMixin):
        for _input, expected_input in zip(inputs, self.tool.inputs.values()):
            input_type = expected_input["type"]
            if isinstance(input_type, list):
-                _inputs.append([AGENT_TYPE_MAPPING[_input_type](_input) for _input_type in input_type])
+                _inputs.append(
                    [
                        AGENT_TYPE_MAPPING[_input_type](_input)
                        for _input_type in input_type
                    ]
                )
            else:
                _inputs.append(AGENT_TYPE_MAPPING[input_type](_input))
@ -91,7 +99,10 @@ class PythonInterpreterTester(unittest.TestCase):
        code = "print = '3'"
        with pytest.raises(InterpreterError) as e:
            evaluate_python_code(code, {"print": print}, state={})
-        assert "Cannot assign to name 'print': doing this would erase the existing tool!" in str(e)
+        assert (
            "Cannot assign to name 'print': doing this would erase the existing tool!"
            in str(e)
        )
    def test_evaluate_call(self):
        code = "y = add_two(x)"
@ -117,7 +128,9 @@ class PythonInterpreterTester(unittest.TestCase):
        state = {"x": 3}
        result = evaluate_python_code(code, {"add_two": add_two}, state=state)
        self.assertDictEqual(result, {"x": 3, "y": 5})
-        self.assertDictEqual(state, {"x": 3, "test_dict": {"x": 3, "y": 5}, "print_outputs": ""})
+        self.assertDictEqual(
            state, {"x": 3, "test_dict": {"x": 3, "y": 5}, "print_outputs": ""}
        )
    def test_evaluate_expression(self):
        code = "x = 3\ny = 5"
@ -133,7 +146,9 @@ class PythonInterpreterTester(unittest.TestCase):
        result = evaluate_python_code(code, {}, state=state)
        # evaluate returns the value of the last assignment.
        assert result == "This is x: 3."
-        self.assertDictEqual(state, {"x": 3, "text": "This is x: 3.", "print_outputs": ""})
+        self.assertDictEqual(
            state, {"x": 3, "text": "This is x: 3.", "print_outputs": ""}
        )
    def test_evaluate_if(self):
        code = "if x <= 3:\n    y = 2\nelse:\n    y = 5"
@ -174,11 +189,15 @@ class PythonInterpreterTester(unittest.TestCase):
        state = {"x": 3}
        result = evaluate_python_code(code, {"add_two": add_two}, state=state)
        assert result == 5
-        self.assertDictEqual(state, {"x": 3, "test_dict": {"x": 3, "y": 5}, "print_outputs": ""})
+        self.assertDictEqual(
            state, {"x": 3, "test_dict": {"x": 3, "y": 5}, "print_outputs": ""}
        )
        code = "vendor = {'revenue': 31000, 'rent': 50312}; vendor['ratio'] = round(vendor['revenue'] / vendor['rent'], 2)"
        state = {}
-        evaluate_python_code(code, {"min": min, "print": print, "round": round}, state=state)
+        evaluate_python_code(
            code, {"min": min, "print": print, "round": round}, state=state
        )
        assert state["vendor"] == {"revenue": 31000, "rent": 50312, "ratio": 0.62}
    def test_subscript_string_with_string_index_raises_appropriate_error(self):
@ -292,7 +311,16 @@ print(check_digits)
 """
        state = {}
        evaluate_python_code(
-            code, {"range": range, "print": print, "sum": sum, "enumerate": enumerate, "int": int, "str": str}, state
+            code,
            {
                "range": range,
                "print": print,
                "sum": sum,
                "enumerate": enumerate,
                "int": int,
                "str": str,
            },
            state,
        )
    def test_listcomp(self):
@ -325,7 +353,9 @@ print(check_digits)
        assert result == {0: 0, 1: 1, 2: 4}
        code = "{num: name for num, name in {101: 'a', 102: 'b'}.items() if name not in ['a']}"
-        result = evaluate_python_code(code, {"print": print}, state={}, authorized_imports=["pandas"])
+        result = evaluate_python_code(
            code, {"print": print}, state={}, authorized_imports=["pandas"]
        )
        assert result == {102: "b"}
        code = """
@ -373,7 +403,9 @@ else:
    best_city = "Manhattan"
    best_city
    """
-        result = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={"a": 1, "b": 2, "c": 3, "d": 4, "e": 5})
+        result = evaluate_python_code(
            code, BASE_PYTHON_TOOLS, state={"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}
        )
        assert result == "Brooklyn"
        code = """if d > e and a < b:
@ -384,7 +416,9 @@ else:
    best_city = "Manhattan"
    best_city
    """
-        result = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={"a": 1, "b": 2, "c": 3, "d": 4, "e": 5})
+        result = evaluate_python_code(
            code, BASE_PYTHON_TOOLS, state={"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}
        )
        assert result == "Sacramento"
    def test_if_conditions(self):
@ -400,7 +434,9 @@ if char.isalpha():
        result = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={})
        assert result == 2.0
-        code = "from random import choice, seed\nseed(12)\nchoice(['win', 'lose', 'draw'])"
+        code = (
            "from random import choice, seed\nseed(12)\nchoice(['win', 'lose', 'draw'])"
        )
        result = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={})
        assert result == "lose"
@ -434,10 +470,14 @@ if char.isalpha():
        # Test submodules are handled properly, thus not raising error
        code = "import numpy.random as rd\nrng = rd.default_rng(12345)\nrng.random()"
-        result = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={}, authorized_imports=["numpy"])
+        result = evaluate_python_code(
            code, BASE_PYTHON_TOOLS, state={}, authorized_imports=["numpy"]
        )
        code = "from numpy.random import default_rng as d_rng\nrng = d_rng(12345)\nrng.random()"
-        result = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={}, authorized_imports=["numpy"])
+        result = evaluate_python_code(
            code, BASE_PYTHON_TOOLS, state={}, authorized_imports=["numpy"]
        )
    def test_additional_imports(self):
        code = "import numpy as np"
@ -554,7 +594,11 @@ cat_sound = cat.sound()
 cat_str = str(cat)
    """
        state = {}
-        evaluate_python_code(code, {"print": print, "len": len, "super": super, "str": str, "sum": sum}, state=state)
+        evaluate_python_code(
            code,
            {"print": print, "len": len, "super": super, "str": str, "sum": sum},
            state=state,
        )
        # Assert results
        assert state["dog1_sound"] == "The dog barks."
@ -588,7 +632,11 @@ except ValueError as e:
    exception_message = str(e)
    """
        state = {}
-        evaluate_python_code(code, {"print": print, "len": len, "super": super, "str": str, "sum": sum}, state=state)
+        evaluate_python_code(
            code,
            {"print": print, "len": len, "super": super, "str": str, "sum": sum},
            state=state,
        )
        assert state["exception_message"] == "An error occurred"
    def test_print(self):
@ -600,7 +648,9 @@ except ValueError as e:
    def test_types_as_objects(self):
        code = "type_a = float(2); type_b = str; type_c = int"
        state = {}
-        result = evaluate_python_code(code, {"float": float, "str": str, "int": int}, state=state)
+        result = evaluate_python_code(
            code, {"float": float, "str": str, "int": int}, state=state
        )
        assert result is int
    def test_tuple_id(self):
@ -731,7 +781,9 @@ def add_one(n, shift):
 add_one(1, 1)
 """
        state = {}
-        result = evaluate_python_code(code, {"print": print, "range": range, "ord": ord, "chr": chr}, state=state)
+        result = evaluate_python_code(
            code, {"print": print, "range": range, "ord": ord, "chr": chr}, state=state
        )
        assert result == 2
        # test returning None
@ -742,7 +794,9 @@ def returns_none(a):
 returns_none(1)
 """
        state = {}
-        result = evaluate_python_code(code, {"print": print, "range": range, "ord": ord, "chr": chr}, state=state)
+        result = evaluate_python_code(
            code, {"print": print, "range": range, "ord": ord, "chr": chr}, state=state
        )
        assert result is None
    def test_nested_for_loop(self):
@ -758,7 +812,9 @@ out = [i for sublist in all_res for i in sublist]
 out[:10]
 """
        state = {}
-        result = evaluate_python_code(code, {"print": print, "range": range}, state=state)
+        result = evaluate_python_code(
            code, {"print": print, "range": range}, state=state
        )
        assert result == [0, 0, 1, 0, 1, 2, 0, 1, 2, 3]
    def test_pandas(self):
@ -773,7 +829,9 @@ parts_with_5_set_count = df[df['SetCount'] == 5.0]
 parts_with_5_set_count[['Quantity', 'SetCount']].values[1]
 """
        state = {}
-        result = evaluate_python_code(code, {}, state=state, authorized_imports=["pandas"])
+        result = evaluate_python_code(
            code, {}, state=state, authorized_imports=["pandas"]
        )
        assert np.array_equal(result, [-1, 5])
        code = """
@ -785,7 +843,9 @@ print("HH0")
 # Filter the DataFrame to get only the rows with outdated atomic numbers
 filtered_df = df.loc[df['AtomicNumber'].isin([104])]
 """
-        result = evaluate_python_code(code, {"print": print}, state={}, authorized_imports=["pandas"])
+        result = evaluate_python_code(
            code, {"print": print}, state={}, authorized_imports=["pandas"]
        )
        assert np.array_equal(result.values[0], [104, 1])
        code = """import pandas as pd
@ -818,7 +878,9 @@ coords_barcelona = (41.3869, 2.1660)
 distance_geneva_barcelona = haversine(*coords_geneva, *coords_barcelona)
 """
-        result = evaluate_python_code(code, {"print": print, "map": map}, state={}, authorized_imports=["math"])
+        result = evaluate_python_code(
            code, {"print": print, "map": map}, state={}, authorized_imports=["math"]
        )
        assert round(result, 1) == 622395.4
    def test_for(self):
--- a/tests/test_speech_to_text.py
+++ b/tests/test_speech_to_text.py
@ -1,36 +0,0 @@
 # coding=utf-8
 # Copyright 2023 HuggingFace Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import unittest
 import numpy as np
 from transformers import load_tool
 from .test_tools_common import ToolTesterMixin
 class SpeechToTextToolTester(unittest.TestCase, ToolTesterMixin):
    def setUp(self):
        self.tool = load_tool("speech_to_text")
        self.tool.setup()
    def test_exact_match_arg(self):
        result = self.tool(np.ones(3000))
        self.assertEqual(result, " Thank you.")
    def test_exact_match_kwarg(self):
        result = self.tool(audio=np.ones(3000))
        self.assertEqual(result, " Thank you.")
--- a/tests/test_text_to_speech.py
+++ b/tests/test_text_to_speech.py
@ -1,50 +0,0 @@
 # coding=utf-8
 # Copyright 2023 HuggingFace Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import unittest
 from transformers import load_tool
 from transformers.utils import is_torch_available
 if is_torch_available():
    import torch
 from transformers.testing_utils import require_torch
 from .test_tools_common import ToolTesterMixin
@require_torch
 class TextToSpeechToolTester(unittest.TestCase, ToolTesterMixin):
    def setUp(self):
        self.tool = load_tool("text_to_speech")
        self.tool.setup()
    def test_exact_match_arg(self):
        # SpeechT5 isn't deterministic
        torch.manual_seed(0)
        result = self.tool("hey")
        resulting_tensor = result.to_raw()
        self.assertTrue(len(resulting_tensor.detach().shape) == 1)
        self.assertTrue(resulting_tensor.detach().shape[0] > 1000)
    def test_exact_match_kwarg(self):
        # SpeechT5 isn't deterministic
        torch.manual_seed(0)
        result = self.tool("hey")
        resulting_tensor = result.to_raw()
        self.assertTrue(len(resulting_tensor.detach().shape) == 1)
        self.assertTrue(resulting_tensor.detach().shape[0] > 1000)
--- a/tests/test_tools_common.py
+++ b/tests/test_tools_common.py
@ -20,7 +20,12 @@ import numpy as np
 import pytest
 from transformers import is_torch_available, is_vision_available
-from transformers.agents.agent_types import AGENT_TYPE_MAPPING, AgentAudio, AgentImage, AgentText
+from transformers.agents.agent_types import (
    AGENT_TYPE_MAPPING,
    AgentAudio,
    AgentImage,
    AgentText,
 )
 from transformers.agents.tools import Tool, tool
 from transformers.testing_utils import get_tests_dir
--- a/tests/test_translation.py
+++ b/tests/test_translation.py
@ -32,7 +32,9 @@ class TranslationToolTester(unittest.TestCase, ToolTesterMixin):
        self.assertEqual(result, "- Hé, comment ça va?")
    def test_exact_match_kwarg(self):
-        result = self.tool(text="Hey, what's up?", src_lang="English", tgt_lang="French")
+        result = self.tool(
            text="Hey, what's up?", src_lang="English", tgt_lang="French"
        )
        self.assertEqual(result, "- Hé, comment ça va?")
    def test_call(self):