Coverage for haystack/tools/tool.py: 97%
154 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import asyncio
6import inspect
7from collections.abc import Callable
8from dataclasses import asdict, dataclass
9from typing import Any
11from jsonschema import Draft202012Validator
12from jsonschema.exceptions import SchemaError
14from haystack.core.serialization import generate_qualified_class_name
15from haystack.tools.errors import ToolInvocationError
16from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
19@dataclass
20class Tool:
21 """
22 Data class representing a Tool that Language Models can prepare a call for.
24 Accurate definitions of the textual attributes such as `name` and `description`
25 are important for the Language Model to correctly prepare the call.
27 For resource-intensive operations like establishing connections to remote services or
28 loading models, implement an optional `warm_up()` method. It is called before invocation
29 and during pipeline/agent setup, and must be idempotent. Native async tools can implement
30 `warm_up_async()` instead. Implement `close()` and/or `close_async()` to release owned resources;
31 these methods must be safe to call before warm-up and multiple times.
33 :param name:
34 Name of the Tool.
35 :param description:
36 Description of the Tool.
37 :param parameters:
38 A JSON schema defining the parameters expected by the Tool.
39 :param function:
40 The synchronous function invoked by `Tool.invoke`. Must be a regular function — coroutine functions should
41 be passed to `async_function` instead. Either `function` or `async_function` (or both) must be set.
42 :param async_function:
43 Optional coroutine function awaited by `Tool.invoke_async`. When only `async_function` is set, `invoke` raises
44 a `ToolInvocationError`. When only `function` is set, `invoke_async` falls back to running `function` in a
45 worker thread via `asyncio.to_thread`.
46 :param outputs_to_string:
47 Optional dictionary defining how tool outputs should be converted into string(s) or results.
48 If not provided, the tool result is converted to a string using a default handler.
50 `outputs_to_string` supports two formats:
52 1. Single output format - use "source", "handler", and/or "raw_result" at the root level:
53 ```python
54 {
55 "source": "docs", "handler": format_documents, "raw_result": False
56 }
57 ```
58 - `source`: If provided, only the specified output key is sent to the handler. If not provided, the whole
59 tool result is sent to the handler.
60 - `handler`: A function that takes the tool output (or the extracted source value) and returns the
61 final result.
62 - `raw_result`: If `True`, the result is returned raw without string conversion, but applying the `handler`
63 if provided. This is intended for tools that return images. In this mode, the Tool function or the
64 `handler` must return a list of `TextContent`/`ImageContent` objects to ensure compatibility with Chat
65 Generators.
67 2. Multiple output format - map keys to individual configurations:
68 ```python
69 {
70 "formatted_docs": {"source": "docs", "handler": format_documents},
71 "summary": {"source": "summary_text", "handler": str.upper}
72 }
73 ```
74 Each key maps to a dictionary that can contain "source" and/or "handler".
75 Note that `raw_result` is not supported in the multiple output format.
76 :param inputs_from_state:
77 Optional dictionary mapping state keys to tool parameter names.
78 Example: `{"repository": "repo"}` maps state's "repository" to tool's "repo" parameter.
79 :param outputs_to_state:
80 Optional dictionary defining how tool outputs map to keys within state as well as optional handlers.
81 If the source is provided only the specified output key is sent to the handler.
82 Example:
83 ```python
84 {
85 "documents": {"source": "docs", "handler": custom_handler}
86 }
87 ```
88 If the source is omitted the whole tool result is sent to the handler.
89 Example:
90 ```python
91 {
92 "documents": {"handler": custom_handler}
93 }
94 ```
95 :raises ValueError: If neither `function` nor `async_function` is provided, if `function` is a
96 coroutine function, if `async_function` is not a coroutine function, if `parameters` is not a
97 valid JSON schema, or if the `outputs_to_state`, `outputs_to_string`, or `inputs_from_state`
98 configurations are invalid.
99 :raises TypeError: If any configuration value in `outputs_to_state`, `outputs_to_string`, or
100 `inputs_from_state` has the wrong type.
101 """
103 name: str
104 description: str
105 parameters: dict[str, Any]
106 function: Callable | None = None
107 outputs_to_string: dict[str, Any] | None = None
108 inputs_from_state: dict[str, str] | None = None
109 outputs_to_state: dict[str, dict[str, Any]] | None = None
110 async_function: Callable | None = None
112 def __post_init__(self) -> None: # noqa: C901, PLR0912
113 # At least one of function / async_function must be set.
114 if self.function is None and self.async_function is None:
115 raise ValueError(f"Tool '{self.name}' requires at least one of `function` or `async_function` to be set.")
117 # `function` must be a regular (sync) function. Coroutine functions belong on `async_function`.
118 if self.function is not None and inspect.iscoroutinefunction(self.function):
119 raise ValueError(
120 f"`function` must be a synchronous function. "
121 f"The function '{self.function.__name__}' is a coroutine function. "
122 f"Pass it as `async_function` instead."
123 )
125 # `async_function` must be a coroutine function defined with `async def`.
126 if self.async_function is not None and not inspect.iscoroutinefunction(self.async_function):
127 raise ValueError(
128 f"`async_function` must be a coroutine function defined with `async def`. "
129 f"Got '{getattr(self.async_function, '__name__', repr(self.async_function))}'."
130 )
132 # Check that the parameters define a valid JSON schema
133 try:
134 Draft202012Validator.check_schema(self.parameters)
135 except SchemaError as e:
136 raise ValueError("The provided parameters do not define a valid JSON schema") from e
138 # Validate outputs structure if provided
139 if self.outputs_to_state is not None:
140 for key, config in self.outputs_to_state.items():
141 if not isinstance(config, dict):
142 raise TypeError(f"outputs_to_state configuration for key '{key}' must be a dictionary")
143 if "source" in config and not isinstance(config["source"], str):
144 raise ValueError(f"outputs_to_state source for key '{key}' must be a string.")
145 if "handler" in config and not callable(config["handler"]):
146 raise ValueError(f"outputs_to_state handler for key '{key}' must be callable")
148 # Validate that outputs_to_state source keys exist as valid tool outputs
149 valid_outputs: set[str] | None = self._get_valid_outputs()
150 if valid_outputs is not None:
151 for state_key, config in self.outputs_to_state.items():
152 source = config.get("source")
153 if source is not None and source not in valid_outputs:
154 raise ValueError(
155 f"outputs_to_state: '{self.name}' maps state key '{state_key}' to unknown output '{source}'"
156 f"Valid outputs are: {valid_outputs}."
157 )
159 if self.outputs_to_string is not None:
160 if "source" in self.outputs_to_string and not isinstance(self.outputs_to_string["source"], str):
161 raise ValueError("outputs_to_string source must be a string.")
162 if "handler" in self.outputs_to_string and not callable(self.outputs_to_string["handler"]):
163 raise ValueError("outputs_to_string handler must be callable")
164 if "raw_result" in self.outputs_to_string and not isinstance(self.outputs_to_string["raw_result"], bool):
165 raise ValueError("outputs_to_string raw_result must be a boolean.")
167 if (
168 "source" in self.outputs_to_string
169 or "handler" in self.outputs_to_string
170 or "raw_result" in self.outputs_to_string
171 ):
172 # Single output configuration
173 for key in self.outputs_to_string:
174 if key not in {"source", "handler", "raw_result"}:
175 raise ValueError(
176 "Invalid outputs_to_string config. "
177 "When using 'source', 'handler' or 'raw_result' at the root level, no other keys are "
178 " allowed. Use individual output configs instead."
179 )
180 else:
181 # Multiple outputs configuration
182 for key, config in self.outputs_to_string.items():
183 if not isinstance(config, dict):
184 raise TypeError(f"outputs_to_string configuration for key '{key}' must be a dictionary")
185 if "raw_result" in config:
186 raise ValueError(
187 f"Invalid outputs_to_string configuration for key '{key}': "
188 f"'raw_result' is not supported in the multiple output format."
189 )
190 if "source" not in config:
191 raise ValueError(
192 f"Invalid outputs_to_string configuration for key '{key}': "
193 f"each output must have a 'source' defined."
194 )
195 if "source" in config and not isinstance(config["source"], str):
196 raise ValueError(f"outputs_to_string source for key '{key}' must be a string.")
197 if "handler" in config and not callable(config["handler"]):
198 raise ValueError(f"outputs_to_string handler for key '{key}' must be callable")
200 # Validate that inputs_from_state parameter names exist as valid tool parameters
201 if self.inputs_from_state is not None:
202 valid_inputs = self._get_valid_inputs()
203 for state_key, param_name in self.inputs_from_state.items():
204 if not isinstance(param_name, str):
205 raise TypeError(
206 f"inputs_from_state values must be str, not {type(param_name).__name__}. "
207 f"Got {param_name!r} for key '{state_key}'."
208 )
209 if valid_inputs and param_name not in valid_inputs:
210 raise ValueError(
211 f"inputs_from_state maps '{state_key}' to unknown parameter '{param_name}'. "
212 f"Valid parameters are: {valid_inputs}."
213 )
215 def _get_valid_inputs(self) -> set[str]:
216 """
217 Return the set of valid input parameter names that this tool accepts.
219 Used to validate that `inputs_from_state` only references parameters that actually exist.
220 This prevents typos and catches configuration errors at tool construction time.
222 By default, introspects the function signature to get ALL parameters, including those
223 that may be excluded from the JSON schema (e.g., parameters mapped from state).
224 Falls back to schema properties if introspection fails.
226 Subclasses like ComponentTool override this to return component input socket names.
228 :returns: Set of valid input parameter names for validation.
229 """
230 # Combine parameters from both function signature and schema for robustness
231 # Function signature includes all parameters (even those excluded from schema)
232 # Schema properties provide the validated parameter set
233 valid_params: set[str] = set()
235 # Try to get parameters from function introspection.
236 # Prefer `function`; fall back to `async_function` for async-only tools.
237 introspection_target = self.function if self.function is not None else self.async_function
238 if introspection_target is not None:
239 try:
240 sig = inspect.signature(introspection_target)
241 valid_params.update(sig.parameters.keys())
242 except (ValueError, TypeError):
243 pass # Introspection failed, will rely on schema
245 # Add parameters from schema (union with function params)
246 valid_params.update(self.parameters.get("properties", {}).keys())
248 return valid_params
250 def _get_valid_outputs(self) -> set[str] | None:
251 """
252 Return the set of valid output names that this tool produces.
254 Used to validate that `outputs_to_state` only references outputs that actually exist.
255 This prevents typos and catches configuration errors at tool construction time.
257 By default, returns None because regular function-based tools don't have a formal
258 output schema. When None is returned, output validation is skipped.
260 Subclasses like ComponentTool override this to return component output socket names,
261 enabling validation for tools where outputs are known.
263 :returns: Set of valid output names for validation, or None to skip validation.
264 """
265 return None
267 @property
268 def tool_spec(self) -> dict[str, Any]:
269 """
270 Return the Tool specification to be used by the Language Model.
271 """
272 return {"name": self.name, "description": self.description, "parameters": self.parameters}
274 def invoke(self, **kwargs: Any) -> Any:
275 """
276 Warm up the Tool if supported and invoke it synchronously with the provided keyword arguments.
278 :raises ToolInvocationError: If the Tool has no sync `function`, or if the underlying call
279 raises an exception.
280 """
281 if self.function is None:
282 raise ToolInvocationError(
283 f"Tool `{self.name}` has no sync `function` and can only be invoked via `invoke_async` "
284 f"(use `Agent.run_async`).",
285 tool_name=self.name,
286 )
288 try:
289 if hasattr(self, "warm_up"):
290 self.warm_up()
291 result = self.function(**kwargs)
292 except Exception as e:
293 raise ToolInvocationError(
294 f"Failed to invoke Tool `{self.name}` with parameters {kwargs}. Error: {e}", tool_name=self.name
295 ) from e
296 return result
298 async def invoke_async(self, **kwargs: Any) -> Any:
299 """
300 Invoke the Tool asynchronously with the provided keyword arguments.
302 If `async_function` is set, prepare with `warm_up_async` (falling back to `warm_up`) before awaiting it.
303 Otherwise, synchronous preparation and invocation run in a worker thread via `asyncio.to_thread`,
304 which propagates the current context to the worker.
306 :raises ToolInvocationError: If the underlying call raises an exception.
307 """
308 if self.async_function is None:
309 return await asyncio.to_thread(self.invoke, **kwargs)
311 try:
312 if hasattr(self, "warm_up_async"):
313 await self.warm_up_async()
314 elif hasattr(self, "warm_up"):
315 self.warm_up()
316 return await self.async_function(**kwargs)
317 except Exception as e:
318 raise ToolInvocationError(
319 f"Failed to invoke Tool `{self.name}` with parameters {kwargs}. Error: {e}", tool_name=self.name
320 ) from e
322 def to_dict(self) -> dict[str, Any]:
323 """
324 Serializes the Tool to a dictionary.
326 :returns:
327 Dictionary with serialized data.
328 """
329 data = asdict(self)
330 data["function"] = serialize_callable(self.function) if self.function is not None else None
331 data["async_function"] = serialize_callable(self.async_function) if self.async_function is not None else None
333 if self.outputs_to_state is not None:
334 data["outputs_to_state"] = _serialize_outputs_to_state(self.outputs_to_state)
336 if self.outputs_to_string is not None:
337 data["outputs_to_string"] = _serialize_outputs_to_string(self.outputs_to_string)
339 return {"type": generate_qualified_class_name(type(self)), "data": data}
341 @classmethod
342 def from_dict(cls, data: dict[str, Any]) -> "Tool":
343 """
344 Deserializes the Tool from a dictionary.
346 :param data:
347 Dictionary to deserialize from.
348 :returns:
349 Deserialized Tool.
350 """
351 init_parameters = data["data"]
352 init_parameters["function"] = (
353 deserialize_callable(init_parameters["function"]) if init_parameters.get("function") is not None else None
354 )
355 if init_parameters.get("async_function") is not None:
356 init_parameters["async_function"] = deserialize_callable(init_parameters["async_function"])
357 if "outputs_to_state" in init_parameters and init_parameters["outputs_to_state"]:
358 init_parameters["outputs_to_state"] = _deserialize_outputs_to_state(init_parameters["outputs_to_state"])
360 if init_parameters.get("outputs_to_string") is not None:
361 init_parameters["outputs_to_string"] = _deserialize_outputs_to_string(init_parameters["outputs_to_string"])
363 return cls(**init_parameters)
366def _check_duplicate_tool_names(tools: list[Tool] | None) -> None:
367 """
368 Checks for duplicate tool names and raises a ValueError if they are found.
370 :param tools: The list of tools to check.
371 :raises ValueError: If duplicate tool names are found.
372 """
373 if tools is None:
374 return
375 tool_names = [tool.name for tool in tools]
376 duplicate_tool_names = {name for name in tool_names if tool_names.count(name) > 1}
377 if duplicate_tool_names:
378 raise ValueError(f"Duplicate tool names found: {duplicate_tool_names}")
381def _convert_handler(config: dict[str, Any], converter: Callable[[Any], Any]) -> dict[str, Any]:
382 """
383 Copies a single output config, converting its "handler" entry (if present) via `converter`.
385 :param config: A single output configuration dictionary that may contain a "handler" key.
386 :param converter: `serialize_callable` or `deserialize_callable`, applied to the "handler" value.
387 :returns: A copy of `config` with the "handler" value converted, if present.
388 """
389 new_config = config.copy()
390 if "handler" in config:
391 new_config["handler"] = converter(config["handler"])
392 return new_config
395def _convert_handler_in_configs(
396 configs: dict[str, dict[str, Any]], converter: Callable[[Any], Any]
397) -> dict[str, dict[str, Any]]:
398 """
399 Applies `_convert_handler` to every config in a dictionary of named output configs.
401 :param configs: A mapping of keys to output configuration dictionaries.
402 :param converter: `serialize_callable` or `deserialize_callable`, applied to each "handler" value.
403 :returns: A new mapping with the same keys, each config converted via `_convert_handler`.
404 """
405 return {key: _convert_handler(config, converter) for key, config in configs.items()}
408def _serialize_outputs_to_state(outputs_to_state: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
409 """
410 Serializes the outputs_to_state dictionary, converting any callable handlers to their string representation.
412 :param outputs_to_state: The outputs_to_state dictionary to serialize.
413 :returns: The serialized outputs_to_state dictionary.
414 """
415 return _convert_handler_in_configs(outputs_to_state, serialize_callable)
418def _deserialize_outputs_to_state(outputs_to_state: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
419 """
420 Deserializes the outputs_to_state dictionary, converting any string handlers back to callables.
422 :param outputs_to_state: The outputs_to_state dictionary to deserialize.
423 :returns: The deserialized outputs_to_state dictionary.
424 """
425 return _convert_handler_in_configs(outputs_to_state, deserialize_callable)
428def _serialize_outputs_to_string(outputs_to_string: dict[str, Any]) -> dict[str, Any]:
429 """
430 Serializes the outputs_to_string dictionary, converting any callable handlers to their string representation.
432 :param outputs_to_string: The outputs_to_string dictionary to serialize.
433 :returns: The serialized outputs_to_string dictionary.
434 """
435 if "source" in outputs_to_string or "handler" in outputs_to_string or "raw_result" in outputs_to_string:
436 # Single output configuration
437 return _convert_handler(outputs_to_string, serialize_callable)
439 # Multiple outputs configuration
440 return _convert_handler_in_configs(outputs_to_string, serialize_callable)
443def _deserialize_outputs_to_string(outputs_to_string: dict[str, Any]) -> dict[str, Any]:
444 """
445 Deserializes the outputs_to_string dictionary, converting any string handlers back to callables.
447 :param outputs_to_string: The outputs_to_string dictionary to deserialize.
448 :returns: The deserialized outputs_to_string dictionary.
449 """
450 if "source" in outputs_to_string or "handler" in outputs_to_string or "raw_result" in outputs_to_string:
451 # Single output configuration
452 return _convert_handler(outputs_to_string, deserialize_callable)
454 # Multiple outputs configuration
455 return _convert_handler_in_configs(outputs_to_string, deserialize_callable)