Coverage for haystack/tools/component_tool.py: 95%

109 statements  

« prev     ^ index     » next       coverage.py v7.16.1, created at 2026-10-09 14:34 +0000

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from collections.abc import Callable 

6from typing import Any, get_args, get_origin 

7 

8from pydantic import Field, TypeAdapter, create_model 

9 

10from haystack import logging 

11from haystack.core.component import Component 

12from haystack.core.serialization import component_to_dict, generate_qualified_class_name 

13from haystack.tools import Tool 

14from haystack.tools.errors import SchemaGenerationError 

15from haystack.tools.from_function import _remove_title_from_schema 

16from haystack.tools.parameters_schema_utils import ( 

17 _contains_callable_type, 

18 _get_component_param_descriptions, 

19 _is_state_annotation, 

20 _resolve_type, 

21 _unwrap_optional, 

22) 

23from haystack.tools.tool import ( 

24 _deserialize_outputs_to_state, 

25 _deserialize_outputs_to_string, 

26 _serialize_outputs_to_state, 

27 _serialize_outputs_to_string, 

28) 

29from haystack.utils.deserialization import deserialize_component_inplace 

30from haystack.utils.type_serialization import _is_union_type 

31 

32logger = logging.getLogger(__name__) 

33 

34 

35class ComponentTool(Tool): 

36 """ 

37 A Tool that wraps Haystack components, allowing them to be used as tools by LLMs. 

38 

39 ComponentTool automatically generates LLM-compatible tool schemas from component input sockets, 

40 which are derived from the component's `run` method signature and type hints. 

41 

42 

43 Key features: 

44 - Automatic LLM tool calling schema generation from component input sockets 

45 - Type conversion and validation for component inputs 

46 - Support for types: 

47 - Dataclasses 

48 - Lists of dataclasses 

49 - Basic types (str, int, float, bool, dict) 

50 - Lists of basic types 

51 - Automatic name generation from component class name 

52 - Description extraction from component docstrings 

53 

54 To use ComponentTool, you first need a Haystack component - either an existing one or a new one you create. 

55 You can create a ComponentTool from the component by passing the component to the ComponentTool constructor. 

56 Below is an example of creating a ComponentTool from an existing SerperDevWebSearch component 

57 from the `serperdev-haystack` integration package (`pip install serperdev-haystack`). 

58 

59 ## Usage Example: 

60 <!-- test-ignore --> 

61 ```python 

62 from haystack import component 

63 from haystack.tools import ComponentTool 

64 from haystack.utils import Secret 

65 from haystack.components.agents import Agent 

66 from haystack.components.generators.chat import OpenAIResponsesChatGenerator 

67 from haystack.dataclasses import ChatMessage 

68 from haystack_integrations.components.websearch.serperdev import SerperDevWebSearch 

69 

70 # Create a SerperDev search component 

71 search = SerperDevWebSearch(api_key=Secret.from_env_var("SERPERDEV_API_KEY"), top_k=3) 

72 

73 # Create a tool from the component 

74 tool = ComponentTool( 

75 component=search, 

76 name="web_search", # Optional: defaults to "serper_dev_web_search" 

77 description="Search the web for current information on any topic" # Optional: defaults to component docstring 

78 ) 

79 

80 # Create an Agent with an OpenAIResponsesChatGenerator and the tool 

81 agent = Agent(chat_generator=OpenAIResponsesChatGenerator(), tools=[tool]) 

82 

83 message = ChatMessage.from_user("Use the web search tool to find information about Nikola Tesla") 

84 

85 # Run the Agent 

86 result = agent.run(messages=[message]) 

87 

88 print(result) 

89 ``` 

90 

91 """ 

92 

93 def __init__( 

94 self, 

95 component: Component, 

96 name: str | None = None, 

97 description: str | None = None, 

98 parameters: dict[str, Any] | None = None, 

99 *, 

100 outputs_to_string: dict[str, str | Callable[[Any], str]] | None = None, 

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

102 outputs_to_state: dict[str, dict[str, str | Callable]] | None = None, 

103 ) -> None: 

