Coverage for haystack/tools/tool.py: 97%

150 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-21 13:53 +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, override the `warm_up()` method. This method is called before the Tool 

29 is used and should be idempotent, as it may be called multiple times during 

30 pipeline/agent setup. 

31 

32 :param name: 

33 Name of the Tool. 

34 :param description: 

35 Description of the Tool. 

36 :param parameters: 

37 A JSON schema defining the parameters expected by the Tool. 

38 :param function: 

39 The synchronous function invoked by `Tool.invoke`. Must be a regular function — coroutine functions should 

40 be passed to `async_function` instead. Either `function` or `async_function` (or both) must be set. 

41 :param async_function: 

42 Optional coroutine function awaited by `Tool.invoke_async`. When only `async_function` is set, `invoke` raises 

43 a `ToolInvocationError`. When only `function` is set, `invoke_async` falls back to running `function` in a 

44 worker thread via `asyncio.to_thread`. 

45 :param outputs_to_string: 

46 Optional dictionary defining how tool outputs should be converted into string(s) or results. 

47 If not provided, the tool result is converted to a string using a default handler. 

48 

49 `outputs_to_string` supports two formats: 

50 

51 1. Single output format - use "source", "handler", and/or "raw_result" at the root level: 

52 ```python 

53 { 

54 "source": "docs", "handler": format_documents, "raw_result": False 

55 } 

56 ``` 

57 - `source`: If provided, only the specified output key is sent to the handler. If not provided, the whole 

58 tool result is sent to the handler. 

59 - `handler`: A function that takes the tool output (or the extracted source value) and returns the 

60 final result. 

61 - `raw_result`: If `True`, the result is returned raw without string conversion, but applying the `handler` 

62 if provided. This is intended for tools that return images. In this mode, the Tool function or the 

63 `handler` must return a list of `TextContent`/`ImageContent` objects to ensure compatibility with Chat 

64 Generators. 

65 

66 2. Multiple output format - map keys to individual configurations: 

67 ```python 

68 { 

69 "formatted_docs": {"source": "docs", "handler": format_documents}, 

70 "summary": {"source": "summary_text", "handler": str.upper} 

71 } 

72 ``` 

73 Each key maps to a dictionary that can contain "source" and/or "handler". 

74 Note that `raw_result` is not supported in the multiple output format. 

75 :param inputs_from_state: 

76 Optional dictionary mapping state keys to tool parameter names. 

77 Example: `{"repository": "repo"}` maps state's "repository" to tool's "repo" parameter. 

78 :param outputs_to_state: 

79 Optional dictionary defining how tool outputs map to keys within state as well as optional handlers. 

80 If the source is provided only the specified output key is sent to the handler. 

81 Example: 

82 ```python 

83 { 

84 "documents": {"source": "docs", "handler": custom_handler} 

85 } 

86 ``` 

87 If the source is omitted the whole tool result is sent to the handler. 

88 Example: 

89 ```python 

90 { 

91 "documents": {"handler": custom_handler} 

92 } 

93 ``` 

94 :raises ValueError: If neither `function` nor `async_function` is provided, if `function` is a 

95 coroutine function, if `async_function` is not a coroutine function, if `parameters` is not a 

96 valid JSON schema, or if the `outputs_to_state`, `outputs_to_string`, or `inputs_from_state` 

97 configurations are invalid. 

98 :raises TypeError: If any configuration value in `outputs_to_state`, `outputs_to_string`, or 

99 `inputs_from_state` has the wrong type. 

