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

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

import collections
import types
from collections.abc import Callable, Sequence
from collections.abc import Callable as ABCCallable
from dataclasses import MISSING, fields, is_dataclass
from inspect import getdoc
from types import NoneType
from typing import Any, Union, get_args, get_origin

from docstring_parser import parse
from pydantic import BaseModel, Field, create_model

from haystack import logging
from haystack.dataclasses import ChatMessage
from haystack.utils.type_serialization import _is_union_type

logger = logging.getLogger(__name__)


def _unwrap_optional(type_hint: Any) -> Any:
    """
    Unwrap Optional types (i.e. ``X | None`` or ``Optional[X]``) to get the inner type.

    :param type_hint: The type hint to unwrap.
    :returns: The inner type if ``type_hint`` is ``Optional[X]``, otherwise ``type_hint`` unchanged.
    """
    origin = get_origin(type_hint)
    if origin is Union or origin is types.UnionType:
        non_none = [a for a in get_args(type_hint) if a is not NoneType]
        if len(non_none) == 1:
            return non_none[0]
    return type_hint


def _contains_callable_type(type_hint: Any) -> bool:
    """
    Check if a type hint contains a Callable type, including within Union types.

    The purpose of this function is to help identify Callable types so they can
    be skipped during schema generation.

    :param type_hint: The type hint to check.
    :returns: True if the type contains a Callable, False otherwise.
    """
    origin = get_origin(type_hint)

    # Check if it's a Callable type (direct or parameterized)
    if type_hint in (Callable, ABCCallable) or origin in (Callable, ABCCallable):
        return True

    # Recursively check Union types (both typing.Union and types.UnionType for `X | Y` syntax)
    if origin in (Union, types.UnionType):
        return any(_contains_callable_type(arg) for arg in get_args(type_hint))

    return False


# Schema placeholder models for Tool and Toolset
# These are used during JSON schema generation to represent non-serializable types
class _ToolSchemaPlaceholder(BaseModel):
    """Placeholder model representing a Tool for JSON schema generation."""

    name: str = Field(description="Name of the tool")
    description: str = Field(description="Description of the tool")
    parameters: dict[str, Any] = Field(description="JSON schema of the tool parameters")


class _ToolsetSchemaPlaceholder(BaseModel):
    """Placeholder model representing a Toolset for JSON schema generation."""

    tools: list[_ToolSchemaPlaceholder] = Field(description="List of tools in the toolset")


def _get_param_descriptions(method: Callable) -> tuple[str, dict[str, str]]:
    """
    Extracts parameter descriptions from the method's docstring using docstring_parser.

    :param method: The method to extract parameter descriptions from.
    :returns:
        A tuple including the short description of the method and a dictionary mapping parameter names to their
        descriptions.
    """
    docstring = getdoc(method)
    if not docstring:
        return "", {}

    parsed_doc = parse(docstring)
    param_descriptions = {}
    for param in parsed_doc.params:
        if not param.description:
            logger.warning(
                "Missing description for parameter '{arg_name}'. Please add a description in the component's "
                "run() method docstring using the format ':param {arg_name}: <description>'. "
                "This description helps the LLM understand how to use this parameter.",
                arg_name=param.arg_name,
            )
        param_descriptions[param.arg_name] = param.description.strip() if param.description else ""
    return parsed_doc.short_description or "", param_descriptions


