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

397 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-10-12 09:42 +0000

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import inspect 

6from collections import Counter 

7from dataclasses import dataclass 

8from typing import Any, Literal, cast 

9 

10from haystack import component, logging, tracing 

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

12 State, 

13 _schema_from_dict, 

14 _schema_to_dict, 

15 _validate_schema, 

16 replace_values, 

17) 

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

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

20from haystack.components.agents.utils import ( 

21 _record_context_tokens, 

22 _record_llm_usage, 

23 _record_tool_calls, 

24 _render_prompt_messages, 

25 _select_tools_by_name, 

26 _spawn_tools, 

27 _template_for_role, 

28 _validate_prompt_message_blocks, 

29) 

30from haystack.components.builders import ChatPromptBuilder 

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

32from haystack.core.serialization import ( 

33 component_to_dict, 

34 default_from_dict, 

35 default_to_dict, 

36 generate_qualified_class_name, 

37) 

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

39from haystack.hooks.invocation import _run_hooks, _run_hooks_async 

40from haystack.hooks.protocol import ( 

41 AFTER_RUN, 

42 AFTER_TOOL, 

43 BEFORE_LLM, 

44 BEFORE_RUN, 

45 BEFORE_TOOL, 

46 ON_EXIT, 

47 VALID_HOOK_POINTS, 

48 Hook, 

49 HookPoint, 

50) 

51from haystack.hooks.utils import ( 

52 _deserialize_hooks_dictionary, 

53 _serialize_hooks_dictionary, 

54 close_hooks, 

55 close_hooks_async, 

56 warm_up_hooks, 

57 warm_up_hooks_async, 

58) 

59from haystack.tools import ( 

60 ComponentTool, 

61 Tool, 

62 Toolset, 

63 ToolsType, 

64 _check_duplicate_tool_names, 

65 close_tools, 

66 close_tools_async, 

67 deserialize_tools_or_toolset_inplace, 

68 flatten_tools_or_toolsets, 

69 serialize_tools_or_toolset, 

70 warm_up_tools, 

71 warm_up_tools_async, 

72) 

73from haystack.utils.async_utils import _execute_component_async 

74from haystack.utils.callable_serialization import deserialize_callable, serialize_callable 

75from haystack.utils.deserialization import deserialize_component_inplace 

76 

77logger = logging.getLogger(__name__) 

78 

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

80# incomplete model generation, or the `max_agent_steps` budget running out. 

81_EXIT_REASON_TEXT = "text" 

82_EXIT_REASON_LENGTH = "length" 

83_EXIT_REASON_CONTENT_FILTER = "content_filter" 

84_EXIT_REASON_MAX_STEPS = "max_agent_steps" 

85 

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

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

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

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

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

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

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

93} 

94 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

110} 

111 

112 

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

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

115 sig = inspect.signature(instance.run) 

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

117 

118 

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

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

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

122 

123 

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

125 """ 

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

127 

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

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

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

131 """ 

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

133 if hook_point not in VALID_HOOK_POINTS: 

134 raise ValueError( 

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

136 ) 

137 for h in hook_list: 

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

139 if callable(h): 

140 raise TypeError( 

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

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

143 ) 

144 raise TypeError( 

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

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

147 ) 

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

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

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

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

152 raise ValueError( 

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

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

155 ) 

156 

157 

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

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

160 should_continue = state.data["continue_run"] 

161 state.set("continue_run", False) 

162 return should_continue 

163 

164 

165def _get_model_exit_reason(messages: list[ChatMessage]) -> str | None: 

166 """ 

167 Return the exit reason for a terminal assistant reply without tool calls. 

168 

169 Incomplete generation reasons take precedence over text so callers can distinguish a partial response from a 

170 complete answer. An empty response without a recognized terminal reason does not trigger an exit, preserving the 

171 Agent's recovery behavior for malformed tool calls that a Chat Generator discarded. 

172 """ 

173 # If the messages list is empty or the last message has tool calls, don't exit. 

174 if not messages or any(message.tool_call for message in messages): 

175 return None 

176 

177 last = messages[-1] 

178 

179 # If the last message is not from the assistant, don't exit. 

180 if not last.is_from(ChatRole.ASSISTANT): 

181 return None 

182 

183 # If the finish reason on the last message is length or content_filter, exit with that reason. 

184 if last.meta.get("finish_reason") == _EXIT_REASON_LENGTH: 

185 return _EXIT_REASON_LENGTH 

186 if last.meta.get("finish_reason") == _EXIT_REASON_CONTENT_FILTER: 

187 return _EXIT_REASON_CONTENT_FILTER 

188 

189 # If the last message has text, exit with the text reason. 

190 if last.text: 

