Coverage for haystack/tools/pipeline_tool.py: 93%

30 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 

5from collections.abc import Callable 

6from typing import Any 

7 

8from haystack import Pipeline, SuperComponent, logging 

9from haystack.core.serialization import generate_qualified_class_name 

10from haystack.tools.component_tool import ComponentTool 

11from haystack.tools.tool import ( 

12 _deserialize_outputs_to_state, 

13 _deserialize_outputs_to_string, 

14 _serialize_outputs_to_state, 

15 _serialize_outputs_to_string, 

16) 

17 

18logger = logging.getLogger(__name__) 

19 

20 

21class PipelineTool(ComponentTool): 

22 """ 

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

24 

25 PipelineTool automatically generates LLM-compatible tool schemas from pipeline input sockets, 

26 which are derived from the underlying components in the pipeline. 

27 

28 Key features: 

29 - Automatic LLM tool calling schema generation from pipeline inputs 

30 - Description extraction of pipeline inputs based on the underlying component docstrings 

31 

32 To use PipelineTool, you first need a Haystack pipeline. 

33 Below is an example of creating a PipelineTool 

34 

35 ## Usage Example: 

36 

37 ```python 

38 from haystack import Document, Pipeline 

39 from haystack.dataclasses import ChatMessage 

40 from haystack.document_stores.in_memory import InMemoryDocumentStore 

41 from haystack.components.embedders import OpenAITextEmbedder, OpenAIDocumentEmbedder 

42 from haystack.components.generators.chat import OpenAIChatGenerator 

43 from haystack.components.retrievers import InMemoryEmbeddingRetriever 

44 from haystack.components.agents import Agent 

45 from haystack.tools import PipelineTool 

46 

47 # Initialize a document store and add some documents 

48 document_store = InMemoryDocumentStore() 

49 document_embedder = OpenAIDocumentEmbedder() 

50 documents = [ 

51 Document(content="Nikola Tesla was a Serbian-American inventor and electrical engineer."), 

52 Document( 

53 content="He is best known for his contributions to the design of the modern alternating current (AC) " 

54 "electricity supply system." 

55 ), 

56 ] 

57 docs_with_embeddings = document_embedder.run(documents=documents)["documents"] 

58 document_store.write_documents(docs_with_embeddings) 

59 

60 # Build a simple retrieval pipeline 

61 retrieval_pipeline = Pipeline() 

62 retrieval_pipeline.add_component("embedder", OpenAITextEmbedder()) 

63 retrieval_pipeline.add_component("retriever", InMemoryEmbeddingRetriever(document_store=document_store)) 

64 

65 retrieval_pipeline.connect("embedder.embedding", "retriever.query_embedding") 

66 

67 # Wrap the pipeline as a tool 

68 retriever_tool = PipelineTool( 

69 pipeline=retrieval_pipeline, 

70 input_mapping={"query": ["embedder.text"]}, 

71 output_mapping={"retriever.documents": "documents"}, 

72 name="document_retriever", 

73 description="For any questions about Nikola Tesla, always use this tool", 

74 ) 

75 

76 # Create an Agent with the tool 

77 agent = Agent( 

78 chat_generator=OpenAIChatGenerator(model="gpt-4.1-mini"), 

79 tools=[retriever_tool] 

80 ) 

81 

82 # Let the Agent handle a query 

83 result = agent.run([ChatMessage.from_user("Who was Nikola Tesla?")]) 

84 

85 # Print result of the tool call 

86 print("Tool Call Result:") 

87 print(result["messages"][2].tool_call_result.result) 

88 print("") 

89 

90 # Print answer 

91 print("Answer:") 

92 print(result["messages"][-1].text) 

93 ``` 

94 """ 

95 

96 def __init__( 

97 self, 

98 pipeline: Pipeline, 

99 *, 

100 name: str, 

101 description: str, 

102 input_mapping: dict[str, list[str]] | None = None, 

103 output_mapping: dict[str, str] | None = None, 

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

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

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

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

108 ) -> None: 

