feat: Add support for async functions inside of Tool (#11426)

This commit is contained in:
Sebastian Husch Lee
2026-06-01 11:42:37 +00:00
committed by GitHub
parent 7dd1eedde3
commit 1cab85a5d5
20 changed files with 482 additions and 70 deletions
+48 -26
View File
@@ -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:
+26 -1
View File
@@ -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,
+7 -2
View File
@@ -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,
+5
View File
@@ -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
+74 -18
View File
@@ -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"])
@@ -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``.
+8 -4
View File
@@ -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
@@ -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,
@@ -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
@@ -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,
@@ -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,
@@ -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,
@@ -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,
@@ -341,6 +341,7 @@ class TestOpenAIChatGenerator:
{
"type": "haystack.tools.tool.Tool",
"data": {
"async_function": None,
"description": "description",
"function": "builtins.print",
"inputs_from_state": None,
@@ -255,6 +255,7 @@ class TestSerDe:
{
"type": "haystack.tools.tool.Tool",
"data": {
"async_function": None,
"description": "description",
"function": "builtins.print",
"inputs_from_state": None,
+49
View File
@@ -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"}
+27 -13
View File
@@ -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"
+19
View File
@@ -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
+4 -4
View File
@@ -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)
+113 -2
View File
@@ -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),