191 return _EXIT_REASON_TEXT 

192 

193 # If we reached here no valid exit reason was found, so don't exit. 

194 return None 

195 

196 

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

198 """ 

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

200 

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

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

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

204 LLM call unless `max_agent_steps` has been reached. 

205 """ 

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

207 if not messages: 

208 return [] 

209 last_message = messages[-1] 

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

211 

212 

213def _get_tool_telemetry_data(tool: Tool) -> dict[str, Any]: 

214 """ 

215 Describe a single tool for telemetry. 

216 

217 Descriptions are deliberately left out: they are free text that can hold personal or confidential data. 

218 """ 

219 data = {"type": generate_qualified_class_name(type(tool)), "name": tool.name} 

220 if isinstance(tool, ComponentTool): 

221 data["component"] = generate_qualified_class_name(type(tool._component)) 

222 return data 

223 

224 

225def _get_tools_telemetry_data(tools: ToolsType) -> dict[str, Any]: 

226 """ 

227 Summarize the tools for telemetry. 

228 

229 Call it only after the tools are warmed up: flattening iterates the Toolsets, and some of them (e.g. 

230 SearchableToolset) warm up on iteration, while others (e.g. MCPToolset) only load their tools in `warm_up()`. 

231 """ 

232 toolsets = [tools] if isinstance(tools, Toolset) else [entry for entry in tools if isinstance(entry, Toolset)] 

233 flat_tools = flatten_tools_or_toolsets(tools=tools) 

234 return { 

235 "count": len(flat_tools), 

236 "tools": [_get_tool_telemetry_data(tool) for tool in flat_tools], 

237 "toolset_types": dict(Counter(generate_qualified_class_name(type(toolset)) for toolset in toolsets)), 

238 } 

239 

240 

241@dataclass(kw_only=True) 

242class _ExecutionContext: 

243 """ 

244 Context for executing the agent. 

245 

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

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

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

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

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

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

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

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

254 """ 

255 

256 state: State 

257 tools: ToolsType 

258 chat_generator_inputs: dict 

259 tool_execution_inputs: dict 

260 counter: int = 0 

261 

262 

263@component 

264class Agent: 

265 """ 

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

267 

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

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

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

271 

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

273 

274 ### Usage examples 

275 

276 This is an example agent that: 

277 1. Searches for tipping customs in France. 

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

279 3. Returns the final answer with its context. 

280 

281 ```python 

282 from haystack.components.agents import Agent 

283 from haystack.components.generators.chat import OpenAIResponsesChatGenerator 

284 from haystack.components.generators.utils import print_streaming_chunk 

285 from haystack.dataclasses import ChatMessage 

286 from haystack.tools import tool 

287 from typing import Annotated, Literal 

288 

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

290 @tool 

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

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

293 # Placeholder: would call actual search API 

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

295 

296 @tool 

297 def calculator( 

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

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

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

301 ) -> float: 

302 '''Perform mathematical calculations.''' 

303 if operation == "multiply": 

304 return a * b 

305 elif operation == "percentage": 

306 return (a / 100) * b 

307 return 0 

308 

309 agent = Agent( 

310 system_prompt=( 

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

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

313 ), 

314 chat_generator=OpenAIResponsesChatGenerator(), 

315 tools=[search, calculator], 

316 streaming_callback=print_streaming_chunk, 

317 ) 

318 

319 result = agent.run( 

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

321 ) 

322 

323 # Access the final response from the Agent 

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

325 ``` 

326 

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

328 

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

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

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

332 

333 ```python 

334 from haystack.components.agents import Agent 

335 from haystack.components.generators.chat import OpenAIResponsesChatGenerator 

336 from haystack.tools import tool 

337 from typing import Annotated 

338 

339 

340 @tool 

341 def translate( 

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

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

344 ) -> str: 

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

346 # Placeholder: would call an actual translation API 

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

348 

349 agent = Agent( 

350 chat_generator=OpenAIResponsesChatGenerator(), 

351 tools=[translate], 

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

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

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

355 {% endmessage %}\"\"\", 

356 ) 

357 

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

359 result = agent.run( 

360 messages=[], 

361 language="French", 

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

363 ) 

364 

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

366 ``` 

367 

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

369 

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

371 

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

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

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

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

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

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

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

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

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

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

382 

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

384 required tool has been called. 

385 

386 ```python 

387 from haystack.components.agents import Agent 

388 from haystack.components.agents.state import State 

389 from haystack.components.generators.chat import OpenAIResponsesChatGenerator 

390 from haystack.dataclasses import ChatMessage 

391 from haystack.hooks import hook 

392 from haystack.tools import tool 

393 from typing import Annotated 

394 

395 

396 @tool 

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

398 \"\"\"Save the final result.\"\"\" 

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

400 return "saved" 

401 

402 

403 @hook 

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

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

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

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

408 

409 

410 agent = Agent( 

411 chat_generator=OpenAIResponsesChatGenerator(), 

412 tools=[save_result], 

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

414 ) 

415 ``` 

416 

417 """ 

