From 1cab85a5d5db902ee00c257ce91777d43366aeb3 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:42:37 +0200 Subject: [PATCH] feat: Add support for async functions inside of Tool (#11426) --- haystack/components/agents/tool_calling.py | 74 +++++++---- haystack/tools/component_tool.py | 27 +++- haystack/tools/from_function.py | 9 +- haystack/tools/pipeline_tool.py | 5 + haystack/tools/tool.py | 92 +++++++++++--- ...sync-tool-invocation-f0b95da0d1cff332.yaml | 20 +++ test/components/agents/test_agent.py | 12 +- test/components/agents/test_agent_hitl.py | 1 + test/components/agents/test_tool_calling.py | 72 +++++++++++ test/components/generators/chat/test_azure.py | 1 + .../generators/chat/test_azure_responses.py | 1 + .../generators/chat/test_hugging_face_api.py | 3 + .../chat/test_hugging_face_local.py | 2 + .../components/generators/chat/test_openai.py | 1 + .../generators/chat/test_openai_responses.py | 1 + test/tools/test_component_tool.py | 49 ++++++++ test/tools/test_from_function.py | 40 ++++-- test/tools/test_pipeline_tool.py | 19 +++ test/tools/test_serde_utils.py | 8 +- test/tools/test_tool.py | 115 +++++++++++++++++- 20 files changed, 482 insertions(+), 70 deletions(-) create mode 100644 releasenotes/notes/add-async-tool-invocation-f0b95da0d1cff332.yaml diff --git a/haystack/components/agents/tool_calling.py b/haystack/components/agents/tool_calling.py index 379320c0c9..e7fba1092a 100644 --- a/haystack/components/agents/tool_calling.py +++ b/haystack/components/agents/tool_calling.py @@ -186,12 +186,35 @@ def _make_context_bound_invoke(tool: Tool, args: dict[str, Any]) -> Callable[[], return _runner +def _make_bounded_invoke_async(tool: Tool, args: dict[str, Any], semaphore: asyncio.Semaphore) -> Callable[[], Any]: + """ + Return a zero-arg async callable that awaits `tool.invoke_async(**args)` while holding `semaphore`. + + Concurrency is bounded uniformly across native-async tools and sync-fallback tools (which dispatch + to a worker thread inside `Tool.invoke_async`). ContextVars naturally inherit into child tasks for + the native-async branch, and `asyncio.to_thread` propagates them for the fallback branch. + + Returns a `ToolInvocationError` instead of raising so that gathered executions can collect failures + without aborting the whole batch. + """ + + async def _runner() -> Any: + async with semaphore: + try: + return await tool.invoke_async(**args) + except ToolInvocationError as e: + return e + + return _runner + + def _get_func_params(tool: Tool) -> dict[str, Any]: """ Return parameter names → annotations for a tool's invocation function. - For ComponentTool, this is the annotated input schema defined on the underlying component. - - For regular Tools, this is the function signature of the `function` callable. + - For regular Tools, this is the function signature of the `function` callable, falling back to `async_function` + for async-only tools. :param tool: The tool to inspect. :returns: A dict mapping parameter names to their type annotations. @@ -201,15 +224,17 @@ def _get_func_params(tool: Tool) -> dict[str, Any]: tool._component.__haystack_input__, Sockets ) return {name: socket.type for name, socket in tool._component.__haystack_input__._sockets_dict.items()} - return {name: param.annotation for name, param in inspect.signature(tool.function).parameters.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] def _inject_state_args(tool: Tool, llm_args: dict[str, Any], state: State) -> dict[str, Any]: """ Merge LLM-provided arguments with state-sourced arguments. - LLM args take precedence. State values are pulled in via `inputs_from_state` mappings or - parameter-name matching, then the live State object is injected for any param annotated as State. + LLM args take precedence. State values are pulled in via `inputs_from_state` mappings or parameter-name matching, + then the live State object is injected for any param annotated as State. :param tool: The tool being invoked, used to determine parameter mappings and State injection. :param llm_args: The arguments provided by the LLM, which take precedence over state values. @@ -434,30 +459,27 @@ async def _run_tool_async( if not tool_call_params: return tool_messages, state - loop = asyncio.get_running_loop() - with ThreadPoolExecutor(max_workers=max_workers) as executor: - tasks = [ - loop.run_in_executor(executor, _make_context_bound_invoke(p["tool"], p["args"])) for p in tool_call_params - ] - results = await asyncio.gather(*tasks) + # `max_workers` + Semaphore bounds concurrency for both sync and async tool calls async tools are awaited directly, + # and sync tools are dispatched to a worker thread inside `Tool.invoke_async`. + semaphore = asyncio.Semaphore(max_workers) + tasks = [_make_bounded_invoke_async(p["tool"], p["args"], semaphore)() for p in tool_call_params] + results = await asyncio.gather(*tasks) - for result, tool_call in zip(results, tool_calls, strict=True): - if isinstance(result, ToolInvocationError): - if raise_on_failure: - raise result - logger.error("{error_exception}", error_exception=result) - tool_messages.append(ChatMessage.from_tool(tool_result=str(result), origin=tool_call, error=True)) - else: - tool = tools_with_names[tool_call.tool_name] - _merge_tool_outputs_into_state(tool, result, state) - tool_messages.append( - _build_tool_result_message(result, tool_call, tool, raise_on_failure=raise_on_failure) - ) + for result, tool_call in zip(results, tool_calls, strict=True): + if isinstance(result, ToolInvocationError): + if raise_on_failure: + raise result + logger.error("{error_exception}", error_exception=result) + tool_messages.append(ChatMessage.from_tool(tool_result=str(result), origin=tool_call, error=True)) + else: + tool = tools_with_names[tool_call.tool_name] + _merge_tool_outputs_into_state(tool, result, state) + tool_messages.append(_build_tool_result_message(result, tool_call, tool, raise_on_failure=raise_on_failure)) - if streaming_callback is not None: - await streaming_callback( # type: ignore[misc] - _create_tool_result_streaming_chunk(tool_messages, tool_call) - ) + if streaming_callback is not None: + await streaming_callback( # type: ignore[misc] + _create_tool_result_streaming_chunk(tool_messages, tool_call) + ) # We emit a final empty chunk with finish_reason "tool_call_results" to signal the end of the tool results stream. if tool_messages and streaming_callback is not None: diff --git a/haystack/tools/component_tool.py b/haystack/tools/component_tool.py index 7e090d16a5..f6e11a72c6 100644 --- a/haystack/tools/component_tool.py +++ b/haystack/tools/component_tool.py @@ -201,6 +201,29 @@ class ComponentTool(Tool): ) 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__ @@ -218,12 +241,14 @@ class ComponentTool(Tool): 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 + # 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, diff --git a/haystack/tools/from_function.py b/haystack/tools/from_function.py index baabe23dad..8001e9cdac 100644 --- a/haystack/tools/from_function.py +++ b/haystack/tools/from_function.py @@ -61,7 +61,9 @@ def create_tool_from_function( ``` :param function: - The function to be converted into a Tool. + The function to be converted into a Tool. May be either a regular function (assigned to the + resulting Tool's `function` field) or a coroutine function defined with `async def` (assigned + to `async_function`). The function must include type hints for all parameters. The function is expected to have basic python input types (str, int, float, bool, list, dict, tuple). Other input types may work but are not guaranteed. @@ -177,11 +179,14 @@ def create_tool_from_function( if param_name in schema["properties"]: schema["properties"][param_name]["description"] = param_description + is_async = inspect.iscoroutinefunction(function) + return Tool( name=name or function.__name__, description=tool_description, parameters=schema, - function=function, + function=None if is_async else function, + async_function=function if is_async else None, inputs_from_state=inputs_from_state, outputs_to_state=outputs_to_state, outputs_to_string=outputs_to_string, diff --git a/haystack/tools/pipeline_tool.py b/haystack/tools/pipeline_tool.py index e2368a6e5e..4b5b9d5710 100644 --- a/haystack/tools/pipeline_tool.py +++ b/haystack/tools/pipeline_tool.py @@ -202,6 +202,11 @@ class PipelineTool(ComponentTool): inputs_from_state=inputs_from_state, outputs_to_state=outputs_to_state, ) + # SuperComponent always defines `run_async` (so ComponentTool will return an `async_function`), but at runtime + # it will raise if the wrapped pipeline is not AsyncPipeline. So we clear async_function here when using + # a sync pipeline so `Tool.invoke_async` falls back to running the sync pipeline in a worker thread. + if not isinstance(pipeline, AsyncPipeline): + self.async_function = None self._unresolved_parameters = parameters self._pipeline = pipeline self._input_mapping = input_mapping diff --git a/haystack/tools/tool.py b/haystack/tools/tool.py index bb02b57c40..0a31de086a 100644 --- a/haystack/tools/tool.py +++ b/haystack/tools/tool.py @@ -2,6 +2,7 @@ # # SPDX-License-Identifier: Apache-2.0 +import asyncio import inspect from collections.abc import Callable from dataclasses import asdict, dataclass @@ -35,8 +36,12 @@ class Tool: :param parameters: A JSON schema defining the parameters expected by the Tool. :param function: - The function that will be invoked when the Tool is called. - Must be a synchronous function; async functions are not supported. + The synchronous function invoked by `Tool.invoke`. Must be a regular function — coroutine functions should + be passed to `async_function` instead. Either `function` or `async_function` (or both) must be set. + :param async_function: + Optional coroutine function awaited by `Tool.invoke_async`. When only `async_function` is set, `invoke` raises + a `ToolInvocationError`. When only `function` is set, `invoke_async` falls back to running `function` in a + worker thread via `asyncio.to_thread`. :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. @@ -86,8 +91,10 @@ class Tool: "documents": {"handler": custom_handler} } ``` - :raises ValueError: If `function` is async, if `parameters` is not a valid JSON schema, or if the - `outputs_to_state`, `outputs_to_string`, or `inputs_from_state` configurations are invalid. + :raises ValueError: If neither `function` nor `async_function` is provided, if `function` is a + coroutine function, if `async_function` is not a coroutine function, if `parameters` is not a + valid JSON schema, or if the `outputs_to_state`, `outputs_to_string`, or `inputs_from_state` + configurations are invalid. :raises TypeError: If any configuration value in `outputs_to_state`, `outputs_to_string`, or `inputs_from_state` has the wrong type. """ @@ -95,18 +102,30 @@ class Tool: name: str description: str parameters: dict[str, Any] - function: Callable + function: Callable | None = None outputs_to_string: dict[str, Any] | None = None inputs_from_state: dict[str, str] | None = None outputs_to_state: dict[str, dict[str, Any]] | None = None + async_function: Callable | None = None def __post_init__(self) -> None: # noqa: C901, PLR0912 - # Check that the function is not a coroutine (async function) - if inspect.iscoroutinefunction(self.function): + # At least one of function / async_function must be set. + if self.function is None and self.async_function is None: + raise ValueError(f"Tool '{self.name}' requires at least one of `function` or `async_function` to be set.") + + # `function` must be a regular (sync) function. Coroutine functions belong on `async_function`. + if self.function is not None and inspect.iscoroutinefunction(self.function): raise ValueError( - f"Async functions are not supported as tools. " + f"`function` must be a synchronous function. " f"The function '{self.function.__name__}' is a coroutine function. " - f"Please use a synchronous function instead." + f"Pass it as `async_function` instead." + ) + + # `async_function` must be a coroutine function defined with `async def`. + if self.async_function is not None and not inspect.iscoroutinefunction(self.async_function): + raise ValueError( + f"`async_function` must be a coroutine function defined with `async def`. " + f"Got '{getattr(self.async_function, '__name__', repr(self.async_function))}'." ) # Check that the parameters define a valid JSON schema @@ -212,12 +231,15 @@ class Tool: # Schema properties provide the validated parameter set valid_params: set[str] = set() - # Try to get parameters from function introspection - try: - sig = inspect.signature(self.function) - valid_params.update(sig.parameters.keys()) - except (ValueError, TypeError): - pass # Introspection failed, will rely on schema + # Try to get parameters from function introspection. + # Prefer `function`; fall back to `async_function` for async-only tools. + introspection_target = self.function if self.function is not None else self.async_function + if introspection_target is not None: + try: + sig = inspect.signature(introspection_target) + valid_params.update(sig.parameters.keys()) + except (ValueError, TypeError): + pass # Introspection failed, will rely on schema # Add parameters from schema (union with function params) valid_params.update(self.parameters.get("properties", {}).keys()) @@ -260,8 +282,18 @@ class Tool: def invoke(self, **kwargs: Any) -> Any: """ - Invoke the Tool with the provided keyword arguments. + Invoke the Tool synchronously with the provided keyword arguments. + + :raises ToolInvocationError: If the Tool has no sync `function`, or if the underlying call + raises an exception. """ + if self.function is None: + raise ToolInvocationError( + f"Tool `{self.name}` has no sync `function` and can only be invoked via `invoke_async` " + f"(use `Agent.run_async`).", + tool_name=self.name, + ) + try: result = self.function(**kwargs) except Exception as e: @@ -270,6 +302,25 @@ class Tool: ) from e return result + async def invoke_async(self, **kwargs: Any) -> Any: + """ + Invoke the Tool asynchronously with the provided keyword arguments. + + If `async_function` is set, it is awaited directly. Otherwise the sync `function` is dispatched to a worker + thread via `asyncio.to_thread`, which propagates the current context to the worker. + + :raises ToolInvocationError: If the underlying call raises an exception. + """ + try: + if self.async_function is not None: + return await self.async_function(**kwargs) + # `function` is guaranteed to be set: __post_init__ enforces at least one of the two. + return await asyncio.to_thread(self.function, **kwargs) # type: ignore[arg-type] + except Exception as e: + raise ToolInvocationError( + f"Failed to invoke Tool `{self.name}` with parameters {kwargs}. Error: {e}", tool_name=self.name + ) from e + def to_dict(self) -> dict[str, Any]: """ Serializes the Tool to a dictionary. @@ -278,7 +329,8 @@ class Tool: Dictionary with serialized data. """ data = asdict(self) - data["function"] = serialize_callable(self.function) + data["function"] = serialize_callable(self.function) if self.function is not None else None + data["async_function"] = serialize_callable(self.async_function) if self.async_function is not None else None if self.outputs_to_state is not None: data["outputs_to_state"] = _serialize_outputs_to_state(self.outputs_to_state) @@ -299,7 +351,11 @@ class Tool: Deserialized Tool. """ init_parameters = data["data"] - init_parameters["function"] = deserialize_callable(init_parameters["function"]) + init_parameters["function"] = ( + deserialize_callable(init_parameters["function"]) if init_parameters.get("function") is not None else None + ) + if init_parameters.get("async_function") is not None: + init_parameters["async_function"] = deserialize_callable(init_parameters["async_function"]) if "outputs_to_state" in init_parameters and init_parameters["outputs_to_state"]: init_parameters["outputs_to_state"] = _deserialize_outputs_to_state(init_parameters["outputs_to_state"]) diff --git a/releasenotes/notes/add-async-tool-invocation-f0b95da0d1cff332.yaml b/releasenotes/notes/add-async-tool-invocation-f0b95da0d1cff332.yaml new file mode 100644 index 0000000000..ce095a8db7 --- /dev/null +++ b/releasenotes/notes/add-async-tool-invocation-f0b95da0d1cff332.yaml @@ -0,0 +1,20 @@ +--- +features: + - | + Add native async tool invocation. ``Tool`` now accepts an optional ``async_function`` (a coroutine function) alongside the existing ``function`` field. + When ``Agent.run_async`` invokes a tool, it now awaits the ``async_function`` of Tool if available. + + ``@tool`` / ``create_tool_from_function`` automatically routes ``async def`` callables to ``async_function``, so simply decorating an ``async def`` produces an async tool: + + .. code-block:: python + + from typing import Annotated + from haystack.tools import tool + + @tool + async def weather(city: Annotated[str, "The name of the city"]) -> str: + """Get the weather for a city.""" + ... + + ``ComponentTool`` automatically wires an async invoker for components that define ``run_async`` (``__haystack_supports_async__`` is ``True``). + ``PipelineTool`` enables the async path when wrapping an ``AsyncPipeline`` and falls back to the threaded sync path when wrapping a regular ``Pipeline``. diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py index 496f16e533..3c8fff1023 100644 --- a/test/components/agents/test_agent.py +++ b/test/components/agents/test_agent.py @@ -260,6 +260,7 @@ class TestAgent: "required": ["location"], }, "function": "test_agent.weather_function", + "async_function": None, "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, @@ -343,6 +344,7 @@ class TestAgent: "required": ["location"], }, "function": "test_agent.weather_function", + "async_function": None, "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, @@ -413,6 +415,7 @@ class TestAgent: "required": ["location"], }, "function": "test_agent.weather_function", + "async_function": None, "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, @@ -503,6 +506,7 @@ class TestAgent: "required": ["location"], }, "function": "test_agent.weather_function", + "async_function": None, "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, @@ -1268,7 +1272,7 @@ class TestAgentTracing: assert set(llm_tags) == {"haystack.agent.step.llm.input", "haystack.agent.step.llm.output"} assert ( llm_tags["haystack.agent.step.llm.input"] - == '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]}' # noqa: E501 + == '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]}' # noqa: E501 ) assert ( llm_tags["haystack.agent.step.llm.output"] @@ -1283,7 +1287,7 @@ class TestAgentTracing: _, run_tags = agent_spans[2] assert run_tags == { "haystack.agent.max_steps": 100, - "haystack.agent.tools": '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]', # noqa: E501 + "haystack.agent.tools": '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]', # noqa: E501 "haystack.agent.exit_conditions": '["text"]', "haystack.agent.state_schema": '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}, "step_count": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "token_usage": {"type": "dict[str, typing.Any]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "tool_call_counts": {"type": "dict[str, int]", "handler": "haystack.components.agents.state.state_utils.replace_values"}}', # noqa: E501 "haystack.agent.input": '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}', # noqa: E501 @@ -1376,7 +1380,7 @@ class TestAgentTracing: assert set(llm_tags) == {"haystack.agent.step.llm.input", "haystack.agent.step.llm.output"} assert ( llm_tags["haystack.agent.step.llm.input"] - == '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]}' # noqa: E501 + == '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]}' # noqa: E501 ) assert ( llm_tags["haystack.agent.step.llm.output"] @@ -1389,7 +1393,7 @@ class TestAgentTracing: _, run_tags = agent_spans[2] assert run_tags == { "haystack.agent.max_steps": 100, - "haystack.agent.tools": '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]', # noqa: E501 + "haystack.agent.tools": '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]', # noqa: E501 "haystack.agent.exit_conditions": '["text"]', "haystack.agent.state_schema": '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}, "step_count": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "token_usage": {"type": "dict[str, typing.Any]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "tool_call_counts": {"type": "dict[str, int]", "handler": "haystack.components.agents.state.state_utils.replace_values"}}', # noqa: E501 "haystack.agent.input": '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}', # noqa: E501 diff --git a/test/components/agents/test_agent_hitl.py b/test/components/agents/test_agent_hitl.py index 1944aa0b33..e6fb6db3d8 100644 --- a/test/components/agents/test_agent_hitl.py +++ b/test/components/agents/test_agent_hitl.py @@ -90,6 +90,7 @@ class TestAgent: "type": "object", }, "function": "test_agent_hitl.addition_tool", + "async_function": None, "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, diff --git a/test/components/agents/test_tool_calling.py b/test/components/agents/test_tool_calling.py index de5830cad9..ffc2d5bea4 100644 --- a/test/components/agents/test_tool_calling.py +++ b/test/components/agents/test_tool_calling.py @@ -2,6 +2,7 @@ # # SPDX-License-Identifier: Apache-2.0 +import asyncio import datetime import json import time @@ -48,6 +49,10 @@ def weather_function(location): return weather_info.get(location, {"weather": "unknown", "temperature": 0, "unit": "celsius"}) +async def async_weather_function(location): + return weather_function(location) + + weather_parameters = {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]} @@ -61,6 +66,16 @@ def weather_tool(): ) +@pytest.fixture +def async_weather_tool(): + return Tool( + name="weather_tool", + description="Provides weather information for a given location.", + parameters=weather_parameters, + async_function=async_weather_function, + ) + + @pytest.fixture def weather_tool_with_outputs_to_state(): return Tool( @@ -875,3 +890,60 @@ class TestUtilities: TextContent(text="weather: mostly sunny"), TextContent(text="temperature: 7 celsius"), ] + + +class TestRunToolAsync: + @pytest.mark.asyncio + async def test_async_tool_is_awaited_directly(self, async_weather_tool): + message = ChatMessage.from_assistant( + tool_calls=[ToolCall(id="1", tool_name="weather_tool", arguments={"location": "Berlin"})] + ) + + with patch.object(asyncio, "to_thread") as to_thread_mock: + tool_messages, _ = await _run_tool_async( + messages=[message], state=State(schema={}), tools=[async_weather_tool] + ) + + to_thread_mock.assert_not_called() + assert json.loads(tool_messages[0].tool_call_results[0].result)["weather"] == "mostly sunny" + + @pytest.mark.asyncio + async def test_sync_tool_is_dispatched_to_thread(self, weather_tool): + message = ChatMessage.from_assistant( + tool_calls=[ToolCall(id="1", tool_name="weather_tool", arguments={"location": "Berlin"})] + ) + + with patch.object(asyncio, "to_thread", wraps=asyncio.to_thread) as to_thread_spy: + tool_messages, _ = await _run_tool_async(messages=[message], state=State(schema={}), tools=[weather_tool]) + + assert to_thread_spy.call_count == 1 + assert tool_messages[0].tool_call_results[0].result is not None + + @pytest.mark.asyncio + async def test_max_workers_serializes_async_tool_calls(self): + active = 0 + max_active = 0 + + async def slow_async(location: str) -> str: + nonlocal active, max_active + active += 1 + max_active = max(max_active, active) + await asyncio.sleep(0.01) + active -= 1 + return location + + slow_tool = Tool( + name="weather_tool", + description="Slow async tool.", + parameters=weather_parameters, + async_function=slow_async, + ) + message = ChatMessage.from_assistant( + tool_calls=[ + ToolCall(id=str(i), tool_name="weather_tool", arguments={"location": f"city-{i}"}) for i in range(4) + ] + ) + + await _run_tool_async(messages=[message], state=State(schema={}), tools=[slow_tool], max_workers=1) + + assert max_active == 1 diff --git a/test/components/generators/chat/test_azure.py b/test/components/generators/chat/test_azure.py index 93f5daf11c..edb2bdd452 100644 --- a/test/components/generators/chat/test_azure.py +++ b/test/components/generators/chat/test_azure.py @@ -426,6 +426,7 @@ class TestAzureOpenAIChatGenerator: "required": ["city"], }, "function": "generators.chat.test_azure.get_weather", + "async_function": None, "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, diff --git a/test/components/generators/chat/test_azure_responses.py b/test/components/generators/chat/test_azure_responses.py index 96e5016913..636a99356c 100644 --- a/test/components/generators/chat/test_azure_responses.py +++ b/test/components/generators/chat/test_azure_responses.py @@ -246,6 +246,7 @@ class TestSerDe: "required": ["city"], }, "function": "generators.chat.test_azure_responses.get_weather", + "async_function": None, "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, diff --git a/test/components/generators/chat/test_hugging_face_api.py b/test/components/generators/chat/test_hugging_face_api.py index 1bb91f8e3e..50a0fcb368 100644 --- a/test/components/generators/chat/test_hugging_face_api.py +++ b/test/components/generators/chat/test_hugging_face_api.py @@ -273,6 +273,7 @@ class TestHuggingFaceAPIChatGenerator: { "type": "haystack.tools.tool.Tool", "data": { + "async_function": None, "description": "description", "function": "builtins.print", "inputs_from_state": None, @@ -346,6 +347,7 @@ class TestHuggingFaceAPIChatGenerator: "description": "description", "parameters": {"x": {"type": "string"}}, "function": "builtins.print", + "async_function": None, }, } ], @@ -1212,6 +1214,7 @@ class TestHuggingFaceAPIChatGenerator: "required": ["city"], }, "function": "generators.chat.test_hugging_face_api.get_weather", + "async_function": None, "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, diff --git a/test/components/generators/chat/test_hugging_face_local.py b/test/components/generators/chat/test_hugging_face_local.py index 06069999b5..3f134abd7e 100644 --- a/test/components/generators/chat/test_hugging_face_local.py +++ b/test/components/generators/chat/test_hugging_face_local.py @@ -210,6 +210,7 @@ class TestHuggingFaceLocalChatGenerator: "description": "useful to determine the weather in a given location", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, "function": "generators.chat.test_hugging_face_local.get_weather", + "async_function": None, }, } ] @@ -783,6 +784,7 @@ class TestHuggingFaceLocalChatGeneratorAsync: "required": ["city"], }, "function": "generators.chat.test_hugging_face_local.get_weather", + "async_function": None, "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, diff --git a/test/components/generators/chat/test_openai.py b/test/components/generators/chat/test_openai.py index 9a34f2ca46..af9ac47cbd 100644 --- a/test/components/generators/chat/test_openai.py +++ b/test/components/generators/chat/test_openai.py @@ -341,6 +341,7 @@ class TestOpenAIChatGenerator: { "type": "haystack.tools.tool.Tool", "data": { + "async_function": None, "description": "description", "function": "builtins.print", "inputs_from_state": None, diff --git a/test/components/generators/chat/test_openai_responses.py b/test/components/generators/chat/test_openai_responses.py index e9a2ec2b9a..08800a1fbe 100644 --- a/test/components/generators/chat/test_openai_responses.py +++ b/test/components/generators/chat/test_openai_responses.py @@ -255,6 +255,7 @@ class TestSerDe: { "type": "haystack.tools.tool.Tool", "data": { + "async_function": None, "description": "description", "function": "builtins.print", "inputs_from_state": None, diff --git a/test/tools/test_component_tool.py b/test/tools/test_component_tool.py index 778835d6dd..235916613b 100644 --- a/test/tools/test_component_tool.py +++ b/test/tools/test_component_tool.py @@ -914,8 +914,10 @@ class TestComponentToolInAgent: def test_deepcopy_with_jinja_based_component(self): builder = PromptBuilder("{{query}}") tool = ComponentTool(component=builder) + assert tool.function is not None result = tool.function(query="Hello") tool_copy = _deepcopy_with_exceptions(tool) + assert tool_copy.function is not None result_from_copy = tool_copy.function(query="Hello") assert "prompt" in result_from_copy assert result_from_copy["prompt"] == result["prompt"] @@ -944,3 +946,50 @@ class TestComponentToolInAgent: result = pipeline.run({"llm": {"messages": [ChatMessage.from_user(text="Hello")], "tools": [tool]}}) assert result["llm"]["replies"][0].text == "A response from the model" + + +@component +class SyncOnlyComponent: + @component.output_types(reply=str) + def run(self, text: str) -> dict[str, str]: + return {"reply": f"sync:{text}"} + + +@component +class DualModeComponent: + @component.output_types(reply=str) + def run(self, text: str) -> dict[str, str]: + return {"reply": f"sync:{text}"} + + @component.output_types(reply=str) + async def run_async(self, text: str) -> dict[str, str]: + return {"reply": f"async:{text}"} + + +@pytest.fixture +def sync_tool(): + return ComponentTool(component=SyncOnlyComponent()) + + +@pytest.fixture +def dual_tool(): + return ComponentTool(component=DualModeComponent()) + + +class TestComponentToolAsync: + def test_async_function_is_wired_only_when_component_has_run_async(self, sync_tool, dual_tool): + assert sync_tool.function is not None + assert sync_tool.async_function is None + assert dual_tool.function is not None + assert dual_tool.async_function is not None + + @pytest.mark.asyncio + async def test_invoke_async_uses_run_async_when_available(self, dual_tool): + assert await dual_tool.invoke_async(text="hi") == {"reply": "async:hi"} + + @pytest.mark.asyncio + async def test_invoke_async_falls_back_to_run_for_sync_only_component(self, sync_tool): + assert await sync_tool.invoke_async(text="hi") == {"reply": "sync:hi"} + + def test_invoke_uses_run_on_dual_mode_component(self, dual_tool): + assert dual_tool.invoke(text="hi") == {"reply": "sync:hi"} diff --git a/test/tools/test_from_function.py b/test/tools/test_from_function.py index b091b0be36..48b5563625 100644 --- a/test/tools/test_from_function.py +++ b/test/tools/test_from_function.py @@ -340,19 +340,33 @@ def test_remove_title_from_schema_handle_no_title_in_top_level(): } -def test_from_function_async_raises_error(): - async def async_get_weather(city: str) -> str: - """Get weather report for a city.""" - return f"Weather report for {city}: 20°C, sunny" - - with pytest.raises(ValueError, match="Async functions are not supported as tools"): - create_tool_from_function(async_get_weather) +async def async_function_with_docstring(city: str) -> str: + """Get weather report for a city.""" + return f"Weather report for {city}: 20°C, sunny" -def test_tool_decorator_async_raises_error(): - with pytest.raises(ValueError, match="Async functions are not supported as tools"): +class TestFromFunctionAsync: + def test_create_tool_from_async_function(self): + tool_obj = create_tool_from_function(async_function_with_docstring) - @tool - async def async_get_weather(city: str) -> str: - """Get weather report for a city.""" - return f"Weather report for {city}: 20°C, sunny" + assert tool_obj.function is None + assert tool_obj.async_function is async_function_with_docstring + assert tool_obj.name == "async_function_with_docstring" + assert tool_obj.parameters == { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + } + + def test_tool_decorator_on_async_function(self): + decorated = tool(async_function_with_docstring) + + assert decorated.function is None + assert decorated.async_function is async_function_with_docstring + assert decorated.name == "async_function_with_docstring" + + @pytest.mark.asyncio + async def test_invoke_async(self): + decorated = tool(async_function_with_docstring) + + assert await decorated.invoke_async(city="Berlin") == "Weather report for Berlin: 20°C, sunny" diff --git a/test/tools/test_pipeline_tool.py b/test/tools/test_pipeline_tool.py index f4e8af68d4..5bce0afccb 100644 --- a/test/tools/test_pipeline_tool.py +++ b/test/tools/test_pipeline_tool.py @@ -407,3 +407,22 @@ class TestPipelineTool: description="A test tool", outputs_to_state={"result": {"source": "nonexistent"}}, ) + + +class TestPipelineToolAsync: + # SuperComponent always defines run_async, but PipelineTool clears it for sync pipelines so that + # invoke_async transparently falls back to running the sync pipeline in a thread. + @pytest.mark.parametrize( + "pipeline_fixture, expects_async_function", [("sample_pipeline", False), ("sample_async_pipeline", True)] + ) + def test_async_function_is_set_only_for_async_pipelines(self, request, pipeline_fixture, expects_async_function): + tool = PipelineTool( + pipeline=request.getfixturevalue(pipeline_fixture), + input_mapping={"query": ["bm25_retriever.query"]}, + output_mapping={"ranker.documents": "documents"}, + name="test_tool", + description="A test tool", + ) + + assert tool.function is not None + assert (tool.async_function is not None) is expects_async_function diff --git a/test/tools/test_serde_utils.py b/test/tools/test_serde_utils.py index 6558a963b8..5bf82a8bce 100644 --- a/test/tools/test_serde_utils.py +++ b/test/tools/test_serde_utils.py @@ -288,7 +288,7 @@ class TestToolSerdeUtils: assert isinstance(data["tools"][0], Tool) assert data["tools"][0].name == "weather" assert data["tools"][0].parameters == weather_parameters - assert data["tools"][0].function("Paris") == "Weather report for Paris: 20°C, sunny" + assert data["tools"][0].function("Paris") == "Weather report for Paris: 20°C, sunny" # type: ignore[misc] # Verify Toolset with calculator tool assert isinstance(data["tools"][1], Toolset) @@ -326,14 +326,14 @@ class TestToolSerdeUtils: assert isinstance(data["tools"][0], Tool) assert data["tools"][0].name == "weather" assert data["tools"][0].parameters == weather_parameters - assert data["tools"][0].function("Berlin") == "Weather report for Berlin: 20°C, sunny" + assert data["tools"][0].function("Berlin") == "Weather report for Berlin: 20°C, sunny" # type: ignore[misc] # Verify Tool 2 (calculator) assert isinstance(data["tools"][1], Tool) assert data["tools"][1].name == "calculator" assert data["tools"][1].parameters == calculator_parameters - assert data["tools"][1].function(5, 3, "add") == 8 - assert data["tools"][1].function(5, 3, "multiply") == 15 + assert data["tools"][1].function(5, 3, "add") == 8 # type: ignore[misc] + assert data["tools"][1].function(5, 3, "multiply") == 15 # type: ignore[misc] # Verify Toolset (with summarizer and formatter) assert isinstance(data["tools"][2], Toolset) diff --git a/test/tools/test_tool.py b/test/tools/test_tool.py index 3a8502aff1..5148cf7060 100644 --- a/test/tools/test_tool.py +++ b/test/tools/test_tool.py @@ -49,10 +49,23 @@ class TestTool: with pytest.raises(ValueError): Tool(name="irrelevant", description="irrelevant", parameters=params, function=get_weather_report) - def test_init_async_function_raises_error(self): - with pytest.raises(ValueError, match="Async functions are not supported as tools"): + def test_init_async_function_passed_as_function_raises_error(self): + with pytest.raises(ValueError, match="`function` must be a synchronous function"): Tool(name="weather", description="Get weather report", parameters=parameters, function=async_get_weather) + def test_init_requires_function_or_async_function(self): + with pytest.raises(ValueError, match="requires at least one of `function` or `async_function`"): + Tool(name="weather", description="Get weather report", parameters=parameters) + + def test_init_sync_function_passed_as_async_function_raises_error(self): + with pytest.raises(ValueError, match="`async_function` must be a coroutine function"): + Tool( + name="weather", + description="Get weather report", + parameters=parameters, + async_function=get_weather_report, + ) + @pytest.mark.parametrize( "outputs_to_state", [ @@ -158,6 +171,7 @@ class TestTool: "description": "Get weather report", "parameters": parameters, "function": "test_tool.get_weather_report", + "async_function": None, "outputs_to_string": {"handler": "test_tool.format_string"}, "inputs_from_state": {"location": "city"}, "outputs_to_state": {"documents": {"source": "docs", "handler": "test_tool.get_weather_report"}}, @@ -316,6 +330,103 @@ class TestTool: ) +@pytest.fixture +def async_tool(): + return Tool( + name="weather", description="Get weather report", parameters=parameters, async_function=async_get_weather + ) + + +@pytest.fixture +def sync_tool(): + return Tool(name="weather", description="Get weather report", parameters=parameters, function=get_weather_report) + + +class TestToolAsync: + @pytest.mark.asyncio + async def test_invoke_async_awaits_async_function(self, async_tool): + assert async_tool.function is None + assert async_tool.async_function is async_get_weather + assert await async_tool.invoke_async(city="Berlin") == "Weather report for Berlin: 20°C, sunny" + + def test_invoke_on_async_only_tool_raises(self, async_tool): + with pytest.raises(ToolInvocationError, match=re.escape("has no sync `function`")): + async_tool.invoke(city="Berlin") + + @pytest.mark.asyncio + async def test_invoke_async_falls_back_to_sync_function(self, sync_tool): + # Sync-only tool: invoke_async dispatches to a worker thread via asyncio.to_thread. + assert await sync_tool.invoke_async(city="Berlin") == "Weather report for Berlin: 20°C, sunny" + + @pytest.mark.asyncio + async def test_async_function_is_preferred_when_both_set(self): + async def from_async(city: str) -> str: + return "async" + + def from_sync(city: str) -> str: + return "sync" + + tool = Tool( + name="weather", + description="Get weather report", + parameters=parameters, + function=from_sync, + async_function=from_async, + ) + + assert await tool.invoke_async(city="Berlin") == "async" + assert tool.invoke(city="Berlin") == "sync" + + @pytest.mark.asyncio + async def test_invoke_async_wraps_exception(self): + async def boom(city: str) -> str: + raise RuntimeError("kaboom") + + tool = Tool(name="weather", description="Get weather report", parameters=parameters, async_function=boom) + + with pytest.raises(ToolInvocationError, match="kaboom"): + await tool.invoke_async(city="Berlin") + + @pytest.mark.parametrize( + "kwargs, expected_function, expected_async_function", + [ + pytest.param({"async_function": async_get_weather}, None, "test_tool.async_get_weather", id="async-only"), + pytest.param( + {"function": get_weather_report, "async_function": async_get_weather}, + "test_tool.get_weather_report", + "test_tool.async_get_weather", + id="both", + ), + ], + ) + def test_serde_roundtrip(self, kwargs, expected_function, expected_async_function): + tool = Tool(name="weather", description="Get weather report", parameters=parameters, **kwargs) + + data = tool.to_dict() + assert data["data"]["function"] == expected_function + assert data["data"]["async_function"] == expected_async_function + + restored = Tool.from_dict(data) + assert restored.function == kwargs.get("function") + assert restored.async_function == kwargs.get("async_function") + + def test_from_dict_legacy_payload_without_async_function_key(self): + # Payload produced before `async_function` existed. + legacy = { + "type": "haystack.tools.tool.Tool", + "data": { + "name": "weather", + "description": "Get weather report", + "parameters": parameters, + "function": "test_tool.get_weather_report", + }, + } + + tool = Tool.from_dict(legacy) + assert tool.function is get_weather_report + assert tool.async_function is None + + def test_check_duplicate_tool_names(): tools = [ Tool(name="weather", description="Get weather report", parameters=parameters, function=get_weather_report),