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
5import inspect
6from collections import Counter
7from dataclasses import dataclass
8from typing import Any, Literal, cast
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
77logger = logging.getLogger(__name__)
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"
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}
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}
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}
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}
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.
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 )
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
165def _get_model_exit_reason(messages: list[ChatMessage]) -> str | None:
166 """
167 Return the exit reason for a terminal assistant reply without tool calls.
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
177 last = messages[-1]
179 # If the last message is not from the assistant, don't exit.
180 if not last.is_from(ChatRole.ASSISTANT):
181 return None
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
189 # If the last message has text, exit with the text reason.
190 if last.text:
191 return _EXIT_REASON_TEXT
193 # If we reached here no valid exit reason was found, so don't exit.
194 return None
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.
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 []
213def _get_tool_telemetry_data(tool: Tool) -> dict[str, Any]:
214 """
215 Describe a single tool for telemetry.
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
225def _get_tools_telemetry_data(tools: ToolsType) -> dict[str, Any]:
226 """
227 Summarize the tools for telemetry.
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 }
241@dataclass(kw_only=True)
242class _ExecutionContext:
243 """
244 Context for executing the agent.
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 """
256 state: State
257 tools: ToolsType
258 chat_generator_inputs: dict
259 tool_execution_inputs: dict
260 counter: int = 0
263@component
264class Agent:
265 """
266 A tool-using Agent powered by a large language model.
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.
272 Without tools, the Agent works like a standard LLM that generates text. It produces one response and then stops.
274 ### Usage examples
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.
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
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."
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
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 )
319 result = agent.run(
320 messages=[ChatMessage.from_user("Calculate the appropriate tip for an €85 meal in France")]
321 )
323 # Access the final response from the Agent
324 # print(result["last_message"].text)
325 ```
327 #### Using a `user_prompt` template with variables
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.
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
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}]"
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 )
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 )
365 print(result["last_message"].text)
366 ```
368 #### Using hooks to influence the run loop
370 Hooks are callables that receive the live `State` and run at specific points in the Agent loop:
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)`.
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.
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
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"
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
410 agent = Agent(
411 chat_generator=OpenAIResponsesChatGenerator(),
412 tools=[save_result],
413 hooks={"on_exit": [require_save]},
414 )
415 ```
417 """
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.
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 )
511 if exit_conditions is None:
512 exit_conditions = ["text"]
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.")
527 hooks = hooks or {}
528 _validate_hooks(hooks)
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
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)
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)
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()
583 def _register_prompt_variables(self) -> None:
584 """
585 Collect variables from both Chat Prompt Builders and register Agent inputs.
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
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 )
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]
614 for var_name in builder.variables:
615 all_variables.setdefault(var_name, []).append(label)
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)
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()
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()
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()
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()
666 def clone(self, **overrides: Any) -> "Agent":
667 """
668 Return a new Agent configured like this one, with the given init parameters replaced.
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})
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)}
686 def to_dict(self) -> dict[str, Any]:
687 """
688 Serialize the component to a dictionary.
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 )
709 @classmethod
710 def from_dict(cls, data: dict[str, Any]) -> "Agent":
711 """
712 Deserialize the agent from a dictionary.
714 :param data: Dictionary to deserialize from.
715 :returns: Deserialized agent.
716 """
717 init_params = data.get("init_parameters", {})
719 deserialize_component_inplace(init_params, key="chat_generator")
721 if init_params.get("state_schema") is not None:
722 init_params["state_schema"] = _schema_from_dict(init_params["state_schema"])
724 if init_params.get("streaming_callback") is not None:
725 init_params["streaming_callback"] = deserialize_callable(init_params["streaming_callback"])
727 deserialize_tools_or_toolset_inplace(init_params, key="tools")
729 if init_params.get("hooks") is not None:
730 init_params["hooks"] = _deserialize_hooks_dictionary(init_params["hooks"])
732 return default_from_dict(cls, data)
734 def _create_agent_span(self, tools: ToolsType) -> Any:
735 """
736 Create a span for the agent run.
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).
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 )
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.
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 []
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
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
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.")
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 )
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 {})
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
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 }
842 return _ExecutionContext(
843 state=state,
844 tools=selected_tools,
845 chat_generator_inputs=generator_inputs,
846 tool_execution_inputs=tool_execution_inputs,
847 )
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.
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)
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))
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)
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 )
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.
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)
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 )
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)
966 return result
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.
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.
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)
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 )
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)
1062 return result
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)
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)
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)
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)
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)
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
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)
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)
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)
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)
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)
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
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.
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.
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 }
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
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.
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)
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)