Coverage for haystack/components/agents/agent.py: 99%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import inspect 

6from dataclasses import dataclass 

7from typing import Any, Literal, cast 

8 

9from haystack import component, logging, tracing 

10from haystack.components.agents.state.state import ( 

11 State, 

12 _schema_from_dict, 

13 _schema_to_dict, 

14 _validate_schema, 

15 replace_values, 

16) 

17from haystack.components.agents.state.state_utils import merge_lists 

18from haystack.components.agents.tool_calling import _run_tool, _run_tool_async 

19from haystack.components.agents.utils import ( 

20 _record_context_tokens, 

21 _record_llm_usage, 

22 _record_tool_calls, 

23 _render_prompt_messages, 

24 _select_tools_by_name, 

25 _spawn_tools, 

26 _template_for_role, 

27 _validate_prompt_message_blocks, 

28) 

29from haystack.components.builders import ChatPromptBuilder 

30from haystack.components.generators.chat.types import ChatGenerator 

31from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict 

32from haystack.dataclasses import ChatMessage, ChatRole, StreamingCallbackT, select_streaming_callback 

33from haystack.hooks.invocation import _run_hooks, _run_hooks_async 

34from haystack.hooks.protocol import ( 

35 AFTER_RUN, 

36 AFTER_TOOL, 

37 BEFORE_LLM, 

38 BEFORE_RUN, 

39 BEFORE_TOOL, 

40 ON_EXIT, 

41 VALID_HOOK_POINTS, 

42 Hook, 

43 HookPoint, 

44) 

45from haystack.hooks.utils import ( 

46 _deserialize_hooks_dictionary, 

47 _serialize_hooks_dictionary, 

48 close_hooks, 

49 close_hooks_async, 

50 warm_up_hooks, 

51 warm_up_hooks_async, 

52) 

53from haystack.tools import ( 

54 Toolset, 

55 ToolsType, 

56 _check_duplicate_tool_names, 

57 deserialize_tools_or_toolset_inplace, 

58 flatten_tools_or_toolsets, 

59 serialize_tools_or_toolset, 

60 warm_up_tools, 

61) 

62from haystack.utils.async_utils import _execute_component_async 

63from haystack.utils.callable_serialization import deserialize_callable, serialize_callable 

64from haystack.utils.deserialization import deserialize_component_inplace 

65 

66logger = logging.getLogger(__name__) 

67 

68# `exit_reason` values the Agent sets when it stops without a tool exit condition: a tool-call-free reply, or the 

69# `max_agent_steps` budget running out. A tool exit condition instead reports the tool's name. 

70_EXIT_REASON_TEXT = "text" 

71_EXIT_REASON_MAX_STEPS = "max_agent_steps" 

72 

73# Run-metadata state keys the Agent populates automatically during a run. Users may not define them in their own 

74# `state_schema`, and they are exposed as Agent outputs only (not inputs). 

75_RUN_METADATA_STATE_KEYS: dict[str, dict[str, Any]] = { 

76 "step_count": {"type": int, "handler": replace_values}, 

77 "token_usage": {"type": dict[str, Any], "handler": replace_values}, 

78 "tool_call_counts": {"type": dict[str, int], "handler": replace_values}, 

79 "exit_reason": {"type": str, "handler": replace_values}, 

80} 

81 

82# Internal state keys the Agent manages for run control and hooks. Like run-metadata keys they are reserved and cannot 

83# be redefined by users, but unlike them they are NOT exposed as Agent inputs or outputs (purely internal state): 

84# - `continue_run`: set by an `on_exit` hook to keep the Agent running instead of stopping (re-read each exit attempt). 

85# - `stop_run`: set by a hook to stop the run, read before each LLM call and used as the `exit_reason`. 

86# - `tools`: the flattened tools available in the current step, so a hook can inspect them (e.g. HITL confirmation). 

87# - `hook_context`: per-run request-scoped resources passed to `run`/`run_async` for hooks to read. 

88# - `context_tokens`: approximate current context-window size, refreshed after each LLM call, for hooks to read 

89# (e.g. a `before_llm` hook that triggers compaction once the context grows too large). Kept internal rather than 

90# exposed as an output because it is a best-effort snapshot; see `_record_context_tokens`. 

91_INTERNAL_STATE_KEYS: dict[str, dict[str, Any]] = { 

92 "continue_run": {"type": bool, "handler": replace_values}, 

93 "stop_run": {"type": str, "handler": replace_values}, 

94 "tools": {"type": list, "handler": replace_values}, 

95 "hook_context": {"type": dict[str, Any], "handler": replace_values}, 

96 "context_tokens": {"type": int, "handler": replace_values}, 

97} 

98 

99 

100def _get_run_method_params(instance: "Agent") -> set[str]: 

101 """Derive the parameter names of the Agent.run method via introspection.""" 

102 sig = inspect.signature(instance.run) 

103 return {name for name, p in sig.parameters.items() if p.kind != inspect.Parameter.VAR_KEYWORD} 

104 

105 

106def _public_outputs(state: State) -> dict[str, Any]: 

107 """Return the State data excluding the internal state keys (i.e. the Agent's user-facing outputs).""" 

108 return {key: value for key, value in state.data.items() if key not in _INTERNAL_STATE_KEYS} 

109 

110 

111def _validate_hooks(hooks: dict[HookPoint, list[Hook]]) -> None: 