100 """ 

101 

102 name: str 

103 description: str 

104 parameters: dict[str, Any] 

105 function: Callable | None = None 

106 outputs_to_string: dict[str, Any] | None = None 

107 inputs_from_state: dict[str, str] | None = None 

108 outputs_to_state: dict[str, dict[str, Any]] | None = None 

109 async_function: Callable | None = None 

110 

111 def __post_init__(self) -> None: # noqa: C901, PLR0912 

112 # At least one of function / async_function must be set. 

113 if self.function is None and self.async_function is None: 

114 raise ValueError(f"Tool '{self.name}' requires at least one of `function` or `async_function` to be set.") 

115 

116 # `function` must be a regular (sync) function. Coroutine functions belong on `async_function`. 

117 if self.function is not None and inspect.iscoroutinefunction(self.function): 

118 raise ValueError( 

119 f"`function` must be a synchronous function. " 

120 f"The function '{self.function.__name__}' is a coroutine function. " 

121 f"Pass it as `async_function` instead." 

122 ) 

123 

124 # `async_function` must be a coroutine function defined with `async def`. 

125 if self.async_function is not None and not inspect.iscoroutinefunction(self.async_function): 

126 raise ValueError( 

127 f"`async_function` must be a coroutine function defined with `async def`. " 

128 f"Got '{getattr(self.async_function, '__name__', repr(self.async_function))}'." 

129 ) 

130 

131 # Check that the parameters define a valid JSON schema 

132 try: 

133 Draft202012Validator.check_schema(self.parameters) 

134 except SchemaError as e: 

135 raise ValueError("The provided parameters do not define a valid JSON schema") from e 

136 

137 # Validate outputs structure if provided 

138 if self.outputs_to_state is not None: 

139 for key, config in self.outputs_to_state.items(): 

140 if not isinstance(config, dict): 

141 raise TypeError(f"outputs_to_state configuration for key '{key}' must be a dictionary") 

142 if "source" in config and not isinstance(config["source"], str): 

143 raise ValueError(f"outputs_to_state source for key '{key}' must be a string.") 

144 if "handler" in config and not callable(config["handler"]): 

145 raise ValueError(f"outputs_to_state handler for key '{key}' must be callable") 

146 

147 # Validate that outputs_to_state source keys exist as valid tool outputs 

148 valid_outputs: set[str] | None = self._get_valid_outputs() 

149 if valid_outputs is not None: 

150 for state_key, config in self.outputs_to_state.items(): 

151 source = config.get("source") 

152 if source is not None and source not in valid_outputs: 

153 raise ValueError( 

154 f"outputs_to_state: '{self.name}' maps state key '{state_key}' to unknown output '{source}'" 

155 f"Valid outputs are: {valid_outputs}." 

156 ) 

157 

158 if self.outputs_to_string is not None: 

159 if "source" in self.outputs_to_string and not isinstance(self.outputs_to_string["source"], str): 

160 raise ValueError("outputs_to_string source must be a string.") 

161 if "handler" in self.outputs_to_string and not callable(self.outputs_to_string["handler"]): 

162 raise ValueError("outputs_to_string handler must be callable") 

163 if "raw_result" in self.outputs_to_string and not isinstance(self.outputs_to_string["raw_result"], bool): 

164 raise ValueError("outputs_to_string raw_result must be a boolean.") 

165 

166 if ( 

167 "source" in self.outputs_to_string 

168 or "handler" in self.outputs_to_string 

169 or "raw_result" in self.outputs_to_string 

170 ): 

171 # Single output configuration 

172 for key in self.outputs_to_string: 

173 if key not in {"source", "handler", "raw_result"}: 

174 raise ValueError( 

175 "Invalid outputs_to_string config. " 

176 "When using 'source', 'handler' or 'raw_result' at the root level, no other keys are " 

177 " allowed. Use individual output configs instead." 

178 ) 

179 else: 

180 # Multiple outputs configuration 

181 for key, config in self.outputs_to_string.items(): 

182 if not isinstance(config, dict): 

183 raise TypeError(f"outputs_to_string configuration for key '{key}' must be a dictionary") 

184 if "raw_result" in config: 

185 raise ValueError( 

186 f"Invalid outputs_to_string configuration for key '{key}': " 

187 f"'raw_result' is not supported in the multiple output format." 

188 ) 

189 if "source" not in config: 

190 raise ValueError( 

191 f"Invalid outputs_to_string configuration for key '{key}': " 

192 f"each output must have a 'source' defined." 

193 ) 

194 if "source" in config and not isinstance(config["source"], str): 

195 raise ValueError(f"outputs_to_string source for key '{key}' must be a string.") 

196 if "handler" in config and not callable(config["handler"]): 

197 raise ValueError(f"outputs_to_string handler for key '{key}' must be callable") 

198 

199 # Validate that inputs_from_state parameter names exist as valid tool parameters 

200 if self.inputs_from_state is not None: 

201 valid_inputs = self._get_valid_inputs() 

202 for state_key, param_name in self.inputs_from_state.items(): 

203 if not isinstance(param_name, str): 

204 raise TypeError( 

205 f"inputs_from_state values must be str, not {type(param_name).__name__}. " 

206 f"Got {param_name!r} for key '{state_key}'." 

207 ) 

208 if valid_inputs and param_name not in valid_inputs: 

209 raise ValueError( 

210 f"inputs_from_state maps '{state_key}' to unknown parameter '{param_name}'. " 

211 f"Valid parameters are: {valid_inputs}." 

212 ) 

213 

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

215 """ 

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

217 

218 Used to validate that `inputs_from_state` only references parameters that actually exist. 

219 This prevents typos and catches configuration errors at tool construction time. 

220 

221 By default, introspects the function signature to get ALL parameters, including those 