104 """ 

105 Create a Tool instance from a Haystack component. 

106 

107 :param component: The Haystack component to wrap as a tool. 

108 :param name: Optional name for the tool (defaults to snake_case of component class name). 

109 :param description: Optional description (defaults to component's docstring). 

110 :param parameters: 

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

112 Will fall back to the parameters defined in the component's run method signature if not provided. 

113 :param outputs_to_string: 

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

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

116 

117 `outputs_to_string` supports two formats: 

118 

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

120 ```python 

121 { 

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

123 } 

124 ``` 

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

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

127 final result. 

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

129 `handler` if provided. This is intended for tools that return images. In this mode, the Tool 

130 function or the `handler` function must return a list of `TextContent`/`ImageContent` objects to 

131 ensure compatibility with Chat Generators. 

132 

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

134 ```python 

135 { 

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

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

138 } 

139 ``` 

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

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

142 :param inputs_from_state: 

143 Optional dictionary mapping state keys to tool parameter names. 

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

145 :param outputs_to_state: 

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

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

148 Example: 

149 ```python 

150 { 

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

152 } 

153 ``` 

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

155 Example: 

156 ```python 

157 { 

158 "documents": {"handler": custom_handler} 

159 } 

160 ``` 

161 :raises TypeError: If the object passed is not a Haystack Component instance. 

162 :raises ValueError: If the component has already been added to a pipeline, or if schema generation fails. 

163 """ 

164 if not isinstance(component, Component): 

165 message = ( 

166 f"Object {component!r} is not a Haystack component. " 

167 "Use ComponentTool only with Haystack component instances." 

168 ) 

169 raise TypeError(message) 

170 

171 if getattr(component, "__haystack_added_to_pipeline__", None): 

172 msg = ( 

173 "Component has been added to a pipeline and can't be used to create a ComponentTool. " 

174 "Create ComponentTool from a non-pipeline component instead." 

175 ) 

176 raise ValueError(msg) 

177 

178 self._unresolved_parameters = parameters 

179 # Create the tools schema from the component run method parameters 

180 tool_schema = parameters or self._create_tool_parameters_schema(component, inputs_from_state or {}) 

181 

182 def component_invoker(**kwargs: Any) -> dict[str, Any]: 

183 """ 

184 Invokes the component using keyword arguments provided by the LLM function calling/tool-generated response. 

185 

186 :param kwargs: The keyword arguments to invoke the component with. 

187 :returns: The result of the component invocation. 

188 """ 

189 input_sockets = component.__haystack_input__._sockets_dict # type: ignore[attr-defined] 

190 converted_kwargs = { 

191 param_name: self._convert_param(param_value, input_sockets[param_name].type) 

192 for param_name, param_value in kwargs.items() 

193 } 

194 logger.debug( 

195 "Invoking component {component_type} with kwargs: {converted_kwargs}", 

196 component_type=type(component), 

197 converted_kwargs=converted_kwargs, 

198 ) 

199 return dict(component.run(**converted_kwargs)) 

200 

201 async def async_component_invoker(**kwargs: Any) -> dict[str, Any]: 

202 """ 

203 Asynchronous counterpart of `component_invoker`. Awaits the component's `run_async`. 

204 

205 :param kwargs: The keyword arguments to invoke the component with. 

206 :returns: The result of the component invocation. 

207 """ 

208 input_sockets = component.__haystack_input__._sockets_dict # type: ignore[attr-defined] 

209 converted_kwargs = { 

210 param_name: self._convert_param(param_value, input_sockets[param_name].type) 

211 for param_name, param_value in kwargs.items() 

212 } 

213 logger.debug( 

214 "Invoking component {component_type} asynchronously with kwargs: {converted_kwargs}", 

215 component_type=type(component), 

216 converted_kwargs=converted_kwargs, 

217 ) 

218 # We know run_async exists at this point b/c we only pass the async invoker if the component has 

219 # __haystack_supports_async__ = True 

220 return dict(await component.run_async(**converted_kwargs)) # type: ignore[attr-defined] 

221 

222 component_supports_async = getattr(component, "__haystack_supports_async__", False) 

223 

224 # Generate a name for the tool if not provided 

225 if not name: 

226 class_name = component.__class__.__name__ 

227 # Convert camelCase/PascalCase to snake_case 

228 name = "".join( 

229 [ 

230 "_" + c.lower() if c.isupper() and i > 0 and not class_name[i - 1].isupper() else c.lower() 

231 for i, c in enumerate(class_name) 

232 ] 

233 ).lstrip("_") 