112 """ 

113 Validate a hooks mapping: known hook points, real Hook objects, and hook-point restrictions. 

114 

115 :param hooks: Mapping of hook point to the hooks registered under it. 

116 :raises ValueError: If a hook point is unknown, or a hook is registered under a hook point it does not support. 

117 :raises TypeError: If a registered hook has no callable `run(state)`. 

118 """ 

119 for hook_point, hook_list in hooks.items(): 

120 if hook_point not in VALID_HOOK_POINTS: 

121 raise ValueError( 

122 f"Invalid hook point '{hook_point}'. Valid hook points are: {', '.join(VALID_HOOK_POINTS)}." 

123 ) 

124 for h in hook_list: 

125 if not callable(getattr(h, "run", None)): 

126 if callable(h): 

127 raise TypeError( 

128 f"Hook registered for hook point '{hook_point}' is callable but is not a Hook object. " 

129 "If it is a function, wrap it with the @hook decorator." 

130 ) 

131 raise TypeError( 

132 f"Hook registered for hook point '{hook_point}' must have a callable 'run(state)', " 

133 f"got an object of type '{type(h).__name__}'." 

134 ) 

135 # A hook may declare `allowed_hook_points` to restrict where it can run (e.g. ConfirmationHook only 

136 # makes sense at "before_tool"). Hooks without it can be registered under any hook point. 

137 allowed_points = getattr(h, "allowed_hook_points", None) 

138 if allowed_points is not None and hook_point not in allowed_points: 

139 raise ValueError( 

140 f"Hook of type '{type(h).__name__}' is registered under hook point '{hook_point}' but only " 

141 f"supports: {', '.join(allowed_points)}." 

142 ) 

143 

144 

145def _consume_continue_run(state: State) -> bool: 

146 """Return the `continue_run` control flag and reset it so it does not carry over to the next exit attempt.""" 

147 should_continue = state.data["continue_run"] 

148 state.set("continue_run", False) 

149 return should_continue 

150 

151 

152def _is_text_exit(messages: list[ChatMessage]) -> bool: 

153 """ 

154 Return whether `messages` end in a plain assistant text reply with no tool calls anywhere in the batch. 

155 

156 This is the "no tool call" exit for the model's own replies. The last message must be a non-empty assistant text 

157 message, so an invalid response (e.g. one with no tool calls and no text) does not trigger an exit. 

158 """ 

159 if not messages: 

160 return False 

161 last = messages[-1] 

162 return not any(m.tool_call for m in messages) and last.is_from(ChatRole.ASSISTANT) and bool(last.text) 

163 

164 

165def _pending_tool_call_messages_from_state(state: State) -> list[ChatMessage]: 

166 """ 

167 Return the pending tool-call message after `before_tool` hooks have run. 

168 

169 `before_tool` hooks may mutate `state.data["messages"]`. After they run, the Agent intentionally inspects only the 

170 current last message in the state. If that message has tool calls, those calls are executed. If it has no tool 

171 calls, no tools run for the step, no tool-based exit condition is triggered, and the Agent loops back to the next 

172 LLM call unless `max_agent_steps` has been reached. 

173 """ 

174 messages = state.data.get("messages") or [] 

175 if not messages: 

176 return [] 

177 last_message = messages[-1] 

178 return [last_message] if last_message.tool_calls else [] 

179 

180 

181@dataclass(kw_only=True) 

182class _ExecutionContext: 

183 """ 

184 Context for executing the agent. 

185 

186 :param state: The current state of the agent, including messages and any additional data. 

187 :param tools: The tools selected for this run, kept unflattened (the original Toolset or list of 

188 Tools/Toolsets). Storing the unflattened form lets each step re-flatten it and pick up tools a dynamic 

189 toolset (e.g. SearchableToolset) discovers over time; flattening would freeze a snapshot. The chat 

190 generator and tool execution receive a freshly flattened snapshot per step. 

191 :param chat_generator_inputs: Runtime inputs to be passed to the chat generator (tools are injected per step). 

192 :param tool_execution_inputs: Runtime inputs to be passed to tool execution (tools are injected per step). 

193 :param counter: A counter to track the number of steps taken in the agent's run. 

194 """ 

195 

196 state: State 

197 tools: ToolsType 

198 chat_generator_inputs: dict 

199 tool_execution_inputs: dict 

200 counter: int = 0 

201 

202 

203@component 

204class Agent: 