418 

419 def __init__( # noqa: PLR0913 

420 self, 

421 *, 

422 chat_generator: ChatGenerator, 

423 tools: ToolsType | None = None, 

424 system_prompt: str | None = None, 

425 user_prompt: str | None = None, 

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

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

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

429 max_agent_steps: int = 100, 

430 streaming_callback: StreamingCallbackT | None = None, 

431 raise_on_tool_invocation_failure: bool = False, 

432 tool_concurrency_limit: int = 4, 

433 tool_streaming_callback_passthrough: bool = False, 

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

435 ) -> None: 

436 """ 

437 Initialize the agent component. 

438 

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

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

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

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

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

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

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

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

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

448 :param required_variables: 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

471 list order. Valid hook points are: 

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

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

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

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

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

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

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

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

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

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

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

483 where no tools run. 

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

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

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

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

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

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

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

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

492 no effect. 

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

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

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

496 not support (via its `allowed_hook_points`). 

497 """ 

498 # --- Validation --- 

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

500 # An empty list carries no tools, so it must not trip this check: `tools` is normalized to `[]` below, and 

501 # both `clone()` and `to_dict()` feed that normalized value straight back into `__init__`. This mirrors the 

502 # equivalent check in `run()`. Only a list is measured; a Toolset is never tested for truthiness here b/c 

503 # that calls __len__, which for SearchableToolset would iterate and prematurely warm it up at init. 

504 tools_provided = tools is not None and (not isinstance(tools, list) or len(tools) > 0) 

505 if tools_provided and not self._chat_generator_supports_tools: 

506 raise TypeError( 

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

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

509 ) 

510 

511 if exit_conditions is None: 

512 exit_conditions = ["text"] 

513 

514 if state_schema is not None: 

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

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

517 if reserved_used: 

518 raise ValueError( 

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

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

521 ) 

522 _validate_schema(state_schema) 

523 _validate_prompt_message_blocks(user_prompt, system_prompt) 

524 if tool_concurrency_limit < 1: 

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

526 

527 hooks = hooks or {} 

528 _validate_hooks(hooks) 

529 

530 # --- Attributes --- 

531 self.chat_generator = chat_generator 

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

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

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

535 self.system_prompt = system_prompt 

536 self.user_prompt = user_prompt 

537 self.required_variables = required_variables 

538 self.exit_conditions = exit_conditions 

539 self.max_agent_steps = max_agent_steps 

540 self.raise_on_tool_invocation_failure = raise_on_tool_invocation_failure 

541 self.streaming_callback = streaming_callback 

542 self.tool_concurrency_limit = tool_concurrency_limit 

543 self.tool_streaming_callback_passthrough = tool_streaming_callback_passthrough 

544 self.hooks = hooks 

545 

546 # --- State schema --- 

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

548 self.state_schema = state_schema or {} 

549 self.resolved_state_schema = dict(self.state_schema) 

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

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

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

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

554 

555 # --- Component I/O --- 

556 self._run_method_params = _get_run_method_params(self) 

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

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

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

560 if param in _INTERNAL_STATE_KEYS: 

561 continue 

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

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

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

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

566 component.set_output_types(self, **output_types) 

567 

568 # --- Prompt builders --- 

569 # required_variables starts empty and is populated by _register_prompt_variables once 

570 # builder.variables are known 

571 self._user_chat_prompt_builder = ( 

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

573 if user_prompt is not None 

574 else None 

575 ) 

576 self._system_chat_prompt_builder: ChatPromptBuilder | None = None 

577 if system_prompt is not None: 

578 self._system_chat_prompt_builder = ChatPromptBuilder( 

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

580 ) 

581 self._register_prompt_variables() 

582 

583 def _register_prompt_variables(self) -> None: 

584 """ 

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

586 

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

588 run parameters, and registers component inputs. 

589 """ 

590 required_variables = self.required_variables 

591 

592 prompt_builders = [ 

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

594 ] 

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

596 logger.warning( 

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

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

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

600 ) 

601 

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

603 for builder, label in [ 

604 (self._system_chat_prompt_builder, "system_prompt"), 

605 (self._user_chat_prompt_builder, "user_prompt"), 

606 ]: 

607 if builder is not None: 

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

609 if required_variables == "*": 

610 builder.required_variables = "*" 

611 elif isinstance(self.required_variables, list): 

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

613 

614 for var_name in builder.variables: 

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

