Skip to content

chore: Moving some agent utils to utils file #12238

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sjrl merged 4 commits into from
Aug 5, 2026
+452 −376
Merged

chore: Moving some agent utils to utils file #12238

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
38d9286
refactoring and putting more functions into utils
sjrl Aug 4, 2026
4646937
properly split the prompt tests
sjrl Aug 4, 2026
5277579
move test class back to where it was before
sjrl Aug 4, 2026
697fa4c
PR comments
sjrl Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
212 changes: 10 additions & 202 deletions haystack/components/agents/agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,6 @@
# SPDX-License-Identifier: Apache-2.0

import inspect
import re
from copy import deepcopy
from dataclasses import dataclass
from typing import Any, Literal, cast

Expand All @@ -18,7 +16,16 @@
)
from haystack.components.agents.state.state_utils import merge_lists
from haystack.components.agents.tool_calling import _run_tool, _run_tool_async
from haystack.components.agents.utils import _record_context_tokens
from haystack.components.agents.utils import (
_record_context_tokens,
_record_llm_usage,
_record_tool_calls,
_render_prompt_messages,
_select_tools_by_name,
_spawn_tools,
_template_for_role,
_validate_prompt_message_blocks,
)
from haystack.components.builders import ChatPromptBuilder
from haystack.components.generators.chat.types import ChatGenerator
from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict
Expand Down Expand Up @@ -59,11 +66,6 @@

logger = logging.getLogger(__name__)

# Regex to detect the Jinja2 chat template syntax
_JINJA2_CHAT_TEMPLATE_RE = re.compile(r"\{%\s*message\s")
# Regex to extract the role from a Jinja2 message block, e.g. {% message role="user" %}
_JINJA2_MESSAGE_ROLE_RE = re.compile(r'\{%\s*message\s+role\s*=\s*["\'](\w+)["\']')

# `exit_reason` values the Agent sets when it stops without a tool exit condition: a tool-call-free reply, or the
# `max_agent_steps` budget running out. A tool exit condition instead reports the tool's name.
_EXIT_REASON_TEXT = "text"
Expand Down Expand Up @@ -94,69 +96,6 @@
}


def _accumulate_usage(current: Any, new: Any) -> Any:
"""
Recursively sum numeric leaf values across two usage-like dicts.

Used to aggregate `ChatMessage.meta["usage"]` payloads across LLM calls in a run. Nested dicts (e.g. OpenAI's
`completion_tokens_details`) are merged recursively; numeric leaves are summed; other types fall back to the new
value.

:param current: The current accumulated usage data.
:param new: The new usage data to merge in.
"""
if isinstance(current, dict) and isinstance(new, dict):
result = dict(current)
for k, v in new.items():
result[k] = _accumulate_usage(result[k], v) if k in result else deepcopy(v)
return result
if isinstance(current, (int, float)) and isinstance(new, (int, float)):
return current + new
return new


def _record_llm_usage(state: State, llm_messages: list[ChatMessage]) -> None:
"""
Aggregate token usage from the latest LLM messages into the State.

Only writes when at least one message reports `meta["usage"]`, so generators that don't surface usage data
leave `token_usage` at its default empty dict rather than overwriting it.

:param state: The Agent's State, used to read the running `token_usage` total and write back the new total.
:param llm_messages: The ChatMessage objects returned from the latest LLM call. Token usage is read from each
message's `meta["usage"]` field, if present.
"""
current = state.data.get("token_usage")
updated = False
for msg in llm_messages:
usage = msg.meta.get("usage")
if isinstance(usage, dict):
current = _accumulate_usage(current or {}, usage)
updated = True
if updated:
state.set("token_usage", current)


def _record_tool_calls(state: State, tool_messages: list[ChatMessage]) -> None:
"""
Increment per-tool call counts in the State for every successfully dispatched tool.

:param state: The Agent's State, used to read the running `tool_call_counts` map and write back the new totals.
:param tool_messages: The ChatMessage objects returned from the latest tool execution. Per-tool counts are
incremented based on each message's `tool_call_result.origin.tool_name`.
"""
counts = state.data.get("tool_call_counts") or {}
updated = False
for tm in tool_messages:
if tm.tool_call_result is None:
continue
name = tm.tool_call_result.origin.tool_name
counts[name] = counts.get(name, 0) + 1
updated = True
if updated:
state.set("tool_call_counts", counts)


def _get_run_method_params(instance: "Agent") -> set[str]:
"""Derive the parameter names of the Agent.run method via introspection."""
sig = inspect.signature(instance.run)
Expand Down Expand Up @@ -238,137 +177,6 @@ def _pending_tool_call_messages_from_state(state: State) -> list[ChatMessage]:
return [last_message] if last_message.tool_calls else []


def _select_tools_by_name(configured_tools: ToolsType, names: list[str]) -> list[Tool | Toolset]:
"""
Select configured tools by name for a single run.

Standalone Tools are kept when their name is requested. A Toolset that exposes a requested name is replaced by a
per-run `spawn()` (an isolated copy) with the requested names registered as its `_selected_tool_names`, so
dynamic toolsets such as SearchableToolset preserve their behavior (search/lazy-loading) over the selected subset
without mutating the shared, configured Toolset.

:param configured_tools: The tools configured on the Agent.
:param names: The requested tool names.
:returns: The selected standalone Tools and/or spawned, selection-scoped Toolsets.
:raises ValueError: If no tools were configured, or if any requested name is not a valid tool name.
"""
if not configured_tools:
raise ValueError("No tools were configured for the Agent at initialization.")