205 """ 

206 A tool-using Agent powered by a large language model. 

207 

208 The Agent processes messages and calls tools until it meets an exit condition. 

209 You can set one or more exit conditions to control when it stops. 

210 For example, it can stop after generating a response or after calling a tool. 

211 

212 Without tools, the Agent works like a standard LLM that generates text. It produces one response and then stops. 

213 

214 ### Usage examples 

215 

216 This is an example agent that: 

217 1. Searches for tipping customs in France. 

218 2. Uses a calculator to compute tips based on its findings. 

219 3. Returns the final answer with its context. 

220 

221 ```python 

222 from haystack.components.agents import Agent 

223 from haystack.components.generators.chat import OpenAIChatGenerator 

224 from haystack.components.generators.utils import print_streaming_chunk 

225 from haystack.dataclasses import ChatMessage 

226 from haystack.tools import tool 

227 from typing import Annotated, Literal 

228 

229 # Tool functions - in practice, these would have real implementations 

230 @tool 

231 def search(query: Annotated[str, "The search query"]) -> str: 

232 '''Search for information on the web.''' 

233 # Placeholder: would call actual search API 

234 return "In France, a 15% service charge is typically included, but leaving 5-10% extra is appreciated." 

235 

236 @tool 

237 def calculator( 

238 operation: Annotated[Literal["multiply", "percentage"], "The mathematical operation to perform"], 

239 a: Annotated[float, "First number"], 

240 b: Annotated[float, "Second number"], 

241 ) -> float: 

242 '''Perform mathematical calculations.''' 

243 if operation == "multiply": 

244 return a * b 

245 elif operation == "percentage": 

246 return (a / 100) * b 

247 return 0 

248 

249 agent = Agent( 

250 system_prompt=( 

251 "You are a helpful assistant. Use the 'search' tool to find information " 

252 "about a user's question and the 'calculator' tool to perform math." 

253 ), 

254 chat_generator=OpenAIChatGenerator(), 

255 tools=[search, calculator], 

256 streaming_callback=print_streaming_chunk, 

257 ) 

258 

259 result = agent.run( 

260 messages=[ChatMessage.from_user("Calculate the appropriate tip for an €85 meal in France")] 

261 ) 

262 

263 # Access the final response from the Agent 

264 # print(result["last_message"].text) 

265 ``` 

266 

267 #### Using a `user_prompt` template with variables 

268 

269 You can define a reusable `user_prompt` with Jinja2 template variables so the Agent can be invoked 

270 with different inputs without manually constructing `ChatMessage` objects each time. 

271 This is especially useful when embedding the Agent in a pipeline. 

272 

273 ```python 

274 from haystack.components.agents import Agent 

275 from haystack.components.generators.chat import OpenAIChatGenerator 

276 from haystack.tools import tool 

277 from typing import Annotated 

278 

279 

280 @tool 

281 def translate( 

282 text: Annotated[str, "The text to translate"], 

283 target_language: Annotated[str, "The language to translate to"], 

284 ) -> str: 

285 \"\"\"Translate text to a target language.\"\"\" 

286 # Placeholder: would call an actual translation API 

287 return f"[Translated '{text}' to {target_language}]" 

288 

289 agent = Agent( 

290 chat_generator=OpenAIChatGenerator(), 

291 tools=[translate], 

292 system_prompt="You are a helpful translation assistant.", 

293 user_prompt=\"\"\"{% message role="user"%} 

294 Translate the following document to {{ language }}: {{ document }} 

295 {% endmessage %}\"\"\", 

296 ) 

297 

298 # The template variables 'language' and 'document' become inputs to the run method 

299 result = agent.run( 

300 messages=[], 

301 language="French", 

302 document="The weather is lovely today and the sun is shining.", 

303 ) 

304 

305 print(result["last_message"].text) 

306 ``` 

307 

308 #### Using hooks to influence the run loop 

309 

310 Hooks are callables that receive the live `State` and run at specific points in the Agent loop: 

311 

312 - `before_llm`: runs before each chat-generator call. 

313 - `before_tool`: runs after the model requests tool calls, before any tools run. After these hooks run, the Agent 

314 re-reads the current last message from `state.data["messages"]`. If that message has tool calls, those calls are 

315 executed. If it has no tool calls, no tools run for that step, no tool-based exit condition is triggered, and the 

316 Agent loops back to the next LLM call unless `max_agent_steps` has been reached. 

317 - `after_tool`: runs after tools execute, once their result messages are in `state.data["messages"]`, before the 

318 exit check and the next LLM call. Use it to rewrite the freshly produced tool-result messages (e.g. offload, 

319 redact, truncate, or summarize results). It does not run on the plain-text exit step, where no tools run. 

320 - `on_exit`: runs when the Agent is about to stop on an exit condition. An `on_exit` hook can keep the Agent 

321 running by setting `state.set("continue_run", True)`. 

322 

323 Use the `@hook` decorator to build a hook from a function. This `on_exit` hook keeps the Agent running until a 

324 required tool has been called. 

325 

326 ```python 

327 from haystack.components.agents import Agent 

328 from haystack.components.agents.state import State 

329 from haystack.components.generators.chat import OpenAIChatGenerator 

330 from haystack.dataclasses import ChatMessage 

331 from haystack.hooks import hook 

332 from haystack.tools import tool 

333 from typing import Annotated 

334 

335 

336 @tool 

337 def save_result(content: Annotated[str, "The result to save"]) -> str: 

338 \"\"\"Save the final result.\"\"\" 

339 # Placeholder: would persist `content` to a database or the file system 

340 return "saved" 

341 

342 

343 @hook 

344 def require_save(state: State) -> None: 

345 if state.get("tool_call_counts", {}).get("save_result", 0) == 0: 

346 state.set("messages", [ChatMessage.from_system("Call `save_result` before finishing.")]) 

347 state.set("continue_run", True) # keep the Agent running instead of stopping 

348 

349 

350 agent = Agent( 

351 chat_generator=OpenAIChatGenerator(), 

352 tools=[save_result], 

353 hooks={"on_exit": [require_save]}, 

354 ) 

355 ``` 

356 

357 """ 

358 

359 def __init__( # noqa: PLR0913 

360 self, 

361 *, 

362 chat_generator: ChatGenerator, 

363 tools: ToolsType | None = None, 

364 system_prompt: str | None = None, 

365 user_prompt: str | None = None, 

366 required_variables: list[str] | Literal["*"] | None = "*", 

367 exit_conditions: list[str] | None = None, 

368 state_schema: dict[str, Any] | None = None, 

369 max_agent_steps: int = 100, 

370 streaming_callback: StreamingCallbackT | None = None, 

371 raise_on_tool_invocation_failure: bool = False, 

372 tool_concurrency_limit: int = 4, 

373 tool_streaming_callback_passthrough: bool = False, 

374 hooks: dict[HookPoint, list[Hook]] | None = None, 

375 ) -> None: 