616 

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

618 prompt_source = " and ".join(sources) 

619 if var_name in self.resolved_state_schema: 

620 raise ValueError( 

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

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

623 ) 

624 if var_name in self._run_method_params: 

625 raise ValueError( 

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

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

628 ) 

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

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

631 else: 

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

633 

634 def warm_up(self) -> None: 

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

636 warm_up_tools(tools=self.tools) 

637 warm_up_hooks(self.hooks) 

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

639 self.chat_generator.warm_up() 

640 

641 async def warm_up_async(self) -> None: 

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

643 await warm_up_tools_async(tools=self.tools) 

644 await warm_up_hooks_async(self.hooks) 

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

646 await self.chat_generator.warm_up_async() 

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

648 self.chat_generator.warm_up() 

649 

650 def close(self) -> None: 

651 """Release tools, hooks, and chat generator resources.""" 

652 close_tools(tools=self.tools) 

653 close_hooks(self.hooks) 

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

655 self.chat_generator.close() 

656 

657 async def close_async(self) -> None: 

658 """Release async tools, hooks, and chat generator resources.""" 

659 await close_tools_async(tools=self.tools) 

660 await close_hooks_async(self.hooks) 

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

662 await self.chat_generator.close_async() 

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

664 self.chat_generator.close() 

665 

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

667 """ 

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

669 

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

671 :returns: The new Agent. 

672 """ 

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

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

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

676 

677 def _get_telemetry_data(self) -> dict[str, Any]: 

678 """ 

679 Data that is sent to Posthog for usage analytics. 

680 """ 

681 chat_generator_data: dict[str, Any] = {"type": generate_qualified_class_name(type(self.chat_generator))} 

682 if hasattr(self.chat_generator, "_get_telemetry_data"): 

683 chat_generator_data.update(self.chat_generator._get_telemetry_data()) 

684 return {"chat_generator": chat_generator_data, "tools": _get_tools_telemetry_data(self.tools)} 

685 

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

687 """ 

688 Serialize the component to a dictionary. 

689 

690 :returns: Dictionary with serialized data. 

691 """ 

692 return default_to_dict( 

693 self, 

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

695 tools=serialize_tools_or_toolset(self.tools), 

696 system_prompt=self.system_prompt, 

697 user_prompt=self.user_prompt, 

698 required_variables=self.required_variables, 

699 exit_conditions=self.exit_conditions, 

700 state_schema=_schema_to_dict(self.state_schema), 

701 max_agent_steps=self.max_agent_steps, 

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

703 raise_on_tool_invocation_failure=self.raise_on_tool_invocation_failure, 

704 tool_concurrency_limit=self.tool_concurrency_limit, 

705 tool_streaming_callback_passthrough=self.tool_streaming_callback_passthrough, 

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

707 ) 

708 

709 @classmethod 

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

711 """ 

712 Deserialize the agent from a dictionary. 

713 

714 :param data: Dictionary to deserialize from. 

715 :returns: Deserialized agent. 

716 """ 

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

718 

719 deserialize_component_inplace(init_params, key="chat_generator") 

720 

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

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

723 

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

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

726 

727 deserialize_tools_or_toolset_inplace(init_params, key="tools") 

728 

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

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

731 

732 return default_from_dict(cls, data) 

733 

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

735 """ 

736 Create a span for the agent run. 

737 

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

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

740 

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

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

743 """ 

744 parent_span = tracing.tracer.current_span() 

745 return tracing.tracer.trace( 

746 "haystack.agent.run", 

747 tags={ 

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

749 "haystack.agent.tools": tools, 

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

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

752 }, 

753 parent_span=parent_span, 

754 ) 

755 

756 def _initialize_fresh_execution( 

757 self, 

758 messages: list[ChatMessage], 

759 streaming_callback: StreamingCallbackT | None, 

760 requires_async: bool, 

761 *, 

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

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

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

765 **kwargs: Any, 

766 ) -> _ExecutionContext: 

767 """ 

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

769 

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

771 :param streaming_callback: Optional callback for streaming responses. 

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

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

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

775 precedence, keys set only at initialization are kept. 

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

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

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

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

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

781 """ 

782 messages = messages or [] 

783 

784 if self._user_chat_prompt_builder is not None: 

785 user_messages = _render_prompt_messages( 

786 prompt_builder=self._user_chat_prompt_builder, 

787 expected_role=ChatRole.USER, 

788 prompt_label="user_prompt", 

789 kwargs=kwargs, 

790 ) 

791 messages = messages + user_messages 

792 

793 if self._system_chat_prompt_builder is not None: 

