deepset-ai--haystack/haystack/tools/component_tool.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

from collections.abc import Callable
from typing import Any, get_args, get_origin

from pydantic import Field, TypeAdapter, create_model

from haystack import logging
from haystack.components.agents.state.state import State
from haystack.core.component import Component
from haystack.core.serialization import (
    component_from_dict,
    component_to_dict,
    generate_qualified_class_name,
    import_class_by_name,
)
from haystack.tools import Tool
from haystack.tools.errors import SchemaGenerationError
from haystack.tools.from_function import _remove_title_from_schema
from haystack.tools.parameters_schema_utils import (
    _contains_callable_type,
    _get_component_param_descriptions,
    _resolve_type,
    _unwrap_optional,
)
from haystack.tools.tool import (
    _deserialize_outputs_to_state,
    _deserialize_outputs_to_string,
    _serialize_outputs_to_state,
    _serialize_outputs_to_string,
)
from haystack.utils.type_serialization import _is_union_type

logger = logging.getLogger(__name__)


class ComponentTool(Tool):
    """
    A Tool that wraps Haystack components, allowing them to be used as tools by LLMs.

    ComponentTool automatically generates LLM-compatible tool schemas from component input sockets,
    which are derived from the component's `run` method signature and type hints.


    Key features:
    - Automatic LLM tool calling schema generation from component input sockets
    - Type conversion and validation for component inputs
    - Support for types:
        - Dataclasses
        - Lists of dataclasses
        - Basic types (str, int, float, bool, dict)
        - Lists of basic types
    - Automatic name generation from component class name
    - Description extraction from component docstrings

    To use ComponentTool, you first need a Haystack component - either an existing one or a new one you create.
    You can create a ComponentTool from the component by passing the component to the ComponentTool constructor.
    Below is an example of creating a ComponentTool from an existing SerperDevWebSearch component
    from the `serperdev-haystack` integration package (`pip install serperdev-haystack`).

    ## Usage Example:
    <!-- test-ignore -->
    ```python
    from haystack import component
    from haystack.tools import ComponentTool
    from haystack.utils import Secret
    from haystack.components.agents import Agent
    from haystack.components.generators.chat import OpenAIChatGenerator
    from haystack.dataclasses import ChatMessage
    from haystack_integrations.components.websearch.serperdev import SerperDevWebSearch

    # Create a SerperDev search component
    search = SerperDevWebSearch(api_key=Secret.from_env_var("SERPERDEV_API_KEY"), top_k=3)

    # Create a tool from the component
    tool = ComponentTool(
        component=search,
        name="web_search",  # Optional: defaults to "serper_dev_web_search"
        description="Search the web for current information on any topic"  # Optional: defaults to component docstring
    )

    # Create an Agent with an OpenAIChatGenerator and the tool
    agent = Agent(chat_generator=OpenAIChatGenerator(), tools=[tool])

    message = ChatMessage.from_user("Use the web search tool to find information about Nikola Tesla")

    # Run the Agent
    result = agent.run(messages=[message])

    print(result)
    ```

    """

    def __init__(
        self,
        component: Component,
        name: str | None = None,
        description: str | None = None,
        parameters: dict[str, Any] | None = None,
        *,
        outputs_to_string: dict[str, str | Callable[[Any], str]] | None = None,
        inputs_from_state: dict[str, str] | None = None,
        outputs_to_state: dict[str, dict[str, str | Callable]] | None = None,
    ) -> None:
        """
        Create a Tool instance from a Haystack component.

        :param component: The Haystack component to wrap as a tool.
        :param name: Optional name for the tool (defaults to snake_case of component class name).
        :param description: Optional description (defaults to component's docstring).
        :param parameters:
            A JSON schema defining the parameters expected by the Tool.
            Will fall back to the parameters defined in the component's run method signature if not provided.
        :param outputs_to_string:
            Optional dictionary defining how tool outputs should be converted into string(s) or results.
            If not provided, the tool result is converted to a string using a default handler.

            `outputs_to_string` supports two formats:

            1. Single output format - use "source", "handler", and/or "raw_result" at the root level:
                ```python
                {
                    "source": "docs", "handler": format_documents, "raw_result": False
                }
                ```
                - `source`: If provided, only the specified output key is sent to the handler.
                - `handler`: A function that takes the tool output (or the extracted source value) and returns the
                  final result.
                - `raw_result`: If `True`, the result is returned raw without string conversion, but applying the
                   `handler` if provided. This is intended for tools that return images. In this mode, the Tool
                   function or the `handler` function must return a list of `TextContent`/`ImageContent` objects to
                   ensure compatibility with Chat Generators.

            2. Multiple output format - map keys to individual configurations:
                ```python
                {
                    "formatted_docs": {"source": "docs", "handler": format_documents},
                    "summary": {"source": "summary_text", "handler": str.upper}
                }
                ```
                Each key maps to a dictionary that can contain "source" and/or "handler".
                Note that `raw_result` is not supported in the multiple output format.
        :param inputs_from_state:
            Optional dictionary mapping state keys to tool parameter names.
            Example: `{"repository": "repo"}` maps state's "repository" to tool's "repo" parameter.
        :param outputs_to_state:
            Optional dictionary defining how tool outputs map to keys within state as well as optional handlers.
            If the source is provided only the specified output key is sent to the handler.
            Example:
            ```python
            {
                "documents": {"source": "docs", "handler": custom_handler}
            }
            ```
            If the source is omitted the whole tool result is sent to the handler.
            Example:
            ```python
            {
                "documents": {"handler": custom_handler}
            }
            ```
        :raises TypeError: If the object passed is not a Haystack Component instance.
        :raises ValueError: If the component has already been added to a pipeline, or if schema generation fails.
        """
        if not isinstance(component, Component):
            message = (
                f"Object {component!r} is not a Haystack component. "
                "Use ComponentTool only with Haystack component instances."
            )
            raise TypeError(message)

        if getattr(component, "__haystack_added_to_pipeline__", None):
            msg = (
                "Component has been added to a pipeline and can't be used to create a ComponentTool. "
                "Create ComponentTool from a non-pipeline component instead."
            )
            raise ValueError(msg)

        self._unresolved_parameters = parameters
        # Create the tools schema from the component run method parameters
        tool_schema = parameters or self._create_tool_parameters_schema(component, inputs_from_state or {})

        def component_invoker(**kwargs: Any) -> dict[str, Any]:
            """
            Invokes the component using keyword arguments provided by the LLM function calling/tool-generated response.

            :param kwargs: The keyword arguments to invoke the component with.
            :returns: The result of the component invocation.
            """
            input_sockets = component.__haystack_input__._sockets_dict  # type: ignore[attr-defined]
            converted_kwargs = {
                param_name: self._convert_param(param_value, input_sockets[param_name].type)
                for param_name, param_value in kwargs.items()
            }
            logger.debug(
                "Invoking component {component_type} with kwargs: {converted_kwargs}",
                component_type=type(component),
                converted_kwargs=converted_kwargs,
            )
            return dict(component.run(**converted_kwargs))

        async def async_component_invoker(**kwargs: Any) -> dict[str, Any]:
            """
            Asynchronous counterpart of `component_invoker`. Awaits the component's `run_async`.

            :param kwargs: The keyword arguments to invoke the component with.
            :returns: The result of the component invocation.
            """
            input_sockets = component.__haystack_input__._sockets_dict  # type: ignore[attr-defined]
            converted_kwargs = {
                param_name: self._convert_param(param_value, input_sockets[param_name].type)
                for param_name, param_value in kwargs.items()
            }
            logger.debug(
                "Invoking component {component_type} asynchronously with kwargs: {converted_kwargs}",
                component_type=type(component),
                converted_kwargs=converted_kwargs,
            )
            # We know run_async exists at this point b/c we only pass the async invoker if the component has
            # __haystack_supports_async__ = True
            return dict(await component.run_async(**converted_kwargs))  # type: ignore[attr-defined]

        component_supports_async = getattr(component, "__haystack_supports_async__", False)

        # Generate a name for the tool if not provided
        if not name:
            class_name = component.__class__.__name__
            # Convert camelCase/PascalCase to snake_case
            name = "".join(
                [
                    "_" + c.lower() if c.isupper() and i > 0 and not class_name[i - 1].isupper() else c.lower()
                    for i, c in enumerate(class_name)
                ]
            ).lstrip("_")

        description = description or component.__doc__ or name

        # Store component before calling super().__init__() so _get_valid_outputs() can access it
        self._component = component
        self._is_warmed_up = False

        # Create the Tool instance with the component invoker as the function to be called and the schema.
        # When the wrapped component exposes a `run_async`, also pass the async invoker.
        super().__init__(
            name=name,
            description=description,
            parameters=tool_schema,
            function=component_invoker,
            async_function=async_component_invoker if component_supports_async else None,
            inputs_from_state=inputs_from_state,
            outputs_to_state=outputs_to_state,
            outputs_to_string=outputs_to_string,
        )

    def _get_valid_inputs(self) -> set[str]:
        """
        Return valid input parameter names from the component's input sockets.

        Used to validate `inputs_from_state` against the component's actual inputs.
        This ensures users don't reference non-existent component inputs.

        :returns: Set of component input socket names.
        """
        return set(self._component.__haystack_input__._sockets_dict.keys())  # type: ignore[attr-defined]

    def _get_valid_outputs(self) -> set[str]:
        """
        Return valid output names from the component's output sockets.

        Used to validate `outputs_to_state` against the component's actual outputs.
        This ensures users don't reference non-existent component outputs.

        :returns: Set of component output socket names.
        """
        return set(self._component.__haystack_output__._sockets_dict.keys())  # type: ignore[attr-defined]

    def warm_up(self) -> None:
        """
        Prepare the ComponentTool for use.
        """
        if not self._is_warmed_up:
            if hasattr(self._component, "warm_up"):
                self._component.warm_up()
            self._is_warmed_up = True

    def to_dict(self) -> dict[str, Any]:
        """
        Serializes the ComponentTool to a dictionary.
        """
        serialized: dict[str, Any] = {
            "component": component_to_dict(obj=self._component, name=self.name),
            "name": self.name,
            "description": self.description,
            "parameters": self._unresolved_parameters,
            "inputs_from_state": self.inputs_from_state,
            "outputs_to_state": _serialize_outputs_to_state(self.outputs_to_state) if self.outputs_to_state else None,
            "outputs_to_string": _serialize_outputs_to_string(self.outputs_to_string)
            if self.outputs_to_string
            else None,
        }

        return {"type": generate_qualified_class_name(type(self)), "data": serialized}

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ComponentTool":
        """
        Deserializes the ComponentTool from a dictionary.
        """
        inner_data = data["data"]
        component_class = import_class_by_name(inner_data["component"]["type"])
        component = component_from_dict(cls=component_class, data=inner_data["component"], name=inner_data["name"])

        if "outputs_to_state" in inner_data and inner_data["outputs_to_state"]:
            inner_data["outputs_to_state"] = _deserialize_outputs_to_state(inner_data["outputs_to_state"])

        if inner_data.get("outputs_to_string") is not None:
            inner_data["outputs_to_string"] = _deserialize_outputs_to_string(inner_data["outputs_to_string"])

        return cls(
            component=component,
            name=inner_data["name"],
            description=inner_data["description"],
            parameters=inner_data.get("parameters", None),
            outputs_to_string=inner_data.get("outputs_to_string", None),
            inputs_from_state=inner_data.get("inputs_from_state", None),
            outputs_to_state=inner_data.get("outputs_to_state", None),
        )

    def _create_tool_parameters_schema(self, component: Component, inputs_from_state: dict[str, Any]) -> dict[str, Any]:
        """
        Creates an OpenAI tools schema from a component's run method parameters.

        :param component: The component to create the schema from.
        :raises SchemaGenerationError: If schema generation fails
        :returns: OpenAI tools schema for the component's run method parameters.
        """
        param_descriptions = _get_component_param_descriptions(component)

        # collect fields (types and defaults) and descriptions from function parameters
        fields: dict[str, Any] = {}

        for input_name, socket in component.__haystack_input__._sockets_dict.items():  # type: ignore[attr-defined]
            if inputs_from_state is not None and input_name in list(inputs_from_state.values()):
                continue
            input_type = socket.type

            # Skip Callable types since Pydantic cannot generate JSON schemas for them
            if _contains_callable_type(input_type):
                continue

            # Skip State-typed parameters - Agent tool execution injects them at runtime
            if _unwrap_optional(input_type) is State:
                continue

            description = param_descriptions.get(input_name, f"Input '{input_name}' for the component.")

            # if the parameter has not a default value, Pydantic requires an Ellipsis (...)
            # to explicitly indicate that the parameter is required
            default = ... if socket.is_mandatory else socket.default_value
            resolved_type = _resolve_type(input_type)
            fields[input_name] = (resolved_type, Field(default=default, description=description))

        parameters_schema: dict[str, Any] = {}
        try:
            # No `__doc__`: it would surface as a top-level `description` on the parameters schema,
            # which LLM providers ignore. The component description feeds the tool-level description.
            model = create_model(component.run.__name__, **fields)
            parameters_schema = model.model_json_schema()
        except Exception as e:
            raise SchemaGenerationError(
                f"Failed to create JSON schema for the run method of Component '{component.__class__.__name__}'"
            ) from e

        # we don't want to include title keywords in the schema, as they contain redundant information
        # there is no programmatic way to prevent Pydantic from adding them, so we remove them later
        # see https://github.com/pydantic/pydantic/discussions/8504
        _remove_title_from_schema(parameters_schema)

        return parameters_schema

    def _convert_param(self, param_value: Any, param_type: type) -> Any:
        """
        Converts a single parameter value to the expected type.

        :param param_value: The value to convert.
        :param param_type: The expected type of the parameter.

        :returns:
            The converted parameter value.
        """
        # We unwrap optional types so we can support types like messages: list[ChatMessage] | None
        unwrapped_param_type = _unwrap_optional(param_type)

        # We handle union types (e.g. list[ChatMessage] | str) by extracting the list[T] arm that has
        # from_dict, so the conversion below works the same as for plain list[T]. Other arms like str
        # need no special handling and fall through to Pydantic or the plain return at the end.
        if _is_union_type(unwrapped_param_type):
            list_arms = [
                a
                for a in get_args(unwrapped_param_type)
                if get_origin(a) is list and get_args(a) and hasattr(get_args(a)[0], "from_dict")
            ]
            unwrapped_param_type = list_arms[0] if list_arms else unwrapped_param_type

        # We support calling from_dict on target types that have it, even if they are wrapped in a list.
        # This allows us to support lists of dataclasses as well as single dataclass inputs.
        target_type = (
            get_args(unwrapped_param_type)[0] if get_origin(unwrapped_param_type) is list else unwrapped_param_type
        )
        if hasattr(target_type, "from_dict"):
            if isinstance(param_value, list):
                return [target_type.from_dict(item) if isinstance(item, dict) else item for item in param_value]
            if isinstance(param_value, dict):
                return target_type.from_dict(param_value)
            return param_value

        # Use the original type for pydantic validation
        return TypeAdapter(param_type).validate_python(param_value)