376 """ 

377 Initialize the agent component. 

378 

379 :param chat_generator: An instance of the chat generator that your agent should use. It must support tools. 

380 :param tools: A list of Tool and/or Toolset objects, or a single Toolset that the agent can use. 

381 :param system_prompt: System prompt for the agent. Can be a plain string template or a Jinja2 message template. 

382 For details on the supported template syntax, refer to the 

383 [documentation](https://docs.haystack.deepset.ai/docs/chatpromptbuilder#string-templates). 

384 :param user_prompt: User prompt for the agent. Can be a plain string template or a Jinja2 message template. 

385 If provided, this is appended to the messages provided at runtime. 

386 For details on the supported template syntax, refer to the 

387 [documentation](https://docs.haystack.deepset.ai/docs/chatpromptbuilder#string-templates). 

388 :param required_variables: 

389 Lists the variables that must be provided as inputs to `user_prompt` or `system_prompt`. 

390 If a required variable is not provided at run time, an exception is raised. 

391 If set to `"*"`, all variables found in the prompts are required. Defaults to `"*"`. 

392 Set to `None` to make all variables optional; missing ones render as empty strings. 

393 :param exit_conditions: List of conditions that will cause the agent to return. 

394 Can include "text" if the agent should return when it generates a message without tool calls, 

395 or tool names that will cause the agent to return once the tool was executed. Defaults to ["text"]. 

396 :param state_schema: A dictionary defining the agent's runtime state. Each key maps to a type config 

397 with `"type"` (required) and an optional `"handler"` for merging values across tool calls. 

398 Tools can read from and write to state keys using `inputs_from_state` and `outputs_to_state`. 

399 :param max_agent_steps: Maximum number of steps the agent will run before stopping. Defaults to 100. 

400 A step is one chat-generator call plus the execution of every tool call the model requested in 

401 that call (if any). If the agent reaches this number of steps it stops and returns the current state. 

402 :param streaming_callback: A callback that will be invoked when a response is streamed from the LLM. 

403 The same callback can be configured to emit tool results when a tool is called. 

404 :param raise_on_tool_invocation_failure: Should the agent raise an exception when a tool invocation fails? 

405 If set to False, the exception will be turned into a chat message and passed to the LLM. 

406 :param tool_concurrency_limit: Maximum number of tool calls to execute at the same time. 

407 Defaults to 4. Set to 1 to disable parallel tool execution. 

408 :param tool_streaming_callback_passthrough: If True, pass the streaming callback to tools that accept it. 

409 :param hooks: A dictionary mapping a hook point to a list of hooks the Agent runs at that point. Each hook 

410 receives the live `State` and influences the run by mutating it in place; hooks for a hook point run in 

411 list order. Valid hook points are: 

412 - "before_run": Runs once per run, after the state is initialized and before the first chat-generator 

413 call. Use it to rewrite the initial messages or seed state (e.g. turn the user query into a task 

414 brief) without re-running on every step like "before_llm" does. 

415 - "before_llm": Runs before each chat-generator call. 

416 - "before_tool": Runs after the model requests tool calls, before any tools run. After these hooks run, 

417 the Agent re-reads the current last message from `state.data["messages"]`. If that message contains tool 

418 calls, those calls are executed. If it does not, no tools run for that step, no tool-based exit condition 

419 is triggered, and the Agent loops back to the next LLM call unless `max_agent_steps` has been reached. 

420 - "after_tool": Runs after tools execute, once their result messages are in `state.data["messages"]`, 

421 before the exit check and the next LLM call. Use it to rewrite the freshly produced tool-result messages 

422 (e.g. offload, redact, truncate, or summarize results). It does not run on the plain-text exit step, 

423 where no tools run. 

424 - "on_exit": Runs when the Agent is about to stop on an exit condition. An "on_exit" hook can keep the 

425 Agent running by setting the `continue_run` control flag (`state.set("continue_run", True)`), usually 

426 alongside a message telling the model what to do next. "on_exit" hooks run when the Agent stops on an 

427 exit condition, but not when it stops because `max_agent_steps` is reached. 

428 - "after_run": Runs once per run, after the step loop has ended and before the Agent builds its return 

429 value — regardless of whether the run stopped on an exit condition or because `max_agent_steps` was 

430 reached (unlike "on_exit"). Mutations to the state (e.g. appending a final message) are reflected in 

431 the returned `messages` / `last_message` and `state_schema` outputs. Setting `continue_run` here has 

432 no effect. 

433 :raises TypeError: If the chat_generator does not support tools parameter in its run method. 

434 :raises ValueError: If any `user_prompt` variable overlaps with the `state_schema` or `run` method parameters, 

435 if a hook is registered under an unknown hook point, or if a hook is registered under a hook point it does 

436 not support (via its `allowed_hook_points`). 

437 """ 

438 # --- Validation --- 

439 self._chat_generator_supports_tools: bool = "tools" in inspect.signature(chat_generator.run).parameters 

440 # We use an explicit None check for tools b/c testing for truthiness calls __len__, which for SearchableToolset 

441 # would iterate and prematurely warm it up at init. 

442 if tools is not None and not self._chat_generator_supports_tools: 

443 raise TypeError( 

444 f"{type(chat_generator).__name__} does not accept tools parameter in its run method. " 

445 "The Agent component requires a chat generator that supports tools when tools are provided." 

446 ) 