794 system_messages = _render_prompt_messages( 

795 prompt_builder=self._system_chat_prompt_builder, 

796 expected_role=ChatRole.SYSTEM, 

797 prompt_label="system_prompt", 

798 kwargs=kwargs, 

799 ) 

800 messages = system_messages + messages 

801 

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

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

804 

805 selected_tools = self._select_tools(tools=tools) 

806 flat_tools = flatten_tools_or_toolsets(tools=selected_tools) 

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

808 if flat_tools and not self._chat_generator_supports_tools: 

809 raise TypeError( 

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

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

812 ) 

813 

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

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

816 state.set("messages", messages) 

817 state.set("step_count", 0) 

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

819 state.set("context_tokens", 0) 

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

821 state.set("exit_reason", None) 

822 state.set("continue_run", False) 

823 state.set("tools", flat_tools) 

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

825 

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

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

828 ) 

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

830 if streaming_callback is not None: 

831 generator_inputs["streaming_callback"] = streaming_callback 

832 if generation_kwargs is not None: 

833 generator_inputs["generation_kwargs"] = generation_kwargs 

834 

835 tool_execution_inputs: dict[str, Any] = { 

836 "raise_on_failure": self.raise_on_tool_invocation_failure, 

837 "streaming_callback": streaming_callback, 

838 "max_workers": self.tool_concurrency_limit, 

839 "enable_streaming_callback_passthrough": self.tool_streaming_callback_passthrough, 

840 } 

841 

842 return _ExecutionContext( 

843 state=state, 

844 tools=selected_tools, 

845 chat_generator_inputs=generator_inputs, 

846 tool_execution_inputs=tool_execution_inputs, 

847 ) 

848 

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

850 """ 

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

852 

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

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

855 :returns: Selected tools for the current run. 

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

857 or if any provided tool name is not valid. 

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

859 """ 

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

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

862 if tools is None: 

863 return _spawn_tools(tools=self.tools) 

864 

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

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

867 

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

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

870 return _spawn_tools(tools=selected) 

871 

872 raise TypeError( 

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

874 ) 

875 

876 def run( 

877 self, 

878 messages: list[ChatMessage], 

879 streaming_callback: StreamingCallbackT | None = None, 

880 *, 

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

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

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

884 **kwargs: Any, 

885 ) -> dict[str, Any]: 

886 """ 

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

888 

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

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

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

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

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

894 precedence, keys set only at initialization are kept. 

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

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

897 Tool and Toolset objects passed here are warmed up automatically; the caller is responsible for 

898 closing them if they hold resources. 

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

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

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

902 example a ConfirmationHook driving non-blocking user interaction. 

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

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

905 :returns: 

906 A dictionary with the following keys: 

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

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

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

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

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

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

913 `meta["usage"]`. 

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

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

916 `ConditionalRouter`). One of: `"text"` (the model returned a complete reply with no tool calls), 

917 `"length"` or `"content_filter"` (the model returned an incomplete reply, which may contain partial 

918 text), the name of the tool that satisfied a tool exit condition (its result is the last tool message in 

919 `messages` whose `tool_call_result.origin.tool_name` matches it, which may not be `last_message` when 

920 the model called several tools at once), or `"max_agent_steps"` (the Agent hit `max_agent_steps` before 

921 meeting an exit condition), or a custom reason a hook supplied through the `stop_run` state key. 

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

923 """ 

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

925 self.warm_up() 

926 # warm up tools passed at runtime 

927 if isinstance(tools, Toolset): 

928 tools_to_warm_up: ToolsType = tools 

929 else: 

930 tools_to_warm_up = [tool for tool in tools or [] if not isinstance(tool, str)] 

931 warm_up_tools(tools=tools_to_warm_up) 

932 

933 exe_context = self._initialize_fresh_execution( 

934 messages=messages, 

935 streaming_callback=streaming_callback, 

936 requires_async=False, 

937 generation_kwargs=generation_kwargs, 

938 tools=tools, 

939 hook_context=hook_context, 

940 **kwargs, 

941 ) 

942 

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

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

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

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

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

948 while exe_context.counter < self.max_agent_steps: 

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

950 break 

951 else: 

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

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

954 logger.warning( 

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

956 max_agent_steps=self.max_agent_steps, 

957 ) 

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

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

960 result = _public_outputs(state=exe_context.state) 

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

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

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

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

965 

966 return result 

967 

968 async def run_async( 

969 self, 

970 messages: list[ChatMessage], 

971 streaming_callback: StreamingCallbackT | None = None, 

972 *, 

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

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

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

976 **kwargs: Any, 

977 ) -> dict[str, Any]: 

