diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py index 4361b89f8b..b44552b463 100644 --- a/haystack/components/agents/agent.py +++ b/haystack/components/agents/agent.py @@ -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 @@ -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 @@ -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" @@ -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) @@ -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: """ diff --git a/haystack/components/agents/utils.py b/haystack/components/agents/utils.py index 1542ae26d7..c34743a2cb 100644 --- a/haystack/components/agents/utils.py +++ b/haystack/components/agents/utils.py @@ -2,10 +2,14 @@ # # SPDX-License-Identifier: Apache-2.0 +import re +from copy import deepcopy from typing import Any from haystack.components.agents.state.state import State -from haystack.dataclasses import ChatMessage +from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder +from haystack.dataclasses import ChatMessage, ChatRole +from haystack.tools import Tool, Toolset, ToolsType # Input/output token key conventions across chat generators: most report OpenAI-style # `prompt_tokens`/`completion_tokens`; OpenAIResponsesChatGenerator reports `input_tokens`/`output_tokens`. @@ -13,6 +17,149 @@ _OUTPUT_TOKEN_KEYS = ("completion_tokens", "output_tokens") +# --------------------------- +# Run metadata helpers +# --------------------------- + + +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) + + +# --------------------------- +# Tool helpers +# --------------------------- + + +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] + + +# --------------------------- +# Context token helpers +# --------------------------- + + def _first_numeric(usage: dict[str, Any], keys: tuple[str, ...]) -> int: """ Return the first numeric value found under `keys` in `usage`, or 0 if none is present. @@ -60,3 +207,79 @@ def _record_context_tokens(state: State, llm_messages: list[ChatMessage]) -> Non tokens = _context_tokens_from_usage(usage) if tokens: state.set("context_tokens", tokens) + + +# --------------------------- +# Prompt helpers +# --------------------------- + +# 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+)["\']') + + +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 diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py index aaacacc29e..01616b3217 100644 --- a/test/components/agents/test_agent.py +++ b/test/components/agents/test_agent.py @@ -16,7 +16,7 @@ from openai.types.chat import ChatCompletionChunk, chat_completion_chunk from haystack import Document, Pipeline, component, tracing -from haystack.components.agents.agent import Agent, _accumulate_usage, _select_tools_by_name +from haystack.components.agents.agent import Agent from haystack.components.agents.state import State, merge_lists, replace_values from haystack.components.agents.tool_calling import _run_tool from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder @@ -1404,51 +1404,6 @@ async def test_max_steps_exit_async(self, weather_tool): assert result["exit_reason"] == "max_agent_steps" -class TestAccumulateUsage: - """Unit tests for the `_accumulate_usage` helper used to merge ChatGenerator usage dicts.""" - - def test_sums_flat_numeric_keys(self): - current = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} - new = {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5} - assert _accumulate_usage(current, new) == {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20} - - def test_merges_nested_detail_dicts_recursively(self): - current = {"prompt_tokens": 10, "completion_tokens_details": {"reasoning_tokens": 2, "audio_tokens": 0}} - new = { - "prompt_tokens": 4, - "completion_tokens_details": {"reasoning_tokens": 3, "audio_tokens": 1}, - "prompt_tokens_details": {"cached_tokens": 6}, - } - assert _accumulate_usage(current, new) == { - "prompt_tokens": 14, - "completion_tokens_details": {"reasoning_tokens": 5, "audio_tokens": 1}, - "prompt_tokens_details": {"cached_tokens": 6}, - } - - def test_adds_keys_missing_in_current(self): - assert _accumulate_usage({"prompt_tokens": 5}, {"completion_tokens": 7}) == { - "prompt_tokens": 5, - "completion_tokens": 7, - } - - def test_empty_current_dict_returns_copy_of_new(self): - new = {"prompt_tokens": 5, "details": {"cached_tokens": 1}} - result = _accumulate_usage({}, new) - assert result == new - # Nested dicts must be deep-copied so future merges don't mutate the source. - new["details"]["cached_tokens"] = 999 - assert result["details"]["cached_tokens"] == 1 - - def test_non_dict_non_numeric_falls_back_to_new(self): - # Strings, lists, or any other type that isn't a dict-or-number pair returns `new` unchanged. - assert _accumulate_usage("old-model", "new-model") == "new-model" - assert _accumulate_usage(5, "stringified") == "stringified" - assert _accumulate_usage({"model": "gpt-x"}, {"model": "gpt-y"}) == {"model": "gpt-y"} - - def test_sums_floats(self): - assert _accumulate_usage(1.5, 2.25) == 3.75 - - class TestAgentTracing: def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool): chat_generator = MockChatGeneratorWithoutRunAsync() @@ -1816,55 +1771,6 @@ def test_agent_span_has_parent_when_in_pipeline(self, spying_tracer, weather_too assert agent_span.parent_span == agent_component_span -class TestSelectToolsByName: - """Tests for the _select_tools_by_name helper (runtime tool-name selection).""" - - def test_selects_standalone_tools_by_name(self, weather_tool: Tool, component_tool: Tool): - result = _select_tools_by_name([weather_tool, component_tool], [weather_tool.name]) - assert result == [weather_tool] - - def test_raises_on_invalid_name(self, weather_tool: Tool, component_tool: Tool): - with pytest.raises( - ValueError, match="The following tool names are not valid: {'invalid_tool_name'}. Valid tool names are: ." - ): - _select_tools_by_name([weather_tool, component_tool], ["invalid_tool_name"]) - - def test_raises_when_no_tools_configured(self, weather_tool: Tool): - with pytest.raises(ValueError, match="No tools were configured for the Agent at initialization."): - _select_tools_by_name([], [weather_tool.name]) - - def test_returns_isolated_spawn_with_selection(self, weather_tool: Tool, component_tool: Tool): - """A Toolset exposing a requested name is replaced by an isolated spawn carrying the selection. - - The shared, configured Toolset is not mutated. - """ - toolset = Toolset([weather_tool, component_tool]) - - result = _select_tools_by_name([toolset], [weather_tool.name]) - - assert len(result) == 1 - spawned = result[0] - assert isinstance(spawned, Toolset) - assert spawned is not toolset # an isolated per-run copy - assert spawned._selected_tool_names == {weather_tool.name} - assert [tool.name for tool in spawned] == [weather_tool.name] - # The configured toolset is untouched. - assert toolset._selected_tool_names is None - - def test_mixed_standalone_tools_and_toolsets(self, weather_tool: Tool, component_tool: Tool): - toolset = Toolset([weather_tool]) - - result = _select_tools_by_name([component_tool, toolset], [weather_tool.name, component_tool.name]) - - # The standalone tool is passed through; the toolset is replaced by an isolated spawn with the selection. - assert component_tool in result - spawns = [t for t in result if isinstance(t, Toolset)] - assert len(spawns) == 1 - assert spawns[0] is not toolset - assert spawns[0]._selected_tool_names == {weather_tool.name} - assert toolset._selected_tool_names is None - - class TestAgentToolSelection: def test_tool_selection_new_tool(self, weather_tool: Tool, component_tool: Tool): chat_generator = MockChatGenerator("Hello") @@ -2000,12 +1906,16 @@ def test_system_prompt_incorrect_jinja2_syntax_raises(self, make_agent): with pytest.raises(TemplateSyntaxError): make_agent(system_prompt="{% message role='system' %}Incomplete syntax.") - def test_system_prompt_plain_string(self, make_agent): - agent = make_agent(system_prompt="You are a helpful assistant.") - assert agent._system_chat_prompt_builder is not None - result = agent.run(messages=[ChatMessage.from_user("Hi")]) - assert result["messages"][0].is_from(ChatRole.SYSTEM) - assert result["messages"][0].text == "You are a helpful assistant." + def test_prompt_wrong_role_raises_at_init(self, make_agent): + with pytest.raises(ValueError, match="system_prompt message block must have role 'system'"): + make_agent(system_prompt=_user_msg("This is a user message, not system.")) + with pytest.raises(ValueError, match="user_prompt message block must have role 'user'"): + make_agent(user_prompt=_sys_msg("This is a system message, not user.")) + + def test_dynamic_prompt_role_raises_at_runtime(self, make_agent): + agent = make_agent(user_prompt="{% message role=role_name %}Q: {{question}}{% endmessage %}") + with pytest.raises(ValueError, match="user_prompt must render to a user message"): + agent.run(messages=[], role_name="assistant", question="Will it snow?") def test_system_prompt_plain_string_with_template_variables(self, make_agent): agent = make_agent(system_prompt="You are an assistant for {{company}}. Your role is {{role}}.") @@ -2021,38 +1931,6 @@ def test_system_prompt_plain_string_with_template_variables(self, make_agent): assert "company" in input_names assert "role" in input_names - def test_system_prompt_with_template_variables(self, make_agent): - agent = make_agent(system_prompt=_sys_msg("You are an assistant for {{company}}. Your role is {{role}}.")) - assert agent._system_chat_prompt_builder is not None - assert set(agent._system_chat_prompt_builder.variables) == {"company", "role"} - - result = agent.run(messages=[ChatMessage.from_user("Hi")], company="Acme", role="support agent") - sys_msg = result["messages"][0] - assert sys_msg.is_from(ChatRole.SYSTEM) - assert sys_msg.text == "You are an assistant for Acme. Your role is support agent." - - input_names = set(agent.__haystack_input__._sockets_dict.keys()) - assert "company" in input_names - assert "role" in input_names - - def test_system_prompt_with_meta(self, make_agent): - agent = make_agent( - system_prompt="{% message role='system' meta={'key': 'value'} %}System message with meta{% endmessage %}" - ) - assert agent._system_chat_prompt_builder is not None - - result = agent.run(messages=[ChatMessage.from_user("Hi")]) - messages = result["messages"] - assert messages[0].is_from(ChatRole.SYSTEM) - assert messages[0].text == "System message with meta" - assert messages[0].meta == {"key": "value"} - - def test_user_prompt_only_variables_forwarded_to_builder(self, make_agent): - agent = make_agent(user_prompt=_user_msg("Question: {{question}}")) - # 'irrelevant_kwarg' is not a template variable — must not raise - result = agent.run(messages=[], question="Will it snow?", irrelevant_kwarg="unused") - assert "messages" in result - def test_user_prompt_plain_string_with_template_variables(self, make_agent): agent = make_agent(user_prompt="Question: {{question}}") result = agent.run(messages=[], question="Will it snow?") @@ -2062,23 +1940,6 @@ def test_user_prompt_plain_string_with_template_variables(self, make_agent): input_names = set(agent.__haystack_input__._sockets_dict.keys()) assert "question" in input_names - def test_user_prompt_with_template_variables(self, make_agent): - agent = make_agent( - user_prompt=_user_msg( - "Hello {{name|upper}}, check weather for: " - + "{% for c in cities %}{{c}}{% if not loop.last %}, {% endif %}{% endfor %}" - + " on {{date}}?" - ) - ) - result = agent.run(messages=[], name="Alice", cities=["Berlin", "Paris", "Rome"], date="2024-01-15") - user_messages = [m for m in result["messages"] if m.is_from(ChatRole.USER)] - assert user_messages[0].text == "Hello ALICE, check weather for: Berlin, Paris, Rome on 2024-01-15?" - - input_names = set(agent.__haystack_input__._sockets_dict.keys()) - assert "name" in input_names - assert "cities" in input_names - assert "date" in input_names - def test_user_prompt_appended_after_initial_messages(self, make_agent): agent = make_agent(user_prompt=_user_msg("And now: {{query}}")) initial_messages = [ChatMessage.from_user("First message")] @@ -2102,28 +1963,6 @@ def test_system_prompt_and_user_prompt(self, make_agent): user_messages = [m for m in messages if m.is_from(ChatRole.USER)] assert user_messages[0].text == "Tell me about pipelines in the Haystack context." - def test_prompt_wrong_role_raises_at_init(self, make_agent): - with pytest.raises(ValueError, match="system_prompt message block must have role 'system'"): - make_agent(system_prompt=_user_msg("This is a user message, not system.")) - - with pytest.raises(ValueError, match="user_prompt message block must have role 'user'"): - make_agent(user_prompt=_sys_msg("This is a system message, not user.")) - - def test_dynamic_prompt_role_raises_at_runtime(self, make_agent): - agent = make_agent(user_prompt="{% message role=role_name %}Question: {{question}}{% endmessage %}") - with pytest.raises(ValueError, match="user_prompt must render to a user message"): - agent.run(messages=[], role_name="assistant", question="Will it snow?") - - def test_prompt_multiple_message_blocks_raises_at_init(self, make_agent): - multi_message_prompt = """{% message role='system' %}You are a helpful assistant.{% endmessage %} - {% message role='user' %}How are you?{% endmessage %}""" - - with pytest.raises(ValueError, match="system_prompt must define exactly one message block"): - make_agent(system_prompt=multi_message_prompt) - - with pytest.raises(ValueError, match="user_prompt must define exactly one message block"): - make_agent(user_prompt=multi_message_prompt) - @pytest.mark.integration class TestAgentUserPromptInPipeline: diff --git a/test/components/agents/test_utils.py b/test/components/agents/test_utils.py index 71c448e758..116401e729 100644 --- a/test/components/agents/test_utils.py +++ b/test/components/agents/test_utils.py @@ -5,8 +5,117 @@ import pytest from haystack.components.agents.state import State, replace_values -from haystack.components.agents.utils import _context_tokens_from_usage, _record_context_tokens +from haystack.components.agents.utils import ( + _accumulate_usage, + _context_tokens_from_usage, + _record_context_tokens, + _render_prompt_messages, + _select_tools_by_name, + _template_for_role, + _validate_prompt_message_blocks, +) +from haystack.components.builders import ChatPromptBuilder from haystack.dataclasses import ChatMessage +from haystack.dataclasses.chat_message import ChatRole +from haystack.tools import Tool +from haystack.tools.toolset import Toolset + + +def _user_msg(text: str) -> str: + return f'{{% message role="user" %}}{text}{{% endmessage %}}' + + +def _sys_msg(text: str) -> str: + return f'{{% message role="system" %}}{text}{{% endmessage %}}' + + +def _tool_function(value: str) -> str: + return value + + +@pytest.fixture +def first_tool() -> Tool: + return Tool( + name="first_tool", description="First test tool.", parameters={"type": "object"}, function=_tool_function + ) + + +@pytest.fixture +def second_tool() -> Tool: + return Tool( + name="second_tool", description="Second test tool.", parameters={"type": "object"}, function=_tool_function + ) + + +class TestAccumulateUsage: + def test_sums_flat_numeric_keys(self): + current = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + new = {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5} + assert _accumulate_usage(current, new) == {"prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20} + + def test_merges_nested_detail_dicts_recursively(self): + current = {"prompt_tokens": 10, "completion_tokens_details": {"reasoning_tokens": 2, "audio_tokens": 0}} + new = { + "prompt_tokens": 4, + "completion_tokens_details": {"reasoning_tokens": 3, "audio_tokens": 1}, + "prompt_tokens_details": {"cached_tokens": 6}, + } + assert _accumulate_usage(current, new) == { + "prompt_tokens": 14, + "completion_tokens_details": {"reasoning_tokens": 5, "audio_tokens": 1}, + "prompt_tokens_details": {"cached_tokens": 6}, + } + + def test_adds_keys_missing_in_current(self): + assert _accumulate_usage({"prompt_tokens": 5}, {"completion_tokens": 7}) == { + "prompt_tokens": 5, + "completion_tokens": 7, + } + + def test_copies_new_nested_values(self): + new = {"details": {"cached_tokens": 1}} + result = _accumulate_usage({}, new) + new["details"]["cached_tokens"] = 999 + assert result == {"details": {"cached_tokens": 1}} + + def test_falls_back_to_new_for_non_numeric_values(self): + assert _accumulate_usage("old-model", "new-model") == "new-model" + assert _accumulate_usage(5, "stringified") == "stringified" + assert _accumulate_usage({"model": "old"}, {"model": "new"}) == {"model": "new"} + + def test_sums_floats(self): + assert _accumulate_usage(1.5, 2.25) == 3.75 + + +class TestSelectToolsByName: + def test_selects_standalone_tools_by_name(self, first_tool: Tool, second_tool: Tool): + assert _select_tools_by_name([first_tool, second_tool], [first_tool.name]) == [first_tool] + + def test_raises_for_invalid_name(self, first_tool: Tool): + with pytest.raises(ValueError, match="The following tool names are not valid"): + _select_tools_by_name([first_tool], ["unknown"]) + + def test_raises_when_no_tools_configured(self, first_tool: Tool): + with pytest.raises(ValueError, match="No tools were configured for the Agent at initialization."): + _select_tools_by_name([], [first_tool.name]) + + def test_spawns_toolsets_without_mutating_them(self, first_tool: Tool, second_tool: Tool): + toolset = Toolset([first_tool, second_tool]) + selected = _select_tools_by_name([toolset], [first_tool.name]) + spawned = selected[0] + assert isinstance(spawned, Toolset) + assert spawned is not toolset + assert spawned._selected_tool_names == {first_tool.name} + assert toolset._selected_tool_names is None + + def test_selects_standalone_tools_and_toolsets(self, first_tool: Tool, second_tool: Tool): + toolset = Toolset([first_tool]) + selected = _select_tools_by_name([second_tool, toolset], [first_tool.name, second_tool.name]) + assert second_tool in selected + spawned = next(item for item in selected if isinstance(item, Toolset)) + assert spawned is not toolset + assert spawned._selected_tool_names == {first_tool.name} + assert toolset._selected_tool_names is None class TestContextTokensFromUsage: @@ -138,3 +247,100 @@ def test_missing_or_empty_usage_leaves_value_untouched(self): _record_context_tokens(state, [ChatMessage.from_assistant("no usage here")]) _record_context_tokens(state, [ChatMessage.from_assistant("empty", meta={"usage": {}})]) assert state.get("context_tokens") == 0 + + +class TestPrompts: + def test_system_prompt_plain_string(self): + prompt_builder = ChatPromptBuilder(template=_template_for_role("You are a helpful assistant.", "system")) + messages = _render_prompt_messages( + prompt_builder=prompt_builder, expected_role=ChatRole.SYSTEM, prompt_label="system_prompt", kwargs={} + ) + assert messages[0].is_from(ChatRole.SYSTEM) + assert messages[0].text == "You are a helpful assistant." + + def test_system_prompt_with_template_variables(self): + prompt_builder = ChatPromptBuilder( + template=_template_for_role( + _sys_msg("You are an assistant for {{company}}. Your role is {{role}}."), "system" + ) + ) + messages = _render_prompt_messages( + prompt_builder=prompt_builder, + expected_role=ChatRole.SYSTEM, + prompt_label="system_prompt", + kwargs={"company": "Acme", "role": "support agent"}, + ) + sys_msg = messages[0] + assert sys_msg.is_from(ChatRole.SYSTEM) + assert sys_msg.text == "You are an assistant for Acme. Your role is support agent." + + def test_system_prompt_with_meta(self): + prompt_builder = ChatPromptBuilder( + template="{% message role='system' meta={'key': 'value'} %}System message with meta{% endmessage %}" + ) + messages = _render_prompt_messages( + prompt_builder=prompt_builder, expected_role=ChatRole.SYSTEM, prompt_label="system_prompt", kwargs={} + ) + assert messages[0].is_from(ChatRole.SYSTEM) + assert messages[0].text == "System message with meta" + assert messages[0].meta == {"key": "value"} + + def test_user_prompt_only_variables_forwarded_to_builder(self): + prompt_builder = ChatPromptBuilder(template=_user_msg("Question: {{question}}")) + # 'irrelevant_kwarg' is not a template variable — must not raise + messages = _render_prompt_messages( + prompt_builder=prompt_builder, + expected_role=ChatRole.USER, + prompt_label="user_prompt", + kwargs={"question": "Will it snow?", "irrelevant_kwarg": "unused"}, + ) + assert messages[0].text == "Question: Will it snow?" + + def test_user_prompt_with_template_variables(self): + prompt_builder = ChatPromptBuilder( + template=_user_msg( + "Hello {{name|upper}}, check weather for: " + + "{% for c in cities %}{{c}}{% if not loop.last %}, {% endif %}{% endfor %}" + + " on {{date}}?" + ) + ) + messages = _render_prompt_messages( + prompt_builder=prompt_builder, + expected_role=ChatRole.USER, + prompt_label="user_prompt", + kwargs={"name": "Alice", "cities": ["Berlin", "Paris", "Rome"], "date": "2024-01-15"}, + ) + assert messages[0].text == "Hello ALICE, check weather for: Berlin, Paris, Rome on 2024-01-15?" + + def test_prompt_wrong_role_raises(self): + with pytest.raises(ValueError, match="system_prompt message block must have role 'system'"): + _validate_prompt_message_blocks( + user_prompt=None, system_prompt=_user_msg("This is a user message, not system.") + ) + + with pytest.raises(ValueError, match="user_prompt message block must have role 'user'"): + _validate_prompt_message_blocks( + user_prompt=_sys_msg("This is a system message, not user."), system_prompt=None + ) + + def test_dynamic_prompt_role_raises(self): + prompt_builder = ChatPromptBuilder( + template="{% message role=role_name %}Question: {{question}}{% endmessage %}" + ) + with pytest.raises(ValueError, match="user_prompt must render to a user message"): + _render_prompt_messages( + prompt_builder=prompt_builder, + expected_role=ChatRole.USER, + prompt_label="user_prompt", + kwargs={"role_name": "assistant", "question": "Will it snow?"}, + ) + + def test_prompt_multiple_message_blocks_raises(self): + multi_message_prompt = """{% message role='system' %}You are a helpful assistant.{% endmessage %} + {% message role='user' %}How are you?{% endmessage %}""" + + with pytest.raises(ValueError, match="system_prompt must define exactly one message block"): + _validate_prompt_message_blocks(user_prompt=None, system_prompt=multi_message_prompt) + + with pytest.raises(ValueError, match="user_prompt must define exactly one message block"): + _validate_prompt_message_blocks(user_prompt=multi_message_prompt, system_prompt=None)