447 

448 if exit_conditions is None: 

449 exit_conditions = ["text"] 

450 

451 if state_schema is not None: 

452 reserved_keys = _RUN_METADATA_STATE_KEYS.keys() | _INTERNAL_STATE_KEYS.keys() 

453 reserved_used = sorted(set(state_schema) & reserved_keys) 

454 if reserved_used: 

455 raise ValueError( 

456 f"state_schema keys {reserved_used} are reserved for Agent internal state and " 

457 f"cannot be redefined. Reserved keys: {sorted(reserved_keys)}." 

458 ) 

459 _validate_schema(state_schema) 

460 _validate_prompt_message_blocks(user_prompt, system_prompt) 

461 if tool_concurrency_limit < 1: 

462 raise ValueError("tool_concurrency_limit must be greater than or equal to 1.") 

463 

464 hooks = hooks or {} 

465 _validate_hooks(hooks) 

466 

467 # --- Attributes --- 

468 self.chat_generator = chat_generator 

469 # We use an explicit None check for tools b/c testing for truthiness calls __len__, which for SearchableToolset 

470 # would iterate and prematurely warm it up at init. 

471 self.tools = tools if tools is not None else [] 

472 self.system_prompt = system_prompt 

473 self.user_prompt = user_prompt 

474 self.required_variables = required_variables 

475 self.exit_conditions = exit_conditions 

476 self.max_agent_steps = max_agent_steps 

477 self.raise_on_tool_invocation_failure = raise_on_tool_invocation_failure 

478 self.streaming_callback = streaming_callback 

479 self.tool_concurrency_limit = tool_concurrency_limit 

480 self.tool_streaming_callback_passthrough = tool_streaming_callback_passthrough 

481 self.hooks = hooks 

482 

483 # --- State schema --- 

484 # shallow copy is sufficient: we only add a top-level "messages" key, never mutate nested values 

485 self.state_schema = state_schema or {} 

486 self.resolved_state_schema = dict(self.state_schema) 

487 if self.resolved_state_schema.get("messages") is None: 

488 self.resolved_state_schema["messages"] = {"type": list[ChatMessage], "handler": merge_lists} 

489 for key, config in {**_RUN_METADATA_STATE_KEYS, **_INTERNAL_STATE_KEYS}.items(): 

490 self.resolved_state_schema[key] = dict(config) 

491 

492 # --- Component I/O --- 

493 self._run_method_params = _get_run_method_params(self) 

494 output_types: dict[str, Any] = {"last_message": ChatMessage} 

495 for param, config in self.resolved_state_schema.items(): 

496 # Internal keys are run-control / hook-facing state, not exposed as inputs or outputs. 

497 if param in _INTERNAL_STATE_KEYS: 

498 continue 

499 output_types[param] = config["type"] 

500 # Run-metadata keys are populated by the Agent itself and exposed as outputs only, not inputs. 

501 if param not in self._run_method_params and param not in _RUN_METADATA_STATE_KEYS: 

502 component.set_input_type(self, name=param, type=config["type"], default=None) 

503 component.set_output_types(self, **output_types) 

504 

505 # --- Prompt builders --- 

506 # required_variables starts empty and is populated by _register_prompt_variables once 

507 # builder.variables are known 

508 self._user_chat_prompt_builder = ( 

509 ChatPromptBuilder(template=_template_for_role(user_prompt, "user"), required_variables=[]) 

510 if user_prompt is not None 

511 else None 

512 ) 

513 self._system_chat_prompt_builder: ChatPromptBuilder | None = None 

514 if system_prompt is not None: 

515 self._system_chat_prompt_builder = ChatPromptBuilder( 

516 template=_template_for_role(system_prompt, "system"), required_variables=[] 

517 ) 

518 self._register_prompt_variables() 

519 

520 def _register_prompt_variables(self) -> None: 

521 """ 

522 Collect variables from both Chat Prompt Builders and register Agent inputs. 

523 

524 Sets `required_variables` for both Chat Prompt Builders, checks for conflicts with state schema and 

525 run parameters, and registers component inputs. 

526 """ 

527 required_variables = self.required_variables 

528 

529 prompt_builders = [ 

530 builder for builder in (self._system_chat_prompt_builder, self._user_chat_prompt_builder) if builder 

531 ] 

532 if isinstance(required_variables, list) and not any(builder.variables for builder in prompt_builders): 

533 logger.warning( 

534 "The parameter required_variables is provided but neither user_prompt nor system_prompt " 

535 "contains template variables. Either provide a prompt with Jinja2 template variables " 

536 "or remove required_variables, it has otherwise no effect." 

537 ) 

538 

539 all_variables: dict[str, list[str]] = {} 

540 for builder, label in [ 

541 (self._system_chat_prompt_builder, "system_prompt"), 

542 (self._user_chat_prompt_builder, "user_prompt"), 

543 ]: 

544 if builder is not None: 

545 # set required_variables on the builder, filtered to its own variables 

546 if required_variables == "*": 

547 builder.required_variables = "*" 

548 elif isinstance(self.required_variables, list): 

549 builder.required_variables = [v for v in self.required_variables if v in builder.variables] 

550 

551 for var_name in builder.variables: 

552 all_variables.setdefault(var_name, []).append(label) 

553 

554 for var_name, sources in all_variables.items(): 

555 prompt_source = " and ".join(sources) 

556 if var_name in self.resolved_state_schema: 

