c56bef871b
Sync docs with Docusaurus / sync (push) Waiting to run
Tests / Check if changed (push) Waiting to run
Tests / format (push) Blocked by required conditions
Tests / check-imports (push) Blocked by required conditions
Tests / Unit / macos-latest (push) Blocked by required conditions
Tests / Unit / ubuntu-latest (push) Blocked by required conditions
Tests / Unit / windows-latest (push) Blocked by required conditions
Tests / mypy (push) Blocked by required conditions
Tests / Integration / ubuntu-latest (push) Blocked by required conditions
Tests / Integration / macos-latest (push) Blocked by required conditions
Tests / Integration / windows-latest (push) Blocked by required conditions
Tests / notify-slack-on-failure (push) Blocked by required conditions
Tests / Mark tests as completed (push) Blocked by required conditions
Update Platform Components Table / update (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
Docker image release / Build base image (push) Waiting to run
458 lines
22 KiB
Python
458 lines
22 KiB
Python
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import asyncio
|
|
import inspect
|
|
from collections.abc import Callable
|
|
from dataclasses import asdict, dataclass
|
|
from typing import Any
|
|
|
|
from jsonschema import Draft202012Validator
|
|
from jsonschema.exceptions import SchemaError
|
|
|
|
from haystack.core.serialization import generate_qualified_class_name
|
|
from haystack.tools.errors import ToolInvocationError
|
|
from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
|
|
|
|
|
|
@dataclass
|
|
class Tool:
|
|
"""
|
|
Data class representing a Tool that Language Models can prepare a call for.
|
|
|
|
Accurate definitions of the textual attributes such as `name` and `description`
|
|
are important for the Language Model to correctly prepare the call.
|
|
|
|
For resource-intensive operations like establishing connections to remote services or
|
|
loading models, override the `warm_up()` method. This method is called before the Tool
|
|
is used and should be idempotent, as it may be called multiple times during
|
|
pipeline/agent setup.
|
|
|
|
:param name:
|
|
Name of the Tool.
|
|
:param description:
|
|
Description of the Tool.
|
|
:param parameters:
|
|
A JSON schema defining the parameters expected by the Tool.
|
|
:param function:
|
|
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.
|
|
|
|
`outputs_to_string` supports two formats:
|
|
|
|
1. Single output format - use "source", "handler", and/or "raw_result" at the root level:
|
|
```python
|
|
{
|
|
"source": "docs", "handler": format_documents, "raw_result": False
|
|
}
|
|
```
|
|
- `source`: If provided, only the specified output key is sent to the handler. If not provided, the whole
|
|
tool result is sent to the handler.
|
|
- `handler`: A function that takes the tool output (or the extracted source value) and returns the
|
|
final result.
|
|
- `raw_result`: If `True`, the result is returned raw without string conversion, but applying the `handler`
|
|
if provided. This is intended for tools that return images. In this mode, the Tool function or the
|
|
`handler` must return a list of `TextContent`/`ImageContent` objects to ensure compatibility with Chat
|
|
Generators.
|
|
|
|
2. Multiple output format - map keys to individual configurations:
|
|
```python
|
|
{
|
|
"formatted_docs": {"source": "docs", "handler": format_documents},
|
|
"summary": {"source": "summary_text", "handler": str.upper}
|
|
}
|
|
```
|
|
Each key maps to a dictionary that can contain "source" and/or "handler".
|
|
Note that `raw_result` is not supported in the multiple output format.
|
|
:param inputs_from_state:
|
|
Optional dictionary mapping state keys to tool parameter names.
|
|
Example: `{"repository": "repo"}` maps state's "repository" to tool's "repo" parameter.
|
|
:param outputs_to_state:
|
|
Optional dictionary defining how tool outputs map to keys within state as well as optional handlers.
|
|
If the source is provided only the specified output key is sent to the handler.
|
|
Example:
|
|
```python
|
|
{
|
|
"documents": {"source": "docs", "handler": custom_handler}
|
|
}
|
|
```
|
|
If the source is omitted the whole tool result is sent to the handler.
|
|
Example:
|
|
```python
|
|
{
|
|
"documents": {"handler": custom_handler}
|
|
}
|
|
```
|
|
: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.
|
|
"""
|
|
|
|
name: str
|
|
description: str
|
|
parameters: dict[str, Any]
|
|
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
|
|
# 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"`function` must be a synchronous function. "
|
|
f"The function '{self.function.__name__}' is a coroutine function. "
|
|
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
|
|
try:
|
|
Draft202012Validator.check_schema(self.parameters)
|
|
except SchemaError as e:
|
|
raise ValueError("The provided parameters do not define a valid JSON schema") from e
|
|
|
|
# Validate outputs structure if provided
|
|
if self.outputs_to_state is not None:
|
|
for key, config in self.outputs_to_state.items():
|
|
if not isinstance(config, dict):
|
|
raise TypeError(f"outputs_to_state configuration for key '{key}' must be a dictionary")
|
|
if "source" in config and not isinstance(config["source"], str):
|
|
raise ValueError(f"outputs_to_state source for key '{key}' must be a string.")
|
|
if "handler" in config and not callable(config["handler"]):
|
|
raise ValueError(f"outputs_to_state handler for key '{key}' must be callable")
|
|
|
|
# Validate that outputs_to_state source keys exist as valid tool outputs
|
|
valid_outputs: set[str] | None = self._get_valid_outputs()
|
|
if valid_outputs is not None:
|
|
for state_key, config in self.outputs_to_state.items():
|
|
source = config.get("source")
|
|
if source is not None and source not in valid_outputs:
|
|
raise ValueError(
|
|
f"outputs_to_state: '{self.name}' maps state key '{state_key}' to unknown output '{source}'"
|
|
f"Valid outputs are: {valid_outputs}."
|
|
)
|
|
|
|
if self.outputs_to_string is not None:
|
|
if "source" in self.outputs_to_string and not isinstance(self.outputs_to_string["source"], str):
|
|
raise ValueError("outputs_to_string source must be a string.")
|
|
if "handler" in self.outputs_to_string and not callable(self.outputs_to_string["handler"]):
|
|
raise ValueError("outputs_to_string handler must be callable")
|
|
if "raw_result" in self.outputs_to_string and not isinstance(self.outputs_to_string["raw_result"], bool):
|
|
raise ValueError("outputs_to_string raw_result must be a boolean.")
|
|
|
|
if (
|
|
"source" in self.outputs_to_string
|
|
or "handler" in self.outputs_to_string
|
|
or "raw_result" in self.outputs_to_string
|
|
):
|
|
# Single output configuration
|
|
for key in self.outputs_to_string:
|
|
if key not in {"source", "handler", "raw_result"}:
|
|
raise ValueError(
|
|
"Invalid outputs_to_string config. "
|
|
"When using 'source', 'handler' or 'raw_result' at the root level, no other keys are "
|
|
" allowed. Use individual output configs instead."
|
|
)
|
|
else:
|
|
# Multiple outputs configuration
|
|
for key, config in self.outputs_to_string.items():
|
|
if not isinstance(config, dict):
|
|
raise TypeError(f"outputs_to_string configuration for key '{key}' must be a dictionary")
|
|
if "raw_result" in config:
|
|
raise ValueError(
|
|
f"Invalid outputs_to_string configuration for key '{key}': "
|
|
f"'raw_result' is not supported in the multiple output format."
|
|
)
|
|
if "source" not in config:
|
|
raise ValueError(
|
|
f"Invalid outputs_to_string configuration for key '{key}': "
|
|
f"each output must have a 'source' defined."
|
|
)
|
|
if "source" in config and not isinstance(config["source"], str):
|
|
raise ValueError(f"outputs_to_string source for key '{key}' must be a string.")
|
|
if "handler" in config and not callable(config["handler"]):
|
|
raise ValueError(f"outputs_to_string handler for key '{key}' must be callable")
|
|
|
|
# Validate that inputs_from_state parameter names exist as valid tool parameters
|
|
if self.inputs_from_state is not None:
|
|
valid_inputs = self._get_valid_inputs()
|
|
for state_key, param_name in self.inputs_from_state.items():
|
|
if not isinstance(param_name, str):
|
|
raise TypeError(
|
|
f"inputs_from_state values must be str, not {type(param_name).__name__}. "
|
|
f"Got {param_name!r} for key '{state_key}'."
|
|
)
|
|
if valid_inputs and param_name not in valid_inputs:
|
|
raise ValueError(
|
|
f"inputs_from_state maps '{state_key}' to unknown parameter '{param_name}'. "
|
|
f"Valid parameters are: {valid_inputs}."
|
|
)
|
|
|
|
def _get_valid_inputs(self) -> set[str]:
|
|
"""
|
|
Return the set of valid input parameter names that this tool accepts.
|
|
|
|
Used to validate that `inputs_from_state` only references parameters that actually exist.
|
|
This prevents typos and catches configuration errors at tool construction time.
|
|
|
|
By default, introspects the function signature to get ALL parameters, including those
|
|
that may be excluded from the JSON schema (e.g., parameters mapped from state).
|
|
Falls back to schema properties if introspection fails.
|
|
|
|
Subclasses like ComponentTool override this to return component input socket names.
|
|
|
|
:returns: Set of valid input parameter names for validation.
|
|
"""
|
|
# Combine parameters from both function signature and schema for robustness
|
|
# Function signature includes all parameters (even those excluded from schema)
|
|
# Schema properties provide the validated parameter set
|
|
valid_params: set[str] = set()
|
|
|
|
# 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())
|
|
|
|
return valid_params
|
|
|
|
def _get_valid_outputs(self) -> set[str] | None:
|
|
"""
|
|
Return the set of valid output names that this tool produces.
|
|
|
|
Used to validate that `outputs_to_state` only references outputs that actually exist.
|
|
This prevents typos and catches configuration errors at tool construction time.
|
|
|
|
By default, returns None because regular function-based tools don't have a formal
|
|
output schema. When None is returned, output validation is skipped.
|
|
|
|
Subclasses like ComponentTool override this to return component output socket names,
|
|
enabling validation for tools where outputs are known.
|
|
|
|
:returns: Set of valid output names for validation, or None to skip validation.
|
|
"""
|
|
return None
|
|
|
|
@property
|
|
def tool_spec(self) -> dict[str, Any]:
|
|
"""
|
|
Return the Tool specification to be used by the Language Model.
|
|
"""
|
|
return {"name": self.name, "description": self.description, "parameters": self.parameters}
|
|
|
|
def warm_up(self) -> None:
|
|
"""
|
|
Prepare the Tool for use.
|
|
|
|
Override this method to establish connections to remote services, load models,
|
|
or perform other resource-intensive initialization. This method should be idempotent,
|
|
as it may be called multiple times.
|
|
"""
|
|
pass
|
|
|
|
def invoke(self, **kwargs: Any) -> Any:
|
|
"""
|
|
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:
|
|
raise ToolInvocationError(
|
|
f"Failed to invoke Tool `{self.name}` with parameters {kwargs}. Error: {e}", tool_name=self.name
|
|
) 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.
|
|
|
|
:returns:
|
|
Dictionary with serialized data.
|
|
"""
|
|
data = asdict(self)
|
|
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)
|
|
|
|
if self.outputs_to_string is not None:
|
|
data["outputs_to_string"] = _serialize_outputs_to_string(self.outputs_to_string)
|
|
|
|
return {"type": generate_qualified_class_name(type(self)), "data": data}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict[str, Any]) -> "Tool":
|
|
"""
|
|
Deserializes the Tool from a dictionary.
|
|
|
|
:param data:
|
|
Dictionary to deserialize from.
|
|
:returns:
|
|
Deserialized Tool.
|
|
"""
|
|
init_parameters = data["data"]
|
|
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"])
|
|
|
|
if init_parameters.get("outputs_to_string") is not None:
|
|
init_parameters["outputs_to_string"] = _deserialize_outputs_to_string(init_parameters["outputs_to_string"])
|
|
|
|
return cls(**init_parameters)
|
|
|
|
|
|
def _check_duplicate_tool_names(tools: list[Tool] | None) -> None:
|
|
"""
|
|
Checks for duplicate tool names and raises a ValueError if they are found.
|
|
|
|
:param tools: The list of tools to check.
|
|
:raises ValueError: If duplicate tool names are found.
|
|
"""
|
|
if tools is None:
|
|
return
|
|
tool_names = [tool.name for tool in tools]
|
|
duplicate_tool_names = {name for name in tool_names if tool_names.count(name) > 1}
|
|
if duplicate_tool_names:
|
|
raise ValueError(f"Duplicate tool names found: {duplicate_tool_names}")
|
|
|
|
|
|
def _convert_handler(config: dict[str, Any], converter: Callable[[Any], Any]) -> dict[str, Any]:
|
|
"""
|
|
Copies a single output config, converting its "handler" entry (if present) via `converter`.
|
|
|
|
:param config: A single output configuration dictionary that may contain a "handler" key.
|
|
:param converter: `serialize_callable` or `deserialize_callable`, applied to the "handler" value.
|
|
:returns: A copy of `config` with the "handler" value converted, if present.
|
|
"""
|
|
new_config = config.copy()
|
|
if "handler" in config:
|
|
new_config["handler"] = converter(config["handler"])
|
|
return new_config
|
|
|
|
|
|
def _convert_handler_in_configs(
|
|
configs: dict[str, dict[str, Any]], converter: Callable[[Any], Any]
|
|
) -> dict[str, dict[str, Any]]:
|
|
"""
|
|
Applies `_convert_handler` to every config in a dictionary of named output configs.
|
|
|
|
:param configs: A mapping of keys to output configuration dictionaries.
|
|
:param converter: `serialize_callable` or `deserialize_callable`, applied to each "handler" value.
|
|
:returns: A new mapping with the same keys, each config converted via `_convert_handler`.
|
|
"""
|
|
return {key: _convert_handler(config, converter) for key, config in configs.items()}
|
|
|
|
|
|
def _serialize_outputs_to_state(outputs_to_state: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
|
|
"""
|
|
Serializes the outputs_to_state dictionary, converting any callable handlers to their string representation.
|
|
|
|
:param outputs_to_state: The outputs_to_state dictionary to serialize.
|
|
:returns: The serialized outputs_to_state dictionary.
|
|
"""
|
|
return _convert_handler_in_configs(outputs_to_state, serialize_callable)
|
|
|
|
|
|
def _deserialize_outputs_to_state(outputs_to_state: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
|
|
"""
|
|
Deserializes the outputs_to_state dictionary, converting any string handlers back to callables.
|
|
|
|
:param outputs_to_state: The outputs_to_state dictionary to deserialize.
|
|
:returns: The deserialized outputs_to_state dictionary.
|
|
"""
|
|
return _convert_handler_in_configs(outputs_to_state, deserialize_callable)
|
|
|
|
|
|
def _serialize_outputs_to_string(outputs_to_string: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Serializes the outputs_to_string dictionary, converting any callable handlers to their string representation.
|
|
|
|
:param outputs_to_string: The outputs_to_string dictionary to serialize.
|
|
:returns: The serialized outputs_to_string dictionary.
|
|
"""
|
|
if "source" in outputs_to_string or "handler" in outputs_to_string or "raw_result" in outputs_to_string:
|
|
# Single output configuration
|
|
return _convert_handler(outputs_to_string, serialize_callable)
|
|
|
|
# Multiple outputs configuration
|
|
return _convert_handler_in_configs(outputs_to_string, serialize_callable)
|
|
|
|
|
|
def _deserialize_outputs_to_string(outputs_to_string: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Deserializes the outputs_to_string dictionary, converting any string handlers back to callables.
|
|
|
|
:param outputs_to_string: The outputs_to_string dictionary to deserialize.
|
|
:returns: The deserialized outputs_to_string dictionary.
|
|
"""
|
|
if "source" in outputs_to_string or "handler" in outputs_to_string or "raw_result" in outputs_to_string:
|
|
# Single output configuration
|
|
return _convert_handler(outputs_to_string, deserialize_callable)
|
|
|
|
# Multiple outputs configuration
|
|
return _convert_handler_in_configs(outputs_to_string, deserialize_callable)
|