diff --git a/haystack/components/agents/tool_calling.py b/haystack/components/agents/tool_calling.py index b8c66c1854..8031681789 100644 --- a/haystack/components/agents/tool_calling.py +++ b/haystack/components/agents/tool_calling.py @@ -4,7 +4,6 @@ import asyncio import contextvars -import inspect import json from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor @@ -13,6 +12,7 @@ from haystack import logging, tracing from haystack.components.agents.state.state import State from haystack.core.component.sockets import Sockets +from haystack.core.type_utils import _resolve_parameter_types from haystack.dataclasses import ChatMessage, ToolCall from haystack.dataclasses.streaming_chunk import StreamingCallbackT, StreamingChunk, _invoke_streaming_callback from haystack.tools import ComponentTool, Tool, ToolsType, _check_duplicate_tool_names, flatten_tools_or_toolsets @@ -271,7 +271,7 @@ def _get_func_params(tool: Tool) -> dict[str, Any]: return {name: socket.type for name, socket in tool._component.__haystack_input__._sockets_dict.items()} # Tool.__post_init__ guarantees that at least one of `function` / `async_function` is set. target = tool.function if tool.function is not None else tool.async_function - return {name: param.annotation for name, param in inspect.signature(target).parameters.items()} # type: ignore[arg-type] + return _resolve_parameter_types(target) # type: ignore[arg-type] def _inject_state_args(tool: Tool, llm_args: dict[str, Any], state: State) -> dict[str, Any]: diff --git a/haystack/core/component/component.py b/haystack/core/component/component.py index 16468b7b81..a1a5d84d54 100644 --- a/haystack/core/component/component.py +++ b/haystack/core/component/component.py @@ -74,7 +74,6 @@ """ import inspect -import typing from collections.abc import Callable, Coroutine, Iterator, Mapping from contextlib import contextmanager from contextvars import ContextVar @@ -85,6 +84,7 @@ from haystack import logging from haystack.core.errors import ComponentError +from haystack.core.type_utils import _resolve_parameter_types from .sockets import Sockets from .types import InputSocket, OutputSocket, _empty @@ -234,24 +234,15 @@ def inner(method: Callable[..., Any], sockets: Sockets) -> inspect.Signature: from inspect import Parameter run_signature = inspect.signature(method) - try: - # TypeError is raised if the argument is not of a type that can contain annotations - run_hints = typing.get_type_hints(method) - except TypeError: - run_hints = None + # Resolves the annotations of components using postponed evaluation of annotations, where they are stored + # as strings. + param_types = _resolve_parameter_types(method) for param_name, param_info in run_signature.parameters.items(): if param_name == "self" or param_info.kind in (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD): continue - # We prefer the type annotation from inspect.signature, but if it's a string we need to resolve it - # using the hints. The type annotation can be a string if the component is using postponed evaluation - # of annotations. - annotation = param_info.annotation - if isinstance(annotation, str) and run_hints is not None: - annotation = run_hints.get(param_name, annotation) - - socket_kwargs = {"name": param_name, "type": annotation} + socket_kwargs = {"name": param_name, "type": param_types[param_name]} if param_info.default != Parameter.empty: socket_kwargs["default_value"] = param_info.default diff --git a/haystack/core/type_utils.py b/haystack/core/type_utils.py index cd732467ab..7873fc8fb7 100644 --- a/haystack/core/type_utils.py +++ b/haystack/core/type_utils.py @@ -3,9 +3,11 @@ # SPDX-License-Identifier: Apache-2.0 import collections.abc +import inspect +from collections.abc import Callable from enum import Enum from types import NoneType, UnionType -from typing import Any, Union, get_args, get_origin +from typing import Any, Union, get_args, get_origin, get_type_hints from haystack.dataclasses import ChatMessage @@ -28,6 +30,36 @@ class ConversionStrategy(Enum): ConversionStrategyType = ConversionStrategy | None +def _resolve_parameter_types(target: Callable) -> dict[str, Any]: + """ + Map the parameter names of a callable to their type annotations, resolving postponed annotations. + + A callable defined in a module using `from __future__ import annotations` stores its annotations as strings, which + never match the types they refer to. Only string annotations are looked up in the resolved type hints: for the + others the annotation from the signature is kept. + + :param target: The callable to inspect. + :returns: A dict mapping parameter names to their type annotations. Annotations that cannot be resolved, and + parameters without an annotation, are returned as they appear in the signature. + """ + parameters = inspect.signature(target).parameters + if any(isinstance(param.annotation, str) for param in parameters.values()): + try: + hints = get_type_hints(target) + except Exception: + # TypeError is raised for objects that cannot carry annotations, NameError for names that are not + # importable at runtime. Either way we fall back to the unresolved annotations. + hints = {} + # Non-string annotations are kept as they are written: on Python 3.10 `get_type_hints` widens the annotation + # of a parameter defaulting to `None` into an optional. This was changed in Python 3.11, see + # https://docs.python.org/3/whatsnew/3.11.html#typing. + return { + name: hints.get(name, param.annotation) if isinstance(param.annotation, str) else param.annotation + for name, param in parameters.items() + } + return {name: param.annotation for name, param in parameters.items()} + + def _type_name(type_: Any) -> str: """ Util methods to get a nice readable representation of a type. diff --git a/haystack/hooks/from_function.py b/haystack/hooks/from_function.py index 165fba9bfa..3df5712c46 100644 --- a/haystack/hooks/from_function.py +++ b/haystack/hooks/from_function.py @@ -4,10 +4,11 @@ import inspect from collections.abc import Awaitable, Callable -from typing import Any, cast, get_type_hints +from typing import Any, cast from haystack.components.agents.state.state import State from haystack.core.serialization import default_from_dict, default_to_dict +from haystack.core.type_utils import _resolve_parameter_types from haystack.utils.callable_serialization import deserialize_callable, serialize_callable @@ -18,17 +19,8 @@ def _takes_single_state_argument(function: Callable) -> bool: :param function: The callable to inspect. :returns: True if the signature is a single `State`-annotated parameter, False otherwise. """ - params = list(inspect.signature(function).parameters.values()) - if len(params) != 1: - return False - - try: - # get_type_hints resolves postponed annotations, where `State` is stored as a string. If resolution fails, fall - # back to the raw annotation so validation still raises the ValueError in FunctionHook. - annotation = get_type_hints(function).get(params[0].name, params[0].annotation) - except (NameError, TypeError, AttributeError): - annotation = params[0].annotation - return annotation is State + param_types = list(_resolve_parameter_types(function).values()) + return len(param_types) == 1 and param_types[0] is State class FunctionHook: diff --git a/releasenotes/notes/resolve-postponed-tool-annotations-678fb1b0f6bbf752.yaml b/releasenotes/notes/resolve-postponed-tool-annotations-678fb1b0f6bbf752.yaml new file mode 100644 index 0000000000..46454b3728 --- /dev/null +++ b/releasenotes/notes/resolve-postponed-tool-annotations-678fb1b0f6bbf752.yaml @@ -0,0 +1,7 @@ +--- +fixes: + - | + Tool functions defined in a module using ``from __future__ import annotations`` are now inspected correctly by + ``Agent``. Postponed annotations are stored as strings, so a parameter annotated with ``State`` was not recognized + and the live ``State`` object was not injected into the tool call. The annotations are now resolved before they + are inspected. diff --git a/test/components/agents/test_tool_calling.py b/test/components/agents/test_tool_calling.py index e3a09bee9b..31582f13ba 100644 --- a/test/components/agents/test_tool_calling.py +++ b/test/components/agents/test_tool_calling.py @@ -54,6 +54,10 @@ async def async_weather_function(location): return weather_function(location) +def weather_with_postponed_state_annotation(city: str, state: "State") -> str: + return f"Weather in {city}" + + weather_parameters = {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]} @@ -226,6 +230,18 @@ def function_with_optional_state(city: str, state: State | None = None) -> str: assert args["city"] == "Paris" assert args["state"] is state + def test_inject_state_args_injects_state_object_for_postponed_state_annotation(self): + state_tool = Tool( + name="state_tool", + description="A tool whose annotations are postponed.", + parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + function=weather_with_postponed_state_annotation, + ) + state = State(schema={"city": {"type": str}}, data={"city": "Berlin"}) + args = _inject_state_args(tool=state_tool, llm_args={"city": "Paris"}, state=state) + assert args["city"] == "Paris" + assert args["state"] is state + class TestRunTool: def test_run_with_streaming_callback_finish_reason(self, weather_tool): diff --git a/test/core/future_annotations_functions.py b/test/core/future_annotations_functions.py new file mode 100644 index 0000000000..c46c8b1395 --- /dev/null +++ b/test/core/future_annotations_functions.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +# The functions below must live in a module of their own so that `from __future__ import annotations` applies to them +# and their annotations are stored as strings (postponed evaluation). +from __future__ import annotations + +from haystack.dataclasses import Document + + +def retrieve(query: str, documents: list[Document], top_k: int | None = None) -> list[Document]: + return documents[:top_k] diff --git a/test/core/test_type_utils.py b/test/core/test_type_utils.py index b82929af1c..4ca7502d6c 100644 --- a/test/core/test_type_utils.py +++ b/test/core/test_type_utils.py @@ -4,6 +4,8 @@ import sys from enum import Enum +from functools import partial +from inspect import Parameter from pathlib import Path from typing import Any, Callable, Dict, Iterable, List, Literal, Mapping, Optional, Sequence, Set, Tuple, Union @@ -16,10 +18,12 @@ _contains_type, _convert_value, _get_first_item, + _resolve_parameter_types, _type_name, _types_are_compatible, ) from haystack.dataclasses import ByteStream, ChatMessage, Document, GeneratedAnswer +from test.core.future_annotations_functions import retrieve class Class1: @@ -1165,3 +1169,32 @@ def test_union_in_sender_problem(self): is_compatible, strategy = _types_are_compatible(sender=str | ChatMessage, receiver=str) assert not is_compatible assert strategy is None + + +def unresolvable_annotation(document: "Unimportable") -> None: ... # type: ignore[name-defined] # noqa: F821 + + +class TestResolveParameterTypes: + def test_resolves_postponed_annotations(self): + assert _resolve_parameter_types(retrieve) == {"query": str, "documents": list[Document], "top_k": Optional[int]} + + def test_keeps_eagerly_evaluated_annotations(self): + # `top_k` keeps the annotation it was written with, it is not widened to `int | None` because of its default. + def function(documents: list[Document], top_k: int = None) -> None: ... # type: ignore[assignment] + + assert _resolve_parameter_types(function) == {"documents": list[Document], "top_k": int} + + def test_includes_parameters_without_annotation(self): + def function(query, top_k: int) -> None: ... # type: ignore[no-untyped-def] + + assert _resolve_parameter_types(function) == {"query": Parameter.empty, "top_k": int} + + def test_keeps_unresolvable_annotation_as_string(self): + assert _resolve_parameter_types(unresolvable_annotation) == {"document": "Unimportable"} + + def test_keeps_annotations_of_target_that_cannot_carry_them(self): + # A partial is not a module, class, method or function, so its annotations cannot be resolved at all. + assert _resolve_parameter_types(partial(retrieve, "query")) == { + "documents": "list[Document]", + "top_k": "int | None", + }