557 raise ValueError( 

558 f"Variable '{var_name}' from {prompt_source} is already defined in the state schema. " 

559 "Please rename the variable or remove it from the prompt to avoid conflicts." 

560 ) 

561 if var_name in self._run_method_params: 

562 raise ValueError( 

563 f"Variable '{var_name}' from {prompt_source} conflicts with input names in the run method. " 

564 "Please rename the variable or remove it from the prompt to avoid conflicts." 

565 ) 

566 if required_variables == "*" or (isinstance(required_variables, list) and var_name in required_variables): 

567 component.set_input_type(self, name=var_name, type=Any) 

568 else: 

569 component.set_input_type(self, name=var_name, type=Any, default=None) 

570 

571 def warm_up(self) -> None: 

572 """Warm up the tools, hooks, and the underlying chat generator.""" 

573 warm_up_tools(tools=self.tools) 

574 warm_up_hooks(self.hooks) 

575 if hasattr(self.chat_generator, "warm_up"): 

576 self.chat_generator.warm_up() 

577 

578 async def warm_up_async(self) -> None: 

579 """Warm up the tools, hooks, and the underlying chat generator on the serving event loop.""" 

580 warm_up_tools(tools=self.tools) 

581 await warm_up_hooks_async(self.hooks) 

582 if hasattr(self.chat_generator, "warm_up_async"): 

583 await self.chat_generator.warm_up_async() 

584 elif hasattr(self.chat_generator, "warm_up"): 

585 self.chat_generator.warm_up() 

586 

587 def close(self) -> None: 

588 """Release the hooks' and the underlying chat generator's resources.""" 

589 close_hooks(self.hooks) 

590 if hasattr(self.chat_generator, "close"): 

591 self.chat_generator.close() 

592 

593 async def close_async(self) -> None: 

594 """Release the hooks' and the underlying chat generator's async resources.""" 

595 await close_hooks_async(self.hooks) 

596 if hasattr(self.chat_generator, "close_async"): 

597 await self.chat_generator.close_async() 

598 elif hasattr(self.chat_generator, "close"): 

599 self.chat_generator.close() 

600 

601 def clone(self, **overrides: Any) -> "Agent": 

602 """ 

603 Return a new Agent configured like this one, with the given init parameters replaced. 

604 

605 :param overrides: Init parameters to replace, e.g. `agent.clone(system_prompt="...")`. 

606 :returns: The new Agent. 

607 """ 

608 init_params = inspect.signature(type(self).__init__).parameters 

609 params: dict[str, Any] = {name: getattr(self, name) for name in init_params if name != "self"} 

610 return type(self)(**{**params, **overrides}) 

611 

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

613 """ 

614 Serialize the component to a dictionary. 

615 

616 :returns: Dictionary with serialized data. 

617 """ 

618 return default_to_dict( 

619 self, 

620 chat_generator=component_to_dict(obj=self.chat_generator, name="chat_generator"), 

621 tools=serialize_tools_or_toolset(self.tools), 

622 system_prompt=self.system_prompt, 

623 user_prompt=self.user_prompt, 

624 required_variables=self.required_variables, 

625 exit_conditions=self.exit_conditions, 

626 state_schema=_schema_to_dict(self.state_schema), 

627 max_agent_steps=self.max_agent_steps, 

628 streaming_callback=serialize_callable(self.streaming_callback) if self.streaming_callback else None, 

629 raise_on_tool_invocation_failure=self.raise_on_tool_invocation_failure, 

630 tool_concurrency_limit=self.tool_concurrency_limit, 

631 tool_streaming_callback_passthrough=self.tool_streaming_callback_passthrough, 

632 hooks=_serialize_hooks_dictionary(self.hooks) if self.hooks else None, 

633 ) 

634 

635 @classmethod 

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

637 """ 

638 Deserialize the agent from a dictionary. 

639 

640 :param data: Dictionary to deserialize from. 

641 :returns: Deserialized agent. 

642 """ 

643 init_params = data.get("init_parameters", {}) 

644 

645 deserialize_component_inplace(init_params, key="chat_generator") 

646 

647 if init_params.get("state_schema") is not None: 

648 init_params["state_schema"] = _schema_from_dict(init_params["state_schema"]) 

649 

650 if init_params.get("streaming_callback") is not None: 

651 init_params["streaming_callback"] = deserialize_callable(init_params["streaming_callback"]) 

652 

653 deserialize_tools_or_toolset_inplace(init_params, key="tools") 

654 

655 if init_params.get("hooks") is not None: 

656 init_params["hooks"] = _deserialize_hooks_dictionary(init_params["hooks"]) 

657 

658 return default_from_dict(cls, data) 

659 

660 def _create_agent_span(self, tools: ToolsType) -> Any: 

661 """ 

662 Create a span for the agent run. 

663 

664 If the agent is running as part of a pipeline, this span will be nested 

665 under the current active span (the pipeline's component span). 

666 

667 :param tools: The tools selected for this run (init-time tools or the runtime override 

668 resolved by `_initialize_fresh_execution`), so the span reflects the tools actually used. 

669 """ 

670 parent_span = tracing.tracer.current_span() 

671 return tracing.tracer.trace( 

672 "haystack.agent.run", 

673 tags={ 

674 "haystack.agent.max_steps": self.max_agent_steps, 

675 "haystack.agent.tools": tools, 

676 "haystack.agent.exit_conditions": self.exit_conditions, 

677 "haystack.agent.state_schema": _schema_to_dict(self.resolved_state_schema), 

678 }, 

679 parent_span=parent_span, 

680 ) 

