fix: deferred annotation pickup in Agent (#12185)

Author: sjrl

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]:

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
 

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.

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:

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.

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):

test/core/future_annotations_functions.py

@@ -0,0 +1,13 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# 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]

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",
+        }