978 """ 

979 Asynchronously process messages and execute tools until the exit condition is met. 

980 

981 This is the asynchronous version of the `run` method. It follows the same logic but uses 

982 asynchronous operations where possible, such as calling the `run_async` method of the ChatGenerator 

983 if available. 

984 

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

986 :param streaming_callback: An asynchronous callback that will be invoked when a response is streamed from the 

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

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

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

990 precedence, keys set only at initialization are kept. 

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

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

993 Tool and Toolset objects passed here are warmed up automatically; the caller is responsible for 

994 closing them if they hold resources. 

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

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

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

998 example a ConfirmationHook driving non-blocking user interaction. 

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

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

1001 :returns: 

1002 A dictionary with the following keys: 

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

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

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

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

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

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

1009 `meta["usage"]`. 

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

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

1012 `ConditionalRouter`). One of: `"text"` (the model returned a complete reply with no tool calls), 

1013 `"length"` or `"content_filter"` (the model returned an incomplete reply, which may contain partial 

1014 text), the name of the tool that satisfied a tool exit condition (its result is the last tool message in 

1015 `messages` whose `tool_call_result.origin.tool_name` matches it, which may not be `last_message` when 

1016 the model called several tools at once), or `"max_agent_steps"` (the Agent hit `max_agent_steps` before 

1017 meeting an exit condition), or a custom reason a hook supplied through the `stop_run` state key. 

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

1019 """ 

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

1021 await self.warm_up_async() 

1022 # warm up tools passed at runtime 

1023 if isinstance(tools, Toolset): 

1024 tools_to_warm_up: ToolsType = tools 

1025 else: 

1026 tools_to_warm_up = [tool for tool in tools or [] if not isinstance(tool, str)] 

1027 await warm_up_tools_async(tools=tools_to_warm_up) 

1028 

1029 exe_context = self._initialize_fresh_execution( 

1030 messages=messages, 

1031 streaming_callback=streaming_callback, 

1032 requires_async=True, 

1033 tools=tools, 

1034 generation_kwargs=generation_kwargs, 

1035 hook_context=hook_context, 

1036 **kwargs, 

1037 ) 

1038 

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

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

1041 await _run_hooks_async(hooks=self.hooks, hook_point=BEFORE_RUN, state=exe_context.state) 

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

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

1044 while exe_context.counter < self.max_agent_steps: 

1045 if not await self._run_step_async(exe_context=exe_context, agent_span=span): 

1046 break 

1047 else: 

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

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

1050 logger.warning( 

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

1052 max_agent_steps=self.max_agent_steps, 

1053 ) 

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

1055 await _run_hooks_async(hooks=self.hooks, hook_point=AFTER_RUN, state=exe_context.state) 

1056 result = _public_outputs(state=exe_context.state) 

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

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

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

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

1061 

1062 return result 

1063 

1064 def _run_step(self, exe_context: _ExecutionContext, agent_span: tracing.Span) -> bool: 

1065 """Execute one agent step. Returns True to continue the loop, False to stop.""" 

1066 with tracing.tracer.trace( 

1067 "haystack.agent.step", tags={"haystack.agent.step": exe_context.counter}, parent_span=agent_span 

1068 ) as step_span: 

1069 # Re-flatten the tools every step so dynamic toolsets (e.g. SearchableToolset) surface tools discovered in 

1070 # earlier steps. Validate names here so duplicates fail before starting the step. 

1071 current_tools = flatten_tools_or_toolsets(tools=exe_context.tools) 

1072 _check_duplicate_tool_names(tools=current_tools) 

1073 # Expose the current tools to hooks (e.g. ConfirmationHook) via State. 

1074 exe_context.state.set("tools", current_tools, handler_override=replace_values) 

1075 

1076 _run_hooks(hooks=self.hooks, hook_point=BEFORE_LLM, state=exe_context.state) 

1077 # A hook requested a stop: end the run at the step boundary, before spending another LLM call. 

1078 if (reason := exe_context.state.data.get("stop_run")) is not None: 

1079 exe_context.state.set("exit_reason", reason) 

1080 return False 

1081 chat_generator_inputs = { 

1082 "messages": exe_context.state.data["messages"], 

1083 **exe_context.chat_generator_inputs, 

1084 } 

1085 if current_tools: 

1086 chat_generator_inputs["tools"] = current_tools 

1087 with tracing.tracer.trace("haystack.agent.step.llm", parent_span=step_span) as llm_span: 

1088 llm_span.set_content_tag("haystack.agent.step.llm.input", chat_generator_inputs) 

1089 result = self.chat_generator.run(**chat_generator_inputs) 

1090 llm_span.set_content_tag("haystack.agent.step.llm.output", result) 

1091 llm_messages = result["replies"] 

1092 exe_context.state.set("messages", llm_messages) 

1093 _record_llm_usage(state=exe_context.state, llm_messages=llm_messages) 