def _get_component_param_descriptions(component: Any) -> dict[str, str]:
    """
    Get parameter descriptions from a component, handling both regular Components and SuperComponents.

    For regular components, this extracts descriptions from the run method's docstring.
    For SuperComponents, this extracts descriptions from the underlying pipeline components.

    :param component: The component to extract parameter descriptions from
    :returns: A dictionary mapping parameter names to their descriptions
    """
    from haystack.core.super_component.super_component import _SuperComponent

    # Get descriptions from the component's run method
    _, param_descriptions = _get_param_descriptions(component.run)

    # If it's a SuperComponent, enhance the parameter descriptions from the original components
    if isinstance(component, _SuperComponent):
        for super_param_name, pipeline_paths in component.input_mapping.items():
            # Collect descriptions from all mapped components
            descriptions = []
            for path in pipeline_paths:
                try:
                    # Get the component and socket this input is mapped from
                    comp_name, socket_name = component._split_component_path(path)
                    pipeline_component = component.pipeline.get_component(comp_name)

                    # Add parameter description if available
                    _, run_param_descriptions = _get_param_descriptions(pipeline_component.run)
                    if input_param_mapping := run_param_descriptions.get(socket_name):
                        descriptions.append(f"Provided to the '{comp_name}' component as: '{input_param_mapping}'")
                except Exception as e:
                    logger.debug(
                        "Error extracting description for {super_param_name} from {path}: {e}",
                        super_param_name=super_param_name,
                        path=path,
                        e=str(e),
                    )

            # A single SuperComponent input can map to multiple pipeline components, e.g.
            # input_mapping={"combined_input": ["comp_a.query", "comp_b.text"]}
            if descriptions:
                param_descriptions[super_param_name] = ", and ".join(descriptions) + "."

    return param_descriptions


def _dataclass_to_pydantic_model(dc_type: Any) -> type[BaseModel]:
    """
    Convert a Python dataclass to an equivalent Pydantic model.

    :param dc_type: The dataclass type to convert.
    :returns:
        A dynamically generated Pydantic model class with fields and types derived from the dataclass definition.
        Field descriptions are extracted from docstrings when available.
    """
    _, param_descriptions = _get_param_descriptions(dc_type)
    cls = dc_type if isinstance(dc_type, type) else dc_type.__class__

    field_defs: dict[str, Any] = {}
    for field in fields(dc_type):
        f_type = field.type if isinstance(field.type, str) else _resolve_type(field.type)
        default = field.default if field.default is not MISSING else ...
        default = field.default_factory() if callable(field.default_factory) else default

        # Special handling for ChatMessage since pydantic doesn't allow for field names with leading underscores
        field_name = field.name
        if dc_type is ChatMessage and field_name.startswith("_"):
            # We remove the underscore since ChatMessage.from_dict does allow for field names without the underscore
            field_name = field_name[1:]

        description = param_descriptions.get(field_name, f"Field '{field_name}' of '{cls.__name__}'.")
        field_defs[field_name] = (f_type, Field(default, description=description))

    return create_model(cls.__name__, **field_defs)


def _resolve_type(_type: Any) -> Any:  # noqa: PLR0911
    """
    Recursively resolve and convert complex type annotations, transforming dataclasses into Pydantic-compatible types.

    This function walks through nested type annotations (e.g., List, Dict, Union) and converts any dataclass types
    it encounters into corresponding Pydantic models.

    :param _type: The type annotation to resolve. If the type is a dataclass, it will be converted to a Pydantic model.
        For generic types (like list[SomeDataclass]), the inner types are also resolved recursively.

    :returns:
        A fully resolved type, with all dataclass types converted to Pydantic models
    """
    # Special handling for Tool and Toolset types - replace with schema placeholders
    # These types contain Callables which cannot be serialized to JSON Schema
    from haystack.tools.tool import Tool
    from haystack.tools.toolset import Toolset

    if _type is Tool:
        return _ToolSchemaPlaceholder

    if _type is Toolset:
        return _ToolsetSchemaPlaceholder

    if is_dataclass(_type):
        return _dataclass_to_pydantic_model(_type)

    origin = get_origin(_type)
    args = get_args(_type)

    if origin is list:
        return list[_resolve_type(args[0]) if args else Any]  # type: ignore[misc]

    if origin is collections.abc.Sequence:
        return Sequence[_resolve_type(args[0]) if args else Any]  # type: ignore[misc]

    if _is_union_type(origin):
        return Union[tuple(_resolve_type(a) for a in args)]

    if origin is dict:
        return dict[args[0] if args else Any, _resolve_type(args[1]) if args else Any]  # type: ignore[misc]

    return _type