222 that may be excluded from the JSON schema (e.g., parameters mapped from state). 

223 Falls back to schema properties if introspection fails. 

224 

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

226 

227 :returns: Set of valid input parameter names for validation. 

228 """ 

229 # Combine parameters from both function signature and schema for robustness 

230 # Function signature includes all parameters (even those excluded from schema) 

231 # Schema properties provide the validated parameter set 

232 valid_params: set[str] = set() 

233 

234 # Try to get parameters from function introspection. 

235 # Prefer `function`; fall back to `async_function` for async-only tools. 

236 introspection_target = self.function if self.function is not None else self.async_function 

237 if introspection_target is not None: 

238 try: 

239 sig = inspect.signature(introspection_target) 

240 valid_params.update(sig.parameters.keys()) 

241 except (ValueError, TypeError): 

242 pass # Introspection failed, will rely on schema 

243 

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

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

246 

247 return valid_params 

248 

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

250 """ 

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

252 

253 Used to validate that `outputs_to_state` only references outputs that actually exist. 

254 This prevents typos and catches configuration errors at tool construction time. 

255 

256 By default, returns None because regular function-based tools don't have a formal 

257 output schema. When None is returned, output validation is skipped. 

258 

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

260 enabling validation for tools where outputs are known. 

261 

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

263 """ 

264 return None 

265 

266 @property 

267 def tool_spec(self) -> dict[str, Any]: 

268 """ 

269 Return the Tool specification to be used by the Language Model. 

270 """ 

271 return {"name": self.name, "description": self.description, "parameters": self.parameters} 

272 

273 def warm_up(self) -> None: 

274 """ 

275 Prepare the Tool for use. 

276 

277 Override this method to establish connections to remote services, load models, 

278 or perform other resource-intensive initialization. This method should be idempotent, 

279 as it may be called multiple times. 

280 """ 

281 pass 

282 

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

284 """ 

285 Invoke the Tool synchronously with the provided keyword arguments. 

286 

287 :raises ToolInvocationError: If the Tool has no sync `function`, or if the underlying call 

288 raises an exception. 

289 """ 

290 if self.function is None: 

291 raise ToolInvocationError( 

292 f"Tool `{self.name}` has no sync `function` and can only be invoked via `invoke_async` " 

293 f"(use `Agent.run_async`).", 

294 tool_name=self.name, 

295 ) 

296 

297 try: 

298 result = self.function(**kwargs) 

299 except Exception as e: 

300 raise ToolInvocationError( 

301 f"Failed to invoke Tool `{self.name}` with parameters {kwargs}. Error: {e}", tool_name=self.name 

302 ) from e 

303 return result 

304 

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

306 """ 

307 Invoke the Tool asynchronously with the provided keyword arguments. 

308 

309 If `async_function` is set, it is awaited directly. Otherwise the sync `function` is dispatched to a worker 

310 thread via `asyncio.to_thread`, which propagates the current context to the worker. 

311 

312 :raises ToolInvocationError: If the underlying call raises an exception. 

313 """ 

314 try: 

315 if self.async_function is not None: 

316 return await self.async_function(**kwargs) 

317 # `function` is guaranteed to be set: __post_init__ enforces at least one of the two. 

318 return await asyncio.to_thread(self.function, **kwargs) # type: ignore[arg-type] 

319 except Exception as e: 

320 raise ToolInvocationError( 

321 f"Failed to invoke Tool `{self.name}` with parameters {kwargs}. Error: {e}", tool_name=self.name 

322 ) from e 

323 

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

325 """ 

326 Serializes the Tool to a dictionary. 

327 

328 :returns: 

329 Dictionary with serialized data. 

330 """ 

331 data = asdict(self) 

332 data["function"] = serialize_callable(self.function) if self.function is not None else None 

333 data["async_function"] = serialize_callable(self.async_function) if self.async_function is not None else None 

334 

335 if self.outputs_to_state is not None: 

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

337 

338 if self.outputs_to_string is not None: 

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

340 

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

342 

343 @classmethod 

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

345 """ 

346 Deserializes the Tool from a dictionary. 

347 

348 :param data: 

349 Dictionary to deserialize from. 

350 :returns: 

351 Deserialized Tool. 