1094 _record_context_tokens(state=exe_context.state, llm_messages=llm_messages) 

1095 

1096 # Stop when there are no tools, or the model produced a terminal reply without tool calls. 

1097 model_exit_reason = _get_model_exit_reason(messages=llm_messages) 

1098 if not current_tools or model_exit_reason is not None: 

1099 exe_context.counter += 1 

1100 exe_context.state.set("step_count", exe_context.counter) 

1101 exe_context.state.set("exit_reason", model_exit_reason or _EXIT_REASON_TEXT) 

1102 return self._continue_after_exit_hooks(exe_context=exe_context) 

1103 

1104 _run_hooks(hooks=self.hooks, hook_point=BEFORE_TOOL, state=exe_context.state) 

1105 # Re-read the pending tool calls from State so that any rewrites a before_tool hook made (e.g. 

1106 # ConfirmationHook rejecting or modifying calls) are honored by the executor. 

1107 pending_tool_call_messages = _pending_tool_call_messages_from_state(state=exe_context.state) 

1108 

1109 tool_execution_inputs = { 

1110 "messages": pending_tool_call_messages, 

1111 "state": exe_context.state, 

1112 **exe_context.tool_execution_inputs, 

1113 "tools": current_tools, 

1114 } 

1115 tool_messages, exe_context.state = _run_tool(**tool_execution_inputs) 

1116 exe_context.state.set("messages", tool_messages) 

1117 _record_tool_calls(state=exe_context.state, tool_messages=tool_messages) 

1118 _run_hooks(hooks=self.hooks, hook_point=AFTER_TOOL, state=exe_context.state) 

1119 

1120 exe_context.counter += 1 

1121 exe_context.state.set("step_count", exe_context.counter) 

1122 exit_condition_tool = ( 

1123 None 

1124 if self.exit_conditions == ["text"] 

1125 else self._check_exit_conditions(llm_messages=pending_tool_call_messages, tool_messages=tool_messages) 

1126 ) 

1127 if exit_condition_tool is not None: 

1128 exe_context.state.set("exit_reason", exit_condition_tool) 

1129 return self._continue_after_exit_hooks(exe_context=exe_context) 

1130 return True 

1131 

1132 async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: tracing.Span) -> bool: 

1133 """Execute one agent step asynchronously. Returns True to continue the loop, False to stop.""" 

1134 with tracing.tracer.trace( 

1135 "haystack.agent.step", tags={"haystack.agent.step": exe_context.counter}, parent_span=agent_span 

1136 ) as step_span: 

1137 # Re-flatten the tools every step so dynamic toolsets (e.g. SearchableToolset) surface tools discovered in 

1138 # earlier steps. Validate names here so duplicates fail before starting the step. 

1139 current_tools = flatten_tools_or_toolsets(tools=exe_context.tools) 

1140 _check_duplicate_tool_names(tools=current_tools) 

1141 # Expose the current tools to hooks (e.g. ConfirmationHook) via State. 

1142 exe_context.state.set("tools", current_tools, handler_override=replace_values) 

1143 

1144 await _run_hooks_async(hooks=self.hooks, hook_point=BEFORE_LLM, state=exe_context.state) 

1145 # A hook requested a stop: end the run at the step boundary, before spending another LLM call. 

1146 if (reason := exe_context.state.data.get("stop_run")) is not None: 

1147 exe_context.state.set("exit_reason", reason) 

1148 return False 

1149 chat_generator_inputs = { 

1150 "messages": exe_context.state.data["messages"], 

1151 **exe_context.chat_generator_inputs, 

1152 } 

1153 if current_tools: 

1154 chat_generator_inputs["tools"] = current_tools 

1155 with tracing.tracer.trace("haystack.agent.step.llm", parent_span=step_span) as llm_span: 

1156 llm_span.set_content_tag("haystack.agent.step.llm.input", chat_generator_inputs) 

1157 # For sync-only generators, _execute_component_async dispatches to a thread via asyncio.to_thread, 

1158 # which copies the current contextvars context — preserving the active tracing span. 

1159 result = await _execute_component_async(self.chat_generator, **chat_generator_inputs) 

1160 llm_span.set_content_tag("haystack.agent.step.llm.output", result) 

1161 llm_messages = result["replies"] 

1162 exe_context.state.set("messages", llm_messages) 

1163 _record_llm_usage(state=exe_context.state, llm_messages=llm_messages) 

1164 _record_context_tokens(state=exe_context.state, llm_messages=llm_messages) 

1165 

1166 # Stop when there are no tools, or the model produced a terminal reply without tool calls. 

1167 model_exit_reason = _get_model_exit_reason(messages=llm_messages) 

