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 

4 

5import asyncio 

6import inspect 

7from collections.abc import Callable 

8from dataclasses import asdict, dataclass 

9from typing import Any 

10 

11from jsonschema import Draft202012Validator 

12from jsonschema.exceptions import SchemaError 

13 

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 

17 

18 

19@dataclass 

20class Tool: 

21 """ 

22 Data class representing a Tool that Language Models can prepare a call for. 

23 

24 Accurate definitions of the textual attributes such as `name` and `description` 

25 are important for the Language Model to correctly prepare the call. 

26 

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. 

32 

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. 

49 

50 `outputs_to_string` supports two formats: 

51 

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. 

66 

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 """ 

102 

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 

111 

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.") 

116 

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 ) 

124 

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 ) 

131 

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 

137 

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

147 

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 ) 

158 

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.") 

166 

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

199 

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 ) 

214 

215 def _get_valid_inputs(self) -> set[str]: 

216 """ 

217 Return the set of valid input parameter names that this tool accepts. 

218 

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. 

221 

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. 

225 

226 Subclasses like ComponentTool override this to return component input socket names. 

227 

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

234 

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 

244 

245 # Add parameters from schema (union with function params) 

246 valid_params.update(self.parameters.get("properties", {}).keys()) 

247 

248 return valid_params 

249 

250 def _get_valid_outputs(self) -> set[str] | None: 

251 """ 

252 Return the set of valid output names that this tool produces. 

253 

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. 

256 

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. 

259 

260 Subclasses like ComponentTool override this to return component output socket names, 

261 enabling validation for tools where outputs are known. 

262 

263 :returns: Set of valid output names for validation, or None to skip validation. 

264 """ 

265 return None 

266 

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} 

273 

274 def invoke(self, **kwargs: Any) -> Any: 

275 """ 

276 Warm up the Tool if supported and invoke it synchronously with the provided keyword arguments. 

277 

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 ) 

287 

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 

297 

298 async def invoke_async(self, **kwargs: Any) -> Any: 

299 """ 

300 Invoke the Tool asynchronously with the provided keyword arguments. 

301 

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. 

305 

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) 

310 

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 

321 

322 def to_dict(self) -> dict[str, Any]: 

323 """ 

324 Serializes the Tool to a dictionary. 

325 

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 

332 

333 if self.outputs_to_state is not None: 

334 data["outputs_to_state"] = _serialize_outputs_to_state(self.outputs_to_state) 

335 

336 if self.outputs_to_string is not None: 

337 data["outputs_to_string"] = _serialize_outputs_to_string(self.outputs_to_string) 

338 

339 return {"type": generate_qualified_class_name(type(self)), "data": data} 

340 

341 @classmethod 

342 def from_dict(cls, data: dict[str, Any]) -> "Tool": 

343 """ 

344 Deserializes the Tool from a dictionary. 

345 

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"]) 

359 

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"]) 

362 

363 return cls(**init_parameters) 

364 

365 

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. 

369 

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}") 

379 

380 

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`. 

384 

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 

393 

394 

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. 

400 

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()} 

406 

407 

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. 

411 

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) 

416 

417 

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. 

421 

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) 

426 

427 

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. 

431 

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) 

438 

439 # Multiple outputs configuration 

440 return _convert_handler_in_configs(outputs_to_string, serialize_callable) 

441 

442 

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. 

446 

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) 

453 

454 # Multiple outputs configuration 

455 return _convert_handler_in_configs(outputs_to_string, deserialize_callable)