109 """ 

110 Create a Tool instance from a Haystack pipeline. 

111 

112 :param pipeline: The Haystack pipeline to wrap as a tool. 

113 :param name: Name of the tool. 

114 :param description: Description of the tool. 

115 :param input_mapping: A dictionary mapping component input names to pipeline input socket paths. 

116 If not provided, a default input mapping will be created based on all pipeline inputs. 

117 Example: 

118 ```python 

119 input_mapping={ 

120 "query": ["retriever.query", "prompt_builder.query"], 

121 } 

122 ``` 

123 :param output_mapping: A dictionary mapping pipeline output socket paths to component output names. 

124 If not provided, a default output mapping will be created based on all pipeline outputs. 

125 Example: 

126 ```python 

127 output_mapping={ 

128 "retriever.documents": "documents", 

129 "generator.replies": "replies", 

130 } 

131 ``` 

132 :param parameters: 

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

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

135 :param outputs_to_string: 

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

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

138 

139 `outputs_to_string` supports two formats: 

140 

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

142 ```python 

143 { 

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

145 } 

146 ``` 

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

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

149 final result. 

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

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

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

153 ensure compatibility with Chat Generators. 

154 

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

156 ```python 

157 { 

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

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

160 } 

161 ``` 

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

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

164 :param inputs_from_state: 

165 Optional dictionary mapping state keys to tool parameter names. 

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

167 :param outputs_to_state: 

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

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

170 Example: 

171 ```python 

172 { 

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

174 } 

175 ``` 

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

177 Example: 

178 ```python 

179 { 

180 "documents": {"handler": custom_handler} 

181 } 

182 ``` 

183 :raises ValueError: If the provided pipeline is not a valid Haystack Pipeline instance. 

184 """ 

185 if not isinstance(pipeline, Pipeline): 

186 raise TypeError(f"The 'pipeline' parameter must be an instance of Pipeline. Got {type(pipeline)} instead.") 

187 

188 super().__init__( 

189 component=SuperComponent(pipeline=pipeline, input_mapping=input_mapping, output_mapping=output_mapping), 

190 name=name, 

191 description=description, 

192 parameters=parameters, 

193 outputs_to_string=outputs_to_string, 

194 inputs_from_state=inputs_from_state, 

195 outputs_to_state=outputs_to_state, 

196 ) 

197 self._unresolved_parameters = parameters 

198 self._pipeline = pipeline 

199 self._input_mapping = input_mapping 

200 self._output_mapping = output_mapping 

201 

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

203 """ 

204 Serializes the PipelineTool to a dictionary. 

205 

206 :returns: 

207 The serialized dictionary representation of PipelineTool. 

208 """ 

209 serialized: dict[str, Any] = { 

210 "pipeline": self._pipeline.to_dict(), 

211 "name": self.name, 

212 "input_mapping": self._input_mapping, 

213 "output_mapping": self._output_mapping, 

214 "description": self.description, 

215 "parameters": self._unresolved_parameters, 

216 "inputs_from_state": self.inputs_from_state, 

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

218 "outputs_to_string": _serialize_outputs_to_string(self.outputs_to_string) 

219 if self.outputs_to_string 

220 else None, 

221 } 

222 

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

224 

225 @classmethod 

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

227 """ 

228 Deserializes the PipelineTool from a dictionary. 

229 

230 :param data: The dictionary representation of PipelineTool. 

231 :returns: 

232 The deserialized PipelineTool instance. 

233 """ 

234 inner_data = data["data"] 

235 # `is_pipeline_async` is a legacy key kept only for backward compatibility 

236 inner_data.pop("is_pipeline_async", None) 

237 pipeline = Pipeline.from_dict(inner_data["pipeline"]) 

238 

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

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

241 

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

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

244 

245 merged_data = {**inner_data, "pipeline": pipeline} 

246 return cls(**merged_data)