1168 if not current_tools or model_exit_reason is not None: 

1169 exe_context.counter += 1 

1170 exe_context.state.set("step_count", exe_context.counter) 

1171 exe_context.state.set("exit_reason", model_exit_reason or _EXIT_REASON_TEXT) 

1172 return await self._continue_after_exit_hooks_async(exe_context=exe_context) 

1173 

1174 await _run_hooks_async(hooks=self.hooks, hook_point=BEFORE_TOOL, state=exe_context.state) 

1175 # Re-read the pending tool calls from State so that any rewrites a before_tool hook made (e.g. 

1176 # ConfirmationHook rejecting or modifying calls) are honored by the executor. 

1177 pending_tool_call_messages = _pending_tool_call_messages_from_state(state=exe_context.state) 

1178 

1179 tool_execution_inputs = { 

1180 "messages": pending_tool_call_messages, 

1181 "state": exe_context.state, 

1182 **exe_context.tool_execution_inputs, 

1183 "tools": current_tools, 

1184 } 

1185 tool_messages, exe_context.state = await _run_tool_async(**tool_execution_inputs) 

1186 exe_context.state.set("messages", tool_messages) 

1187 _record_tool_calls(state=exe_context.state, tool_messages=tool_messages) 

1188 await _run_hooks_async(hooks=self.hooks, hook_point=AFTER_TOOL, state=exe_context.state) 

1189 

1190 exe_context.counter += 1 

1191 exe_context.state.set("step_count", exe_context.counter) 

1192 exit_condition_tool = ( 

1193 None 

1194 if self.exit_conditions == ["text"] 

1195 else self._check_exit_conditions(llm_messages=pending_tool_call_messages, tool_messages=tool_messages) 

1196 ) 

1197 if exit_condition_tool is not None: 

1198 exe_context.state.set("exit_reason", exit_condition_tool) 

1199 return await self._continue_after_exit_hooks_async(exe_context=exe_context) 

1200 return True 

1201 

1202 def _check_exit_conditions(self, llm_messages: list[ChatMessage], tool_messages: list[ChatMessage]) -> str | None: 

1203 """ 

1204 Decide whether the agent should stop looping and, if so, on which tool. 

1205 

1206 Returns the name of the tool that triggered an exit condition: the model called at least one tool listed in 

1207 `exit_conditions` and that tool did not error. Every tool call in the message is checked, so the order of 

1208 parallel tool calls does not matter. When several exit-condition tools are called in the same step, the first 

1209 one encountered is returned. 

1210 

1211 :param llm_messages: List of messages from the LLM 

1212 :param tool_messages: List of messages from tool execution. 

1213 :return: The name of the tool that satisfied an exit condition, or None if none did (or one errored). 

1214 """ 

1215 errored_tools = { 

1216 tool_msg.tool_call_result.origin.tool_name 

1217 for tool_msg in tool_messages 

1218 if tool_msg.tool_call_result is not None and tool_msg.tool_call_result.error 

1219 } 

1220 

1221 first_match: str | None = None 

1222 for msg in llm_messages: 

1223 for tool_call in msg.tool_calls: 

1224 if tool_call.tool_name not in self.exit_conditions: 

1225 continue 

1226 # An errored exit-condition tool cancels the exit, even if another exit-condition tool succeeded. 

1227 if tool_call.tool_name in errored_tools: 

1228 return None 

1229 first_match = first_match or tool_call.tool_name 

1230 return first_match 

1231 

1232 def _continue_after_exit_hooks(self, exe_context: _ExecutionContext) -> bool: 

1233 """ 

1234 Run `on_exit` hooks and return whether the loop should keep going. 

1235 

1236 A hook keeps the Agent running by setting the `continue_run` control flag (`state.set("continue_run", True)`), 

1237 usually alongside a message telling the model what to do next. The flag is consumed on each exit attempt and 

1238 the loop stays bounded by `max_agent_steps`. 

1239 """ 

1240 if not self.hooks.get(ON_EXIT): 

1241 return False 

1242 exe_context.state.set("continue_run", False) 

1243 _run_hooks(hooks=self.hooks, hook_point=ON_EXIT, state=exe_context.state) 

1244 return _consume_continue_run(exe_context.state) 

1245 

1246 async def _continue_after_exit_hooks_async(self, exe_context: _ExecutionContext) -> bool: 

1247 """Async version of `_continue_after_exit_hooks`.""" 

1248 if not self.hooks.get(ON_EXIT): 

1249 return False 

1250 exe_context.state.set("continue_run", False) 

1251 await _run_hooks_async(hooks=self.hooks, hook_point=ON_EXIT, state=exe_context.state) 

1252 return _consume_continue_run(exe_context.state)