234 

235 description = description or component.__doc__ or name 

236 

237 # Store component before calling super().__init__() so _get_valid_outputs() can access it 

238 self._component = component 

239 

240 # Create the Tool instance with the component invoker as the function to be called and the schema. 

241 # When the wrapped component exposes a `run_async`, also pass the async invoker. 

242 super().__init__( 

243 name=name, 

244 description=description, 

245 parameters=tool_schema, 

246 function=component_invoker, 

247 async_function=async_component_invoker if component_supports_async else None, 

248 inputs_from_state=inputs_from_state, 

249 outputs_to_state=outputs_to_state, 

250 outputs_to_string=outputs_to_string, 

251 ) 

252 

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

254 """ 

255 Return valid input parameter names from the component's input sockets. 

256 

257 Used to validate `inputs_from_state` against the component's actual inputs. 

258 This ensures users don't reference non-existent component inputs. 

259 

260 :returns: Set of component input socket names. 

261 """ 

262 return set(self._component.__haystack_input__._sockets_dict.keys()) # type: ignore[attr-defined] 

263 

264 def _get_valid_outputs(self) -> set[str]: 

265 """ 

266 Return valid output names from the component's output sockets. 

267 

268 Used to validate `outputs_to_state` against the component's actual outputs. 

269 This ensures users don't reference non-existent component outputs. 

270 

271 :returns: Set of component output socket names. 

272 """ 

273 return set(self._component.__haystack_output__._sockets_dict.keys()) # type: ignore[attr-defined] 

274 

275 def warm_up(self) -> None: 

276 """Warm up the wrapped component.""" 

277 if hasattr(self._component, "warm_up"): 

278 self._component.warm_up() 

279 

280 async def warm_up_async(self) -> None: 

281 """Warm up the wrapped component on the serving event loop.""" 

282 if self.async_function is not None and hasattr(self._component, "warm_up_async"): 

283 await self._component.warm_up_async() 

284 elif hasattr(self._component, "warm_up"): 

285 self._component.warm_up() 

286 

287 def close(self) -> None: 

288 """Release the wrapped component's synchronous resources.""" 

289 if hasattr(self._component, "close"): 

290 self._component.close() 

291 

292 async def close_async(self) -> None: 

293 """Release the wrapped component's async resources.""" 

294 if self.async_function is not None and hasattr(self._component, "close_async"): 

295 await self._component.close_async() 

296 elif hasattr(self._component, "close"): 

297 self._component.close() 

298 

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

300 """ 

301 Serializes the ComponentTool to a dictionary. 

302 """ 

303 serialized: dict[str, Any] = { 

304 "component": component_to_dict(obj=self._component, name=self.name), 

305 "name": self.name, 

306 "description": self.description, 

307 "parameters": self._unresolved_parameters, 

308 "inputs_from_state": self.inputs_from_state, 

309 "outputs_to_state": _serialize_outputs_to_state(self.outputs_to_state) if self.outputs_to_state else None, 

310 "outputs_to_string": _serialize_outputs_to_string(self.outputs_to_string) 

311 if self.outputs_to_string 

312 else None, 

313 } 

314 

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

316 

317 @classmethod 

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

319 """ 

320 Deserializes the ComponentTool from a dictionary. 

321 """ 

322 inner_data = data["data"] 

323 deserialize_component_inplace(data=inner_data, key="component") 

324 

325 if "outputs_to_state" in inner_data and inner_data["outputs_to_state"]: 

326 inner_data["outputs_to_state"] = _deserialize_outputs_to_state(inner_data["outputs_to_state"]) 

327 

328 if inner_data.get("outputs_to_string") is not None: 

329 inner_data["outputs_to_string"] = _deserialize_outputs_to_string(inner_data["outputs_to_string"]) 

330 

331 return cls( 

332 component=inner_data["component"], 

333 name=inner_data["name"], 

334 description=inner_data["description"], 

335 parameters=inner_data.get("parameters", None), 

336 outputs_to_string=inner_data.get("outputs_to_string", None), 

337 inputs_from_state=inner_data.get("inputs_from_state", None), 

338 outputs_to_state=inner_data.get("outputs_to_state", None), 

339 ) 