352 """ 

353 init_parameters = data["data"] 

354 init_parameters["function"] = ( 

355 deserialize_callable(init_parameters["function"]) if init_parameters.get("function") is not None else None 

356 ) 

357 if init_parameters.get("async_function") is not None: 

358 init_parameters["async_function"] = deserialize_callable(init_parameters["async_function"]) 

359 if "outputs_to_state" in init_parameters and init_parameters["outputs_to_state"]: 

360 init_parameters["outputs_to_state"] = _deserialize_outputs_to_state(init_parameters["outputs_to_state"]) 

361 

362 if init_parameters.get("outputs_to_string") is not None: 

363 init_parameters["outputs_to_string"] = _deserialize_outputs_to_string(init_parameters["outputs_to_string"]) 

364 

365 return cls(**init_parameters) 

366 

367 

368def _check_duplicate_tool_names(tools: list[Tool] | None) -> None: 

369 """ 

370 Checks for duplicate tool names and raises a ValueError if they are found. 

371 

372 :param tools: The list of tools to check. 

373 :raises ValueError: If duplicate tool names are found. 

374 """ 

375 if tools is None: 

376 return 

377 tool_names = [tool.name for tool in tools] 

378 duplicate_tool_names = {name for name in tool_names if tool_names.count(name) > 1} 

379 if duplicate_tool_names: 

380 raise ValueError(f"Duplicate tool names found: {duplicate_tool_names}") 

381 

382 

383def _convert_handler(config: dict[str, Any], converter: Callable[[Any], Any]) -> dict[str, Any]: 

384 """ 

385 Copies a single output config, converting its "handler" entry (if present) via `converter`. 

386 

387 :param config: A single output configuration dictionary that may contain a "handler" key. 

388 :param converter: `serialize_callable` or `deserialize_callable`, applied to the "handler" value. 

389 :returns: A copy of `config` with the "handler" value converted, if present. 

390 """ 

391 new_config = config.copy() 

392 if "handler" in config: 

393 new_config["handler"] = converter(config["handler"]) 

394 return new_config 

395 

396 

397def _convert_handler_in_configs( 

398 configs: dict[str, dict[str, Any]], converter: Callable[[Any], Any] 

399) -> dict[str, dict[str, Any]]: 

400 """ 

401 Applies `_convert_handler` to every config in a dictionary of named output configs. 

402 

403 :param configs: A mapping of keys to output configuration dictionaries. 

404 :param converter: `serialize_callable` or `deserialize_callable`, applied to each "handler" value. 

405 :returns: A new mapping with the same keys, each config converted via `_convert_handler`. 

406 """ 

407 return {key: _convert_handler(config, converter) for key, config in configs.items()} 

408 

409 

410def _serialize_outputs_to_state(outputs_to_state: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]: 

411 """ 

412 Serializes the outputs_to_state dictionary, converting any callable handlers to their string representation. 

413 

414 :param outputs_to_state: The outputs_to_state dictionary to serialize. 

415 :returns: The serialized outputs_to_state dictionary. 

416 """ 

417 return _convert_handler_in_configs(outputs_to_state, serialize_callable) 

418 

419 

420def _deserialize_outputs_to_state(outputs_to_state: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]: 

421 """ 

422 Deserializes the outputs_to_state dictionary, converting any string handlers back to callables. 

423 

424 :param outputs_to_state: The outputs_to_state dictionary to deserialize. 

425 :returns: The deserialized outputs_to_state dictionary. 

426 """ 

427 return _convert_handler_in_configs(outputs_to_state, deserialize_callable) 

428 

429 

430def _serialize_outputs_to_string(outputs_to_string: dict[str, Any]) -> dict[str, Any]: 

431 """ 

432 Serializes the outputs_to_string dictionary, converting any callable handlers to their string representation. 

433 

434 :param outputs_to_string: The outputs_to_string dictionary to serialize. 

435 :returns: The serialized outputs_to_string dictionary. 

436 """ 

437 if "source" in outputs_to_string or "handler" in outputs_to_string or "raw_result" in outputs_to_string: 

438 # Single output configuration 

439 return _convert_handler(outputs_to_string, serialize_callable) 

440 

441 # Multiple outputs configuration 

442 return _convert_handler_in_configs(outputs_to_string, serialize_callable) 

443 

444 

445def _deserialize_outputs_to_string(outputs_to_string: dict[str, Any]) -> dict[str, Any]: 

446 """ 

447 Deserializes the outputs_to_string dictionary, converting any string handlers back to callables. 

448 

449 :param outputs_to_string: The outputs_to_string dictionary to deserialize. 

450 :returns: The deserialized outputs_to_string dictionary. 

451 """ 

452 if "source" in outputs_to_string or "handler" in outputs_to_string or "raw_result" in outputs_to_string: 

453 # Single output configuration 

454 return _convert_handler(outputs_to_string, deserialize_callable) 

455 

456 # Multiple outputs configuration 

457 return _convert_handler_in_configs(outputs_to_string, deserialize_callable)