requested_names = set(names)
items: list[Tool | Toolset] = (
[configured_tools] if isinstance(configured_tools, Toolset) else list(configured_tools)
)

# Resolve selectable names per item. For Toolsets we use get_selectable_tools() so dynamic toolsets
# (e.g. SearchableToolset) offer their full catalog by name, not just the tools exposed by iteration.
selectable_per_item: list[tuple[Tool | Toolset, set[str]]] = []
valid_tool_names: set[str] = set()
for item in items:
item_names = {tool.name for tool in item.get_selectable_tools()} if isinstance(item, Toolset) else {item.name}
selectable_per_item.append((item, item_names))
valid_tool_names |= item_names

invalid_tool_names = requested_names - valid_tool_names
if invalid_tool_names:
raise ValueError(
f"The following tool names are not valid: {invalid_tool_names}. Valid tool names are: {valid_tool_names}."
)

selected: list[Tool | Toolset] = []
for item, item_names in selectable_per_item:
matched = requested_names & item_names
if not matched:
continue
if isinstance(item, Toolset):
# Apply the selection to a per-run copy so the shared, configured Toolset is never mutated.
spawned = item.spawn()
spawned._selected_tool_names = matched
selected.append(spawned)
else:
selected.append(item)
return selected


def _spawn_tools(tools: ToolsType) -> ToolsType:
"""
Return per-run copies of `tools`, replacing each Toolset with an isolated `spawn()` (Tools are passed through).

This isolates run-scoped Toolset state (e.g. a SearchableToolset's discovered tools and any active name
selection) so that concurrent runs sharing the same configured Toolset — such as parallel sub-agent tool calls
or concurrent requests against one Agent — don't corrupt each other.
"""
if isinstance(tools, Toolset):
return tools.spawn()
return [item.spawn() if isinstance(item, Toolset) else item for item in tools]


def _validate_prompt_message_blocks(user_prompt: str | None, system_prompt: str | None) -> None:
"""
Validate explicit Jinja2 message blocks in Agent prompts.

:param user_prompt: Optional user prompt template.
:param system_prompt: Optional system prompt template.
:raises ValueError: If a prompt contains multiple message blocks or a literal block role is invalid.
"""
if user_prompt is not None:
message_blocks = _JINJA2_CHAT_TEMPLATE_RE.findall(user_prompt)
roles = _JINJA2_MESSAGE_ROLE_RE.findall(user_prompt)
if len(message_blocks) > 1:
raise ValueError(f"user_prompt must define exactly one message block, found {len(message_blocks)}.")
if roles and roles[0] != "user":
raise ValueError(f"user_prompt message block must have role 'user', found role '{roles[0]}'.")

if system_prompt is not None and _JINJA2_CHAT_TEMPLATE_RE.search(system_prompt):
message_blocks = _JINJA2_CHAT_TEMPLATE_RE.findall(system_prompt)
roles = _JINJA2_MESSAGE_ROLE_RE.findall(system_prompt)
if len(message_blocks) > 1:
raise ValueError(f"system_prompt must define exactly one message block, found {len(message_blocks)}.")
if roles and roles[0] != "system":
raise ValueError(f"system_prompt message block must have role 'system', found role '{roles[0]}'.")


def _template_for_role(prompt: str, role: str) -> str:
"""
Convert a prompt into a ChatPromptBuilder string template for the expected role.

:param prompt: Prompt template, with or without an explicit Jinja2 message block.
:param role: Role to use when wrapping a plain string prompt.
:returns: The original message-block template, or a plain string prompt wrapped in one message block.
"""
if _JINJA2_CHAT_TEMPLATE_RE.search(prompt):
return prompt
return f'{{% message role="{role}" %}}{prompt}{{% endmessage %}}'


def _render_prompt_messages(
*, prompt_builder: ChatPromptBuilder, expected_role: ChatRole, prompt_label: str, kwargs: dict[str, Any]
) -> list[ChatMessage]:
"""
Render one Agent prompt and validate the rendered message.

:param prompt_builder: Builder configured with the prompt template.
:param expected_role: Role the rendered message must have.
:param prompt_label: Prompt name used in error messages.
:param kwargs: Runtime values available to the prompt template.
:returns: A single rendered prompt message.
:raises ValueError: If the prompt renders to zero, multiple, or wrong-role messages.
"""
prompt_kwargs = {var: kwargs[var] for var in prompt_builder.variables if var in kwargs}
prompt_messages = prompt_builder.run(**prompt_kwargs)["prompt"]
if len(prompt_messages) != 1:
raise ValueError(
f"{prompt_label} must render to exactly one {expected_role.value} message. "
f"Got {len(prompt_messages)} messages."
)
if not prompt_messages[0].is_from(expected_role):
raise ValueError(
f"{prompt_label} must render to a {expected_role.value} message. "
f"Got a message with role {prompt_messages[0].role}."
)
return prompt_messages


@dataclass(kw_only=True)
class _ExecutionContext:
"""
Expand Down
Loading
Loading