340 

341 def _create_tool_parameters_schema(self, component: Component, inputs_from_state: dict[str, Any]) -> dict[str, Any]: 

342 """ 

343 Creates an OpenAI tools schema from a component's run method parameters. 

344 

345 :param component: The component to create the schema from. 

346 :raises SchemaGenerationError: If schema generation fails 

347 :returns: OpenAI tools schema for the component's run method parameters. 

348 """ 

349 param_descriptions = _get_component_param_descriptions(component) 

350 

351 # collect fields (types and defaults) and descriptions from function parameters 

352 fields: dict[str, Any] = {} 

353 

354 for input_name, socket in component.__haystack_input__._sockets_dict.items(): # type: ignore[attr-defined] 

355 if inputs_from_state is not None and input_name in list(inputs_from_state.values()): 

356 continue 

357 input_type = socket.type 

358 

359 # Skip Callable types since Pydantic cannot generate JSON schemas for them 

360 if _contains_callable_type(input_type): 

361 continue 

362 

363 # Skip State-typed parameters - Agent tool execution injects them at runtime 

364 if _is_state_annotation(input_type): 

365 continue 

366 

367 description = param_descriptions.get(input_name, f"Input '{input_name}' for the component.") 

368 

369 # if the parameter has not a default value, Pydantic requires an Ellipsis (...) 

370 # to explicitly indicate that the parameter is required 

371 default = ... if socket.is_mandatory else socket.default_value 

372 resolved_type = _resolve_type(input_type) 

373 fields[input_name] = (resolved_type, Field(default=default, description=description)) 

374 

375 parameters_schema: dict[str, Any] = {} 

376 try: 

377 # No `__doc__`: it would surface as a top-level `description` on the parameters schema, 

378 # which LLM providers ignore. The component description feeds the tool-level description. 

379 model = create_model(component.run.__name__, **fields) 

380 parameters_schema = model.model_json_schema() 

381 except Exception as e: 

382 raise SchemaGenerationError( 

383 f"Failed to create JSON schema for the run method of Component '{component.__class__.__name__}'" 

384 ) from e 

385 

386 # we don't want to include title keywords in the schema, as they contain redundant information 

387 # there is no programmatic way to prevent Pydantic from adding them, so we remove them later 

388 # see https://github.com/pydantic/pydantic/discussions/8504 

389 _remove_title_from_schema(parameters_schema) 

390 

391 return parameters_schema 

392 

393 def _convert_param(self, param_value: Any, param_type: type) -> Any: 

394 """ 

395 Converts a single parameter value to the expected type. 

396 

397 :param param_value: The value to convert. 

398 :param param_type: The expected type of the parameter. 

399 

400 :returns: 

401 The converted parameter value. 

402 """ 

403 # We unwrap optional types so we can support types like messages: list[ChatMessage] | None 

404 unwrapped_param_type = _unwrap_optional(param_type) 

405 

406 # We handle union types (e.g. list[ChatMessage] | str) by extracting the list[T] arm that has 

407 # from_dict, so the conversion below works the same as for plain list[T]. Other arms like str 

408 # need no special handling and fall through to Pydantic or the plain return at the end. 

409 if _is_union_type(unwrapped_param_type): 

410 list_arms = [ 

411 a 

412 for a in get_args(unwrapped_param_type) 

413 if get_origin(a) is list and get_args(a) and hasattr(get_args(a)[0], "from_dict") 

414 ] 

415 unwrapped_param_type = list_arms[0] if list_arms else unwrapped_param_type 

416 

417 # We support calling from_dict on target types that have it, even if they are wrapped in a list. 

418 # This allows us to support lists of dataclasses as well as single dataclass inputs. 

419 target_type = ( 

420 get_args(unwrapped_param_type)[0] if get_origin(unwrapped_param_type) is list else unwrapped_param_type 

421 ) 

422 if hasattr(target_type, "from_dict"): 

423 if isinstance(param_value, list): 

424 return [target_type.from_dict(item) if isinstance(item, dict) else item for item in param_value] 

425 if isinstance(param_value, dict): 

426 return target_type.from_dict(param_value) 

427 return param_value 

428 

429 # Use the original type for pydantic validation 

430 return TypeAdapter(param_type).validate_python(param_value)