Coverage for haystack/components/agents/utils.py: 98%
120 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import re
6from copy import deepcopy
7from typing import Any
9from haystack.components.agents.state.state import State
10from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder
11from haystack.dataclasses import ChatMessage, ChatRole
12from haystack.tools import Tool, Toolset, ToolsType
14# Input/output token key conventions across chat generators: most report OpenAI-style
15# `prompt_tokens`/`completion_tokens`; OpenAIResponsesChatGenerator reports `input_tokens`/`output_tokens`.
16_INPUT_TOKEN_KEYS = ("prompt_tokens", "input_tokens")
17_OUTPUT_TOKEN_KEYS = ("completion_tokens", "output_tokens")
20# ---------------------------
21# Run metadata helpers
22# ---------------------------
25def _accumulate_usage(current: Any, new: Any) -> Any:
26 """
27 Recursively sum numeric leaf values across two usage-like dicts.
29 Used to aggregate `ChatMessage.meta["usage"]` payloads across LLM calls in a run. Nested dicts (e.g. OpenAI's
30 `completion_tokens_details`) are merged recursively; numeric leaves are summed; other types fall back to the new
31 value.
33 :param current: The current accumulated usage data.
34 :param new: The new usage data to merge in.
35 """
36 if isinstance(current, dict) and isinstance(new, dict):
37 result = dict(current)
38 for k, v in new.items():
39 result[k] = _accumulate_usage(result[k], v) if k in result else deepcopy(v)
40 return result
41 if isinstance(current, (int, float)) and isinstance(new, (int, float)):
42 return current + new
43 return new
46def _record_llm_usage(state: State, llm_messages: list[ChatMessage]) -> None:
47 """
48 Aggregate token usage from the latest LLM messages into the State.
50 Only writes when at least one message reports `meta["usage"]`, so generators that don't surface usage data
51 leave `token_usage` at its default empty dict rather than overwriting it.
53 :param state: The Agent's State, used to read the running `token_usage` total and write back the new total.
54 :param llm_messages: The ChatMessage objects returned from the latest LLM call. Token usage is read from each
55 message's `meta["usage"]` field, if present.
56 """
57 current = state.data.get("token_usage")
58 updated = False
59 for msg in llm_messages:
60 usage = msg.meta.get("usage")
61 if isinstance(usage, dict):
62 current = _accumulate_usage(current or {}, usage)
63 updated = True
64 if updated:
65 state.set("token_usage", current)
68def _record_tool_calls(state: State, tool_messages: list[ChatMessage]) -> None:
69 """
70 Increment per-tool call counts in the State for every successfully dispatched tool.
72 :param state: The Agent's State, used to read the running `tool_call_counts` map and write back the new totals.
73 :param tool_messages: The ChatMessage objects returned from the latest tool execution. Per-tool counts are
74 incremented based on each message's `tool_call_result.origin.tool_name`.
75 """
76 counts = state.data.get("tool_call_counts") or {}
77 updated = False
78 for tm in tool_messages:
79 if tm.tool_call_result is None:
80 continue
81 name = tm.tool_call_result.origin.tool_name
82 counts[name] = counts.get(name, 0) + 1
83 updated = True
84 if updated:
85 state.set("tool_call_counts", counts)
88# ---------------------------
89# Tool helpers
90# ---------------------------
93def _spawn_selection_copy(item: Tool | Toolset, selected_tool_names: set[str]) -> Toolset | None:
94 """
95 Return the per-run copy carrying the selection, or None if the item does not provide one.
97 A Toolset with run-scoped state (e.g. SearchableToolset) overrides `spawn()` to return a copy that
98 applies `selected_tool_names` itself. A plain Toolset returns itself from `spawn()` (it has nothing
99 to isolate), and a standalone Tool has no `spawn()`: in both cases the caller applies the selection.
101 :param item: A configured Tool or Toolset.
102 :param selected_tool_names: The tool names selected for this run.
103 :returns: The selection-carrying per-run copy, or None.
104 """
105 if not isinstance(item, Toolset):
106 return None
107 spawned = item.spawn(selected_tool_names=selected_tool_names)
108 return spawned if spawned is not item else None
111def _select_tools_by_name(configured_tools: ToolsType, names: list[str]) -> list[Tool | Toolset]:
112 """
113 Select configured tools by name for a single run.
115 Standalone Tools are kept when their name is requested. A Toolset with run-scoped state (one overriding
116 `spawn()`, such as SearchableToolset) is replaced by a per-run copy carrying the requested names, so its
117 dynamic behavior (search/lazy-loading) is preserved without mutating the shared, configured Toolset. Any
118 other Toolset is warmed up and reduced to the matching Tools.
120 :param configured_tools: The tools configured on the Agent.
121 :param names: The requested tool names.
122 :returns: The selected Tools and/or selection-scoped Toolset copies.
123 :raises ValueError: If no tools were configured, or if any requested name is not a valid tool name.
124 """
125 if configured_tools is None:
126 raise ValueError("No tools were configured for the Agent at initialization.")
128 requested_names = set(names)
129 items: list[Tool | Toolset] = (
130 [configured_tools] if isinstance(configured_tools, Toolset) else list(configured_tools)
131 )
133 # Resolve the tools each item offers for selection
134 selectable_per_item: list[tuple[Tool | Toolset, list[Tool]]] = []
135 for item in items:
136 selectable = item.get_selectable_tools() if isinstance(item, Toolset) else [item]
137 selectable_per_item.append((item, selectable))
139 valid_tool_names = {tool.name for _, selectable in selectable_per_item for tool in selectable}
140 # A dynamic Toolset may look empty before its catalog is resolved, so emptiness is checked here.
141 if not valid_tool_names:
142 raise ValueError("No tools were configured for the Agent at initialization.")
144 invalid_tool_names = requested_names - valid_tool_names
145 if invalid_tool_names:
146 raise ValueError(
147 f"The following tool names are not valid: {invalid_tool_names}. Valid tool names are: {valid_tool_names}."
148 )
150 selected: list[Tool | Toolset] = []
151 for item, selectable in selectable_per_item:
152 matched = requested_names & {tool.name for tool in selectable}
153 if not matched:
154 continue
155 run_copy = _spawn_selection_copy(item, matched)
156 if run_copy is not None:
157 selected.append(run_copy)
158 else:
159 # Select from `selectable`, the list the names were validated against: iterating a dynamic
160 # Toolset could silently miss tools.
161 selected.extend(tool for tool in selectable if tool.name in matched)
162 return selected
165def _spawn_tools(tools: ToolsType) -> ToolsType:
166 """
167 Return per-run copies of `tools`, replacing each Toolset with its `spawn()` (Tools are passed through).
169 This isolates run-scoped Toolset state (e.g. a SearchableToolset's discovered tools and any active name
170 selection) so that concurrent runs sharing the same configured Toolset — such as parallel sub-agent tool calls
171 or concurrent requests against one Agent — don't corrupt each other. A plain Toolset has no run-scoped state
172 and its `spawn()` returns itself unchanged.
173 """
174 if isinstance(tools, Toolset):
175 return tools.spawn()
176 return [item.spawn() if isinstance(item, Toolset) else item for item in tools]
179# ---------------------------
180# Context token helpers
181# ---------------------------
184def _first_numeric(usage: dict[str, Any], keys: tuple[str, ...]) -> int:
185 """
186 Return the first numeric value found under `keys` in `usage`, or 0 if none is present.
188 :param usage: A ChatMessage `meta["usage"]` payload.
189 :param keys: Candidate keys to check, in priority order.
190 :returns: The first `int`/`float` value (as an `int`), or 0. bool values are skipped (not token counts).
191 """
192 for key in keys:
193 value = usage.get(key)
194 # bool is an int subclass, so exclude it explicitly: True/False is not a token count.
195 if isinstance(value, bool):
196 continue
197 if isinstance(value, (int, float)):
198 return int(value)
199 return 0
202def _context_tokens_from_usage(usage: dict[str, Any]) -> int:
203 """
204 Sum the input and output tokens reported in a single `meta["usage"]` dict.
206 :param usage: A ChatMessage `meta["usage"]` payload.
207 :returns: Input plus output tokens, or 0 if neither key convention is present.
208 """
209 return _first_numeric(usage, _INPUT_TOKEN_KEYS) + _first_numeric(usage, _OUTPUT_TOKEN_KEYS)
212def _record_context_tokens(state: State, llm_messages: list[ChatMessage]) -> None:
213 """
214 Store the approximate current context-window token count from the latest LLM call.
216 A chat-generator call returns a single reply, so only the last message is inspected. Unlike
217 `token_usage`, which accumulates across the run, this value is replaced each call with that reply's
218 prompt-plus-completion tokens. Only writes when usage is reported, so generators that don't surface
219 usage leave the previous value untouched.
221 :param state: The Agent's State, used to write the latest `context_tokens` count.
222 :param llm_messages: The ChatMessage objects returned from the latest LLM call.
223 """
224 if not llm_messages:
225 return
226 usage = llm_messages[-1].meta.get("usage")
227 if isinstance(usage, dict):
228 tokens = _context_tokens_from_usage(usage)
229 if tokens:
230 state.set("context_tokens", tokens)
233# ---------------------------
234# Prompt helpers
235# ---------------------------
237# Regex to detect the Jinja2 chat template syntax
238_JINJA2_CHAT_TEMPLATE_RE = re.compile(r"\{%\s*message\s")
239# Regex to extract the role from a Jinja2 message block, e.g. {% message role="user" %}
240_JINJA2_MESSAGE_ROLE_RE = re.compile(r'\{%\s*message\s+role\s*=\s*["\'](\w+)["\']')
243def _validate_prompt_message_blocks(user_prompt: str | None, system_prompt: str | None) -> None:
244 """
245 Validate explicit Jinja2 message blocks in Agent prompts.
247 :param user_prompt: Optional user prompt template.
248 :param system_prompt: Optional system prompt template.
249 :raises ValueError: If a prompt contains multiple message blocks or a literal block role is invalid.
250 """
251 if user_prompt is not None:
252 message_blocks = _JINJA2_CHAT_TEMPLATE_RE.findall(user_prompt)
253 roles = _JINJA2_MESSAGE_ROLE_RE.findall(user_prompt)
254 if len(message_blocks) > 1:
255 raise ValueError(f"user_prompt must define exactly one message block, found {len(message_blocks)}.")
256 if roles and roles[0] != "user":
257 raise ValueError(f"user_prompt message block must have role 'user', found role '{roles[0]}'.")
259 if system_prompt is not None and _JINJA2_CHAT_TEMPLATE_RE.search(system_prompt):
260 message_blocks = _JINJA2_CHAT_TEMPLATE_RE.findall(system_prompt)
261 roles = _JINJA2_MESSAGE_ROLE_RE.findall(system_prompt)
262 if len(message_blocks) > 1:
263 raise ValueError(f"system_prompt must define exactly one message block, found {len(message_blocks)}.")
264 if roles and roles[0] != "system":
265 raise ValueError(f"system_prompt message block must have role 'system', found role '{roles[0]}'.")
268def _template_for_role(prompt: str, role: str) -> str:
269 """
270 Convert a prompt into a ChatPromptBuilder string template for the expected role.
272 :param prompt: Prompt template, with or without an explicit Jinja2 message block.
273 :param role: Role to use when wrapping a plain string prompt.
274 :returns: The original message-block template, or a plain string prompt wrapped in one message block.
275 """
276 if _JINJA2_CHAT_TEMPLATE_RE.search(prompt):
277 return prompt
278 return f'{{% message role="{role}" %}}{prompt}{{% endmessage %}}'
281def _render_prompt_messages(
282 *, prompt_builder: ChatPromptBuilder, expected_role: ChatRole, prompt_label: str, kwargs: dict[str, Any]
283) -> list[ChatMessage]:
284 """
285 Render one Agent prompt and validate the rendered message.
287 :param prompt_builder: Builder configured with the prompt template.
288 :param expected_role: Role the rendered message must have.
289 :param prompt_label: Prompt name used in error messages.
290 :param kwargs: Runtime values available to the prompt template.
291 :returns: A single rendered prompt message.
292 :raises ValueError: If the prompt renders to zero, multiple, or wrong-role messages.
293 """
294 prompt_kwargs = {var: kwargs[var] for var in prompt_builder.variables if var in kwargs}
295 prompt_messages = prompt_builder.run(**prompt_kwargs)["prompt"]
296 if len(prompt_messages) != 1:
297 raise ValueError(
298 f"{prompt_label} must render to exactly one {expected_role.value} message. "
299 f"Got {len(prompt_messages)} messages."
300 )
301 if not prompt_messages[0].is_from(expected_role):
302 raise ValueError(
303 f"{prompt_label} must render to a {expected_role.value} message. "
304 f"Got a message with role {prompt_messages[0].role}."
305 )
306 return prompt_messages