681 

682 def _initialize_fresh_execution( 

683 self, 

684 messages: list[ChatMessage], 

685 streaming_callback: StreamingCallbackT | None, 

686 requires_async: bool, 

687 *, 

688 generation_kwargs: dict[str, Any] | None = None, 

689 tools: ToolsType | list[str] | None = None, 

690 hook_context: dict[str, Any] | None = None, 

691 **kwargs: Any, 

692 ) -> _ExecutionContext: 

693 """ 

694 Initialize execution context for a fresh run of the agent. 

695 

696 :param messages: List of ChatMessage objects to start the agent with. 

697 :param streaming_callback: Optional callback for streaming responses. 

698 :param requires_async: Whether the agent run requires asynchronous execution. 

699 :param generation_kwargs: Additional keyword arguments for the chat generator. These are merged per key 

700 with the `generation_kwargs` passed at the chat generator's initialization: keys provided here take 

701 precedence, keys set only at initialization are kept. 

702 :param tools: Optional list of Tool objects, a Toolset, or list of tool names to use for this run. 

703 When passing tool names, tools are selected from the Agent's originally configured tools. 

704 :param hook_context: Optional dictionary of request-scoped resources made available to hooks via 

705 `state.data.get("hook_context")`. 

706 :param kwargs: Additional data to pass to the State used by the Agent. 

707 """ 

708 messages = messages or [] 

709 

710 if self._user_chat_prompt_builder is not None: 

711 user_messages = _render_prompt_messages( 

712 prompt_builder=self._user_chat_prompt_builder, 

713 expected_role=ChatRole.USER, 

714 prompt_label="user_prompt", 

715 kwargs=kwargs, 

716 ) 

717 messages = messages + user_messages 

718 

719 if self._system_chat_prompt_builder is not None: 

720 system_messages = _render_prompt_messages( 

721 prompt_builder=self._system_chat_prompt_builder, 

722 expected_role=ChatRole.SYSTEM, 

723 prompt_label="system_prompt", 

724 kwargs=kwargs, 

725 ) 

726 messages = system_messages + messages 

727 

728 if all(m.is_from(ChatRole.SYSTEM) for m in messages): 

729 logger.warning("All messages provided to the Agent component are system messages. This is not recommended.") 

730 

731 selected_tools = self._select_tools(tools=tools) 

732 flat_tools = flatten_tools_or_toolsets(tools=selected_tools) 

733 # Validate tool support once for the run (covers both init-time and runtime tools) 

734 if flat_tools and not self._chat_generator_supports_tools: 

735 raise TypeError( 

736 f"{type(self.chat_generator).__name__} does not accept tools parameter in its run method. " 

737 "The Agent component requires a chat generator that supports tools when tools are provided." 

738 ) 

739 

740 state_kwargs: dict[str, Any] = {key: kwargs[key] for key in self.resolved_state_schema.keys() if key in kwargs} 

741 state = State(schema=self.resolved_state_schema, data=state_kwargs) 

742 state.set("messages", messages) 

743 state.set("step_count", 0) 

744 state.set("token_usage", {}) 

745 state.set("context_tokens", 0) 

746 state.set("tool_call_counts", {tool.name: 0 for tool in flat_tools}) 

747 state.set("exit_reason", None) 

748 state.set("continue_run", False) 

749 state.set("tools", flat_tools) 

750 state.set("hook_context", hook_context or {}) 

751 

752 streaming_callback = select_streaming_callback( # type: ignore[call-overload] 

753 init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=requires_async 

754 ) 

755 generator_inputs: dict[str, Any] = {} 

756 if streaming_callback is not None: 

757 generator_inputs["streaming_callback"] = streaming_callback 

758 if generation_kwargs is not None: 

759 generator_inputs["generation_kwargs"] = generation_kwargs 

760 

761 tool_execution_inputs: dict[str, Any] = { 

762 "raise_on_failure": self.raise_on_tool_invocation_failure, 

763 "streaming_callback": streaming_callback, 

764 "max_workers": self.tool_concurrency_limit, 

765 "enable_streaming_callback_passthrough": self.tool_streaming_callback_passthrough, 

766 } 

767 

768 return _ExecutionContext( 

769 state=state, 

770 tools=selected_tools, 

771 chat_generator_inputs=generator_inputs, 

772 tool_execution_inputs=tool_execution_inputs, 

773 ) 

774 

775 def _select_tools(self, tools: ToolsType | list[str] | None = None) -> ToolsType: 

776 """ 

777 Select tools for the current run based on the provided tools parameter. 

778 

779 :param tools: Optional list of Tool objects, a Toolset, or list of tool names to use for this run. 

780 When passing tool names, tools are selected from the Agent's originally configured tools. 

781 :returns: Selected tools for the current run. 

782 :raises ValueError: If tool names are provided but no tools were configured at initialization, 

783 or if any provided tool name is not valid. 

784 :raises TypeError: If tools is not a list of Tool objects, a Toolset, or a list of tool names (strings). 

785 """ 

786 # Toolsets are spawned per run (see _spawn_tools / _select_tools_by_name) so concurrent runs 

787 # sharing the same configured Toolset don't corrupt each other's run-scoped state. 

788 if tools is None: 

789 return _spawn_tools(tools=self.tools) 

790 

791 if isinstance(tools, list) and all(isinstance(t, str) for t in tools): 

792 return _select_tools_by_name(self.tools, cast(list[str], tools)) 

793 

794 if isinstance(tools, (Toolset, list)): 

795 selected = cast(ToolsType, tools) # mypy can't narrow the Union type from the isinstance checks 

796 # Per-run tools are not covered by the Agent's own warm_up(), so warm them up here. 

797 # warm_up() is expected to be idempotent, so re-warming on every run is cheap. 

798 warm_up_tools(tools=selected) 

799 return _spawn_tools(tools=selected) 

800 

801 raise TypeError( 

802 "tools must be a list of Tool and/or Toolset objects, a Toolset, or a list of tool names (strings)." 

803 ) 

804 

805 def run( 

806 self, 

807 messages: list[ChatMessage], 

808 streaming_callback: StreamingCallbackT | None = None, 

809 *, 

810 generation_kwargs: dict[str, Any] | None = None, 

811 tools: ToolsType | list[str] | None = None, 

812 hook_context: dict[str, Any] | None = None, 

813 **kwargs: Any, 

814 ) -> dict[str, Any]: 

815 """ 

816 Process messages and execute tools until an exit condition is met. 

817 

818 :param messages: List of Haystack ChatMessage objects to process. 

819 :param streaming_callback: A callback that will be invoked when a response is streamed from the LLM. 

820 The same callback can be configured to emit tool results when a tool is called. 

821 :param generation_kwargs: Additional keyword arguments for the chat generator. These are merged per key 

822 with the `generation_kwargs` passed at the chat generator's initialization: keys provided here take 

823 precedence, keys set only at initialization are kept. 

824 :param tools: Optional list of Tool objects, a Toolset, or list of tool names to use for this run. 

825 When passing tool names, tools are selected from the Agent's originally configured tools. 

826 :param hook_context: Optional dictionary of request-scoped resources made available to hooks via 

827 `state.data.get("hook_context")`. Useful in web/server environments to provide per-request objects 

828 (e.g., WebSocket connections, async queues, Redis pub/sub clients) that a hook can use, for 

829 example a ConfirmationHook driving non-blocking user interaction. 

830 :param kwargs: Additional data to pass to the State schema used by the Agent. 

831 The keys must match the schema defined in the Agent's `state_schema`. 

832 :returns: 

833 A dictionary with the following keys: 

834 - "messages": List of all messages exchanged during the agent's run. 

835 - "last_message": The last message exchanged during the agent's run. 

836 - "step_count": The number of steps the agent ran. A step is one chat-generator call plus the 

837 execution of every tool call the model requested in that call (if any). The counter is incremented 

838 after each step completes, including the final step that hits an exit condition or `max_agent_steps`. 

839 - "token_usage": Aggregated token usage from every LLM call in the run, summed from each LLM message's 

840 `meta["usage"]`. 

841 - "tool_call_counts": Mapping of tool name to the number of times that tool was invoked. 

842 - "exit_reason": Why the Agent stopped, useful for routing the output downstream (e.g. with a 

843 `ConditionalRouter`). This is `"text"` when the model returned a reply with no tool calls, the name of 

844 the tool that satisfied a tool exit condition (in which case `last_message` is that tool's result), 

845 `"max_agent_steps"` when the Agent hit `max_agent_steps`, or a custom reason a hook supplied through 

846 the `stop_run` state key. 

847 - Any additional keys defined in the `state_schema`. 

848 """ 

849 agent_inputs = {"messages": messages, "streaming_callback": streaming_callback, **kwargs} 

850 self.warm_up() 

851 

852 exe_context = self._initialize_fresh_execution( 

853 messages=messages, 

854 streaming_callback=streaming_callback, 

855 requires_async=False, 

856 generation_kwargs=generation_kwargs, 

857 tools=tools, 

858 hook_context=hook_context, 

859 **kwargs, 

860 ) 

861 

862 with self._create_agent_span(tools=exe_context.tools) as span: 

863 span.set_content_tag("haystack.agent.input", agent_inputs) 

864 _run_hooks(hooks=self.hooks, hook_point=BEFORE_RUN, state=exe_context.state) 

865 # A before_run hook can restore a saved State, so resume the execution counter from its step count. 

866 exe_context.counter = exe_context.state.data.get("step_count", 0) 

867 while exe_context.counter < self.max_agent_steps: 

868 if not self._run_step(exe_context=exe_context, agent_span=span): 

869 break 

870 else: 

871 # Reached only when the loop ends without a `break`. A `break` means a step already set its own 

872 # `exit_reason`, so this branch runs only when `max_agent_steps` is why the Agent stopped. 

873 logger.warning( 

874 "Agent reached maximum agent steps of {max_agent_steps}, stopping.", 

875 max_agent_steps=self.max_agent_steps, 

876 ) 

877 exe_context.state.set("exit_reason", _EXIT_REASON_MAX_STEPS) 

878 _run_hooks(hooks=self.hooks, hook_point=AFTER_RUN, state=exe_context.state) 

879 result = _public_outputs(state=exe_context.state) 

880 if msgs := result.get("messages"): 

881 result["last_message"] = msgs[-1] 

882 span.set_content_tag("haystack.agent.output", result) 

883 span.set_tag("haystack.agent.steps_taken", exe_context.counter) 

884 

885 return result 

886 

887 async def run_async( 

888 self, 

889 messages: list[ChatMessage], 

890 streaming_callback: StreamingCallbackT | None = None, 

891 *, 

892 generation_kwargs: dict[str, Any] | None = None, 

893 tools: ToolsType | list[str] | None = None, 

894 hook_context: dict[str, Any] | None = None,