Update tools param to Optional[Union[list[Union[Tool, Toolset]], Toolset]]

Author: vblagoje

haystack/components/agents/agent.py

@@ -4,7 +4,7 @@
 
 import inspect
 from dataclasses import dataclass
-from typing import Any, Optional, Union
+from typing import Any, Optional, Union, cast
 
 from haystack import logging, tracing
 from haystack.components.generators.chat.types import ChatGenerator
@@ -25,7 +25,13 @@
 from haystack.dataclasses import ChatMessage, ChatRole
 from haystack.dataclasses.breakpoints import AgentBreakpoint, AgentSnapshot, PipelineSnapshot, ToolBreakpoint
 from haystack.dataclasses.streaming_chunk import StreamingCallbackT, select_streaming_callback
-from haystack.tools import Tool, Toolset, deserialize_tools_or_toolset_inplace, serialize_tools_or_toolset
+from haystack.tools import (
+    Tool,
+    Toolset,
+    deserialize_tools_or_toolset_inplace,
+    flatten_tools_or_toolsets,
+    serialize_tools_or_toolset,
+)
 from haystack.utils import _deserialize_value_with_schema
 from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
 from haystack.utils.deserialization import deserialize_chatgenerator_inplace
@@ -97,7 +103,7 @@ def __init__(
         self,
         *,
         chat_generator: ChatGenerator,
-        tools: Optional[Union[list[Tool], Toolset]] = None,
+        tools: Optional[Union[list[Union[Tool, Toolset]], Toolset]] = None,
         system_prompt: Optional[str] = None,
         exit_conditions: Optional[list[str]] = None,
         state_schema: Optional[dict[str, Any]] = None,
@@ -110,7 +116,7 @@ def __init__(
         Initialize the agent component.
 
         :param chat_generator: An instance of the chat generator that your agent should use. It must support tools.
-        :param tools: List of Tool objects or a Toolset that the agent can use.
+        :param tools: A list of Tool and/or Toolset objects, or a single Toolset that the agent can use.
         :param system_prompt: System prompt for the agent.
         :param exit_conditions: List of conditions that will cause the agent to return.
             Can include "text" if the agent should return when it generates a message without tool calls,
@@ -134,7 +140,7 @@ def __init__(
                 "The Agent component requires a chat generator that supports tools."
             )
 
-        valid_exits = ["text"] + [tool.name for tool in tools or []]
+        valid_exits = ["text"] + [tool.name for tool in flatten_tools_or_toolsets(tools)]
         if exit_conditions is None:
             exit_conditions = ["text"]
         if not all(condition in valid_exits for condition in exit_conditions):
@@ -259,7 +265,7 @@ def _initialize_fresh_execution(
         requires_async: bool,
         *,
         system_prompt: Optional[str] = None,
-        tools: Optional[Union[list[Tool], Toolset, list[str]]] = None,
+        tools: Optional[Union[list[Union[Tool, Toolset]], Toolset, list[str]]] = None,
         **kwargs,
     ) -> _ExecutionContext:
         """
@@ -302,8 +308,8 @@ def _initialize_fresh_execution(
         )
 
     def _select_tools(
-        self, tools: Optional[Union[list[Tool], Toolset, list[str]]] = None
-    ) -> Union[list[Tool], Toolset]:
+        self, tools: Optional[Union[list[Union[Tool, Toolset]], Toolset, list[str]]] = None
+    ) -> Union[list[Union[Tool, Toolset]], Toolset]:
         """
         Select tools for the current run based on the provided tools parameter.
 
@@ -314,32 +320,43 @@ def _select_tools(
             or if any provided tool name is not valid.
         :raises TypeError: If tools is not a list of Tool objects, a Toolset, or a list of tool names (strings).
         """
-        selected_tools: Union[list[Tool], Toolset] = self.tools
-        if isinstance(tools, Toolset) or isinstance(tools, list) and all(isinstance(t, Tool) for t in tools):
-            selected_tools = tools  # type: ignore[assignment] # mypy thinks this could still be list[str]
-        elif isinstance(tools, list) and all(isinstance(t, str) for t in tools):
+        if tools is None:
+            return self.tools
+
+        if isinstance(tools, list) and all(isinstance(t, str) for t in tools):
             if not self.tools:
                 raise ValueError("No tools were configured for the Agent at initialization.")
-            selected_tool_names: list[str] = tools  # type: ignore[assignment] # mypy thinks this could still be list[Tool] or Toolset
-            valid_tool_names = {tool.name for tool in self.tools}
+            available_tools = flatten_tools_or_toolsets(self.tools)
+            selected_tool_names = cast(list[str], tools)
+            valid_tool_names = {tool.name for tool in available_tools}
             invalid_tool_names = {name for name in selected_tool_names if name not in valid_tool_names}
             if invalid_tool_names:
                 raise ValueError(
                     f"The following tool names are not valid: {invalid_tool_names}. "
                     f"Valid tool names are: {valid_tool_names}."
                 )
-            selected_tools = [tool for tool in self.tools if tool.name in selected_tool_names]
-        elif tools is not None:
-            raise TypeError("tools must be a list of Tool objects, a Toolset, or a list of tool names (strings).")
-        return selected_tools
+            return [tool for tool in available_tools if tool.name in selected_tool_names]
+
+        if isinstance(tools, Toolset):
+            return tools
+
+        if isinstance(tools, list) and tools and isinstance(tools[0], Toolset):
+            return cast(list[Union[Tool, Toolset]], tools)
+
+        if isinstance(tools, list) and all(isinstance(t, Tool) for t in tools):
+            return cast(list[Union[Tool, Toolset]], tools)
+
+        raise TypeError(
+            "tools must be a list of Tool and/or Toolset objects, a Toolset, or a list of tool names (strings)."
+        )
 
     def _initialize_from_snapshot(
         self,
         snapshot: AgentSnapshot,
         streaming_callback: Optional[StreamingCallbackT],
         requires_async: bool,
         *,
-        tools: Optional[Union[list[Tool], Toolset, list[str]]] = None,
+        tools: Optional[Union[list[Union[Tool, Toolset]], Toolset, list[str]]] = None,
     ) -> _ExecutionContext:
         """
         Initialize execution context from an AgentSnapshot.
@@ -459,7 +476,7 @@ def run(  # noqa: PLR0915
         break_point: Optional[AgentBreakpoint] = None,
         snapshot: Optional[AgentSnapshot] = None,
         system_prompt: Optional[str] = None,
-        tools: Optional[Union[list[Tool], Toolset, list[str]]] = None,
+        tools: Optional[Union[list[Union[Tool, Toolset]], Toolset, list[str]]] = None,
         **kwargs: Any,
     ) -> dict[str, Any]:
         """
@@ -616,7 +633,7 @@ async def run_async(
         break_point: Optional[AgentBreakpoint] = None,
         snapshot: Optional[AgentSnapshot] = None,
         system_prompt: Optional[str] = None,
-        tools: Optional[Union[list[Tool], Toolset, list[str]]] = None,
+        tools: Optional[Union[list[Union[Tool, Toolset]], Toolset, list[str]]] = None,
         **kwargs: Any,
     ) -> dict[str, Any]:
         """

haystack/components/generators/chat/azure.py

@@ -17,6 +17,7 @@
     Toolset,
     _check_duplicate_tool_names,
     deserialize_tools_or_toolset_inplace,
+    flatten_tools_or_toolsets,
     serialize_tools_or_toolset,
 )
 from haystack.utils import Secret, deserialize_callable, deserialize_secrets_inplace, serialize_callable
@@ -84,7 +85,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         max_retries: Optional[int] = None,
         generation_kwargs: Optional[dict[str, Any]] = None,
         default_headers: Optional[dict[str, str]] = None,
-        tools: Optional[Union[list[Tool], Toolset]] = None,
+        tools: Optional[Union[list[Union[Tool, Toolset]], Toolset]] = None,
         tools_strict: bool = False,
         *,
         azure_ad_token_provider: Optional[Union[AzureADTokenProvider, AsyncAzureADTokenProvider]] = None,
@@ -138,8 +139,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
                   the `response_format` must be a JSON schema and not a Pydantic model.
         :param default_headers: Default headers to use for the AzureOpenAI client.
         :param tools:
-            A list of tools or a Toolset for which the model can prepare calls. This parameter can accept either a
-            list of `Tool` objects or a `Toolset` instance.
+            A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
         :param tools_strict:
             Whether to enable strict schema adherence for tool calls. If set to `True`, the model will follow exactly
             the schema provided in the `parameters` field of the tool definition, but this may increase latency.
@@ -179,7 +179,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         self.default_headers = default_headers or {}
         self.azure_ad_token_provider = azure_ad_token_provider
         self.http_client_kwargs = http_client_kwargs
-        _check_duplicate_tool_names(list(tools or []))
+        _check_duplicate_tool_names(flatten_tools_or_toolsets(tools))
         self.tools = tools
         self.tools_strict = tools_strict
 

haystack/components/generators/chat/hugging_face_api.py

@@ -25,6 +25,7 @@
     Toolset,
     _check_duplicate_tool_names,
     deserialize_tools_or_toolset_inplace,
+    flatten_tools_or_toolsets,
     serialize_tools_or_toolset,
 )
 from haystack.utils import Secret, deserialize_callable, deserialize_secrets_inplace, serialize_callable
@@ -94,7 +95,7 @@ def _convert_hfapi_tool_calls(hfapi_tool_calls: Optional[list["ChatCompletionOut
 
 
 def _convert_tools_to_hfapi_tools(
-    tools: Optional[Union[list[Tool], Toolset]],
+    tools: Optional[Union[list[Union[Tool, Toolset]], Toolset]],
 ) -> Optional[list["ChatCompletionInputTool"]]:
     if not tools:
         return None
@@ -103,7 +104,7 @@ def _convert_tools_to_hfapi_tools(
     parameters_name = "arguments" if hasattr(ChatCompletionInputFunctionDefinition, "arguments") else "parameters"
 
     hf_tools = []
-    for tool in tools:
+    for tool in flatten_tools_or_toolsets(tools):
         hf_tools_args = {"name": tool.name, "description": tool.description, parameters_name: tool.parameters}
 
         hf_tools.append(
@@ -298,7 +299,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         generation_kwargs: Optional[dict[str, Any]] = None,
         stop_words: Optional[list[str]] = None,
         streaming_callback: Optional[StreamingCallbackT] = None,
-        tools: Optional[Union[list[Tool], Toolset]] = None,
+        tools: Optional[Union[list[Union[Tool, Toolset]], Toolset]] = None,
     ):
         """
         Initialize the HuggingFaceAPIChatGenerator instance.
@@ -328,10 +329,10 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         :param streaming_callback:
             An optional callable for handling streaming responses.
         :param tools:
-            A list of tools or a Toolset for which the model can prepare calls.
+            A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
             The chosen model should support tool/function calling, according to the model card.
             Support for tools in the Hugging Face API and TGI is not yet fully refined and you may experience
-            unexpected behavior. This parameter can accept either a list of `Tool` objects or a `Toolset` instance.
+            unexpected behavior.
         """
 
         huggingface_hub_import.check()
@@ -364,7 +365,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
 
         if tools and streaming_callback is not None:
             raise ValueError("Using tools and streaming at the same time is not supported. Please choose one.")
-        _check_duplicate_tool_names(list(tools or []))
+        _check_duplicate_tool_names(flatten_tools_or_toolsets(tools))
 
         # handle generation kwargs setup
         generation_kwargs = generation_kwargs.copy() if generation_kwargs else {}
@@ -423,7 +424,7 @@ def run(
         self,
         messages: list[ChatMessage],
         generation_kwargs: Optional[dict[str, Any]] = None,
-        tools: Optional[Union[list[Tool], Toolset]] = None,
+        tools: Optional[Union[list[Union[Tool, Toolset]], Toolset]] = None,
         streaming_callback: Optional[StreamingCallbackT] = None,
     ):
         """
@@ -452,7 +453,8 @@ def run(
         tools = tools or self.tools
         if tools and self.streaming_callback:
             raise ValueError("Using tools and streaming at the same time is not supported. Please choose one.")
-        _check_duplicate_tool_names(list(tools or []))
+        flat_tools = flatten_tools_or_toolsets(tools)
+        _check_duplicate_tool_names(flat_tools)
 
         # validate and select the streaming callback
         streaming_callback = select_streaming_callback(
@@ -462,9 +464,6 @@ def run(
         if streaming_callback:
             return self._run_streaming(formatted_messages, generation_kwargs, streaming_callback)
 
-        if tools and isinstance(tools, Toolset):
-            tools = list(tools)
-
         hf_tools = _convert_tools_to_hfapi_tools(tools)
 
         return self._run_non_streaming(formatted_messages, generation_kwargs, hf_tools)
@@ -474,7 +473,7 @@ async def run_async(
         self,
         messages: list[ChatMessage],
         generation_kwargs: Optional[dict[str, Any]] = None,
-        tools: Optional[Union[list[Tool], Toolset]] = None,
+        tools: Optional[Union[list[Union[Tool, Toolset]], Toolset]] = None,
         streaming_callback: Optional[StreamingCallbackT] = None,
     ):
         """
@@ -506,17 +505,15 @@ async def run_async(
         tools = tools or self.tools
         if tools and self.streaming_callback:
             raise ValueError("Using tools and streaming at the same time is not supported. Please choose one.")
-        _check_duplicate_tool_names(list(tools or []))
+        flat_tools = flatten_tools_or_toolsets(tools)
+        _check_duplicate_tool_names(flat_tools)
 
         # validate and select the streaming callback
         streaming_callback = select_streaming_callback(self.streaming_callback, streaming_callback, requires_async=True)
 
         if streaming_callback:
             return await self._run_streaming_async(formatted_messages, generation_kwargs, streaming_callback)
 
-        if tools and isinstance(tools, Toolset):
-            tools = list(tools)
-
         hf_tools = _convert_tools_to_hfapi_tools(tools)
 
         return await self._run_non_streaming_async(formatted_messages, generation_kwargs, hf_tools)