refactor: dedupe agent tool steps and colocate OpenAI tool merge helpers

Share client-tool planning and post-invoke logic between sync and async agent steps, and move OpenAI tool-definition merging into generators/utils to avoid a new chat submodule for schema import.

Co-authored-by: Cursor <[email protected]>

haystack/components/agents/agent.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import inspect
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from typing import Any, Literal, cast
 
@@ -326,6 +327,85 @@ class _ExecutionContext:
     counter: int = 0
 
 
+def _stop_step_for_client_tools(exe_context: _ExecutionContext) -> None:
+    exe_context.counter += 1
+    exe_context.state.set("step_count", exe_context.counter)
+    exe_context.state.set("exit_reason", _EXIT_REASON_CLIENT_TOOLS)
+
+
+def _plan_tool_step(
+    exe_context: _ExecutionContext, *, current_tools: list[Tool], client_side_tools: list[Tool]
+) -> tuple[_PendingToolCallPlan, list[ChatMessage]] | None:
+    """
+    Plan tool invocation for the current step.
+
+    :returns: ``(plan, pending_tool_call_messages)``, or ``None`` if the step stops for client-only tool calls.
+    """
+    pending_tool_call_messages = _pending_tool_call_messages_from_state(state=exe_context.state)
+    plan = _plan_pending_tool_calls(
+        pending_tool_call_messages, current_tools=current_tools, client_side_tools=client_side_tools
+    )
+    if plan.stop_for_client_tools and plan.invoke_messages is None:  # only client-side tools to invoke
+        _stop_step_for_client_tools(exe_context)
+        return None
+    return plan, pending_tool_call_messages
+
+
+def _commit_tool_results(exe_context: _ExecutionContext, tool_messages: list[ChatMessage]) -> None:
+    exe_context.state.set("messages", tool_messages)
+    _record_tool_calls(state=exe_context.state, tool_messages=tool_messages)
+
+
+def _finish_tool_step(
+    exe_context: _ExecutionContext,
+    *,
+    exit_conditions: list[str],
+    plan: _PendingToolCallPlan,
+    pending_tool_call_messages: list[ChatMessage],
+    tool_messages: list[ChatMessage],
+    check_exit_conditions: Callable[[list[ChatMessage], list[ChatMessage]], str | None],
+    continue_after_exit: Callable[[], bool],
+) -> bool:
+    """Returns True to continue the agent loop, False to stop."""
+    exe_context.counter += 1
+    exe_context.state.set("step_count", exe_context.counter)
+    if plan.stop_for_client_tools:
+        exe_context.state.set("exit_reason", _EXIT_REASON_CLIENT_TOOLS)
+        return False
+    exit_condition_tool = (
+        None if exit_conditions == ["text"] else check_exit_conditions(pending_tool_call_messages, tool_messages)
+    )
+    if exit_condition_tool is not None:
+        exe_context.state.set("exit_reason", exit_condition_tool)
+        return continue_after_exit()
+    return True
+
+
+async def _finish_tool_step_async(
+    exe_context: _ExecutionContext,
+    *,
+    exit_conditions: list[str],
+    plan: _PendingToolCallPlan,
+    pending_tool_call_messages: list[ChatMessage],
+    tool_messages: list[ChatMessage],
+    check_exit_conditions: Callable[[list[ChatMessage], list[ChatMessage]], str | None],
+    continue_after_exit: Callable[[], Awaitable[bool]],
+) -> bool:
+    """Returns True to continue the agent loop, False to stop."""
+    exe_context.counter += 1
+    exe_context.state.set("step_count", exe_context.counter)
+    if plan.stop_for_client_tools:
+        exe_context.state.set("exit_reason", _EXIT_REASON_CLIENT_TOOLS)
+        return False
+    exit_condition_tool = (
+        None if exit_conditions == ["text"] else check_exit_conditions(pending_tool_call_messages, tool_messages)
+    )
+    if exit_condition_tool is not None:
+        exe_context.state.set("exit_reason", exit_condition_tool)
+        return await continue_after_exit()
+    return True
+
+
 @component
 class Agent:
     """
@@ -1148,43 +1228,31 @@ class Agent:
                 return self._continue_after_exit_hooks(exe_context=exe_context)
 
             _run_hooks(hooks=self.hooks, hook_point=BEFORE_TOOL, state=exe_context.state)
-            # Re-read the pending tool calls from State so that any rewrites a before_tool hook made (e.g.
-            # ConfirmationHook rejecting or modifying calls) are honored by the executor.
-            pending_tool_call_messages = _pending_tool_call_messages_from_state(state=exe_context.state)
-            plan = _plan_pending_tool_calls(
-                pending_tool_call_messages, current_tools=current_tools, client_side_tools=exe_context.client_side_tools
+            planned = _plan_tool_step(
+                exe_context, current_tools=current_tools, client_side_tools=exe_context.client_side_tools
             )
-            if plan.stop_for_client_tools and plan.invoke_messages is None:  # only client-side tools to invoke
-                exe_context.counter += 1
-                exe_context.state.set("step_count", exe_context.counter)
-                exe_context.state.set("exit_reason", _EXIT_REASON_CLIENT_TOOLS)
+            if planned is None:
                 return False
+            plan, pending_tool_call_messages = planned
 
-            tool_execution_inputs = {
-                "messages": plan.invoke_messages,
-                "state": exe_context.state,
+            tool_messages, exe_context.state = _run_tool(
+                messages=plan.invoke_messages,
+                state=exe_context.state,
+                tools=current_tools,
                 **exe_context.tool_execution_inputs,
-                "tools": current_tools,
-            }
-            tool_messages, exe_context.state = _run_tool(**tool_execution_inputs)
-            exe_context.state.set("messages", tool_messages)
-            _record_tool_calls(state=exe_context.state, tool_messages=tool_messages)
+            )
+            _commit_tool_results(exe_context, tool_messages)
             _run_hooks(hooks=self.hooks, hook_point=AFTER_TOOL, state=exe_context.state)
 
-            exe_context.counter += 1
-            exe_context.state.set("step_count", exe_context.counter)
-            if plan.stop_for_client_tools:
-                exe_context.state.set("exit_reason", _EXIT_REASON_CLIENT_TOOLS)
-                return False
-            exit_condition_tool = (
-                None
-                if self.exit_conditions == ["text"]
-                else self._check_exit_conditions(llm_messages=pending_tool_call_messages, tool_messages=tool_messages)
+            return _finish_tool_step(
+                exe_context,
+                exit_conditions=self.exit_conditions,
+                plan=plan,
+                pending_tool_call_messages=pending_tool_call_messages,
+                tool_messages=tool_messages,
+                check_exit_conditions=self._check_exit_conditions,
+                continue_after_exit=lambda: self._continue_after_exit_hooks(exe_context=exe_context),
             )
-            if exit_condition_tool is not None:
-                exe_context.state.set("exit_reason", exit_condition_tool)
-                return self._continue_after_exit_hooks(exe_context=exe_context)
-            return True
 
     async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: tracing.Span) -> bool:
         """Execute one agent step asynchronously. Returns True to continue the loop, False to stop."""
@@ -1230,43 +1298,31 @@ class Agent:
                 return await self._continue_after_exit_hooks_async(exe_context=exe_context)
 
             await _run_hooks_async(hooks=self.hooks, hook_point=BEFORE_TOOL, state=exe_context.state)
-            # Re-read the pending tool calls from State so that any rewrites a before_tool hook made (e.g.
-            # ConfirmationHook rejecting or modifying calls) are honored by the executor.
-            pending_tool_call_messages = _pending_tool_call_messages_from_state(state=exe_context.state)
-            plan = _plan_pending_tool_calls(
-                pending_tool_call_messages, current_tools=current_tools, client_side_tools=exe_context.client_side_tools
+            planned = _plan_tool_step(
+                exe_context, current_tools=current_tools, client_side_tools=exe_context.client_side_tools
             )
-            if plan.stop_for_client_tools and plan.invoke_messages is None:  # only client-side tools to invoke
-                exe_context.counter += 1
-                exe_context.state.set("step_count", exe_context.counter)
-                exe_context.state.set("exit_reason", _EXIT_REASON_CLIENT_TOOLS)
+            if planned is None:
                 return False
+            plan, pending_tool_call_messages = planned
 
-            tool_execution_inputs = {
-                "messages": plan.invoke_messages,
-                "state": exe_context.state,
+            tool_messages, exe_context.state = await _run_tool_async(
+                messages=plan.invoke_messages,
+                state=exe_context.state,
+                tools=current_tools,
                 **exe_context.tool_execution_inputs,
-                "tools": current_tools,
-            }
-            tool_messages, exe_context.state = await _run_tool_async(**tool_execution_inputs)
-            exe_context.state.set("messages", tool_messages)
-            _record_tool_calls(state=exe_context.state, tool_messages=tool_messages)
+            )
+            _commit_tool_results(exe_context, tool_messages)
             await _run_hooks_async(hooks=self.hooks, hook_point=AFTER_TOOL, state=exe_context.state)
 
-            exe_context.counter += 1
-            exe_context.state.set("step_count", exe_context.counter)
-            if plan.stop_for_client_tools:
-                exe_context.state.set("exit_reason", _EXIT_REASON_CLIENT_TOOLS)
-                return False
-            exit_condition_tool = (
-                None
-                if self.exit_conditions == ["text"]
-                else self._check_exit_conditions(llm_messages=pending_tool_call_messages, tool_messages=tool_messages)
+            return await _finish_tool_step_async(
+                exe_context,
+                exit_conditions=self.exit_conditions,
+                plan=plan,
+                pending_tool_call_messages=pending_tool_call_messages,
+                tool_messages=tool_messages,
+                check_exit_conditions=self._check_exit_conditions,
+                continue_after_exit=lambda: self._continue_after_exit_hooks_async(exe_context=exe_context),
             )
-            if exit_condition_tool is not None:
-                exe_context.state.set("exit_reason", exit_condition_tool)
-                return await self._continue_after_exit_hooks_async(exe_context=exe_context)
-            return True
 
     def _check_exit_conditions(self, llm_messages: list[ChatMessage], tool_messages: list[ChatMessage]) -> str | None:
         """

haystack/components/generators/chat/_openai_tool_defs.py

@@ -1,60 +0,0 @@
-# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
-#
-# SPDX-License-Identifier: Apache-2.0
-
-from typing import Any
-
-
-def _openai_tool_definition_key(item: dict[str, Any]) -> str | None:
-    """Return a stable dedupe key for an OpenAI-style tool definition, or None if it cannot be keyed."""
-    tool_type = item.get("type")
-    nested = item.get("function")
-    if tool_type == "function" and isinstance(nested, dict):
-        name = nested.get("name")
-        if isinstance(name, str) and name:
-            return f"function:{name}"
-    name = item.get("name")
-    if isinstance(name, str) and name:
-        type_label = tool_type if isinstance(tool_type, str) else "tool"
-        return f"{type_label}:{name}"
-    return None
-
-
-def _merge_openai_tool_definitions(primary: list[dict[str, Any]], extra: Any) -> list[dict[str, Any]]:
-    """
-    Merge OpenAI API tool definition lists.
-
-    Entries from ``primary`` (typically from the ``tools`` run argument) win over ``extra`` (typically from
-    ``generation_kwargs["tools"]``) when they share the same dedupe key.
-    """
-    merged = list(primary)
-    seen = {key for item in primary if (key := _openai_tool_definition_key(item)) is not None}
-    if not isinstance(extra, list):
-        return merged
-    for item in extra:
-        if not isinstance(item, dict):
-            continue
-        key = _openai_tool_definition_key(item)
-        if key is None:
-            merged.append(item)
-            continue
-        if key in seen:
-            continue
-        seen.add(key)
-        merged.append(item)
-    return merged
-
-
-def _check_duplicate_openai_tool_definition_keys(tool_definitions: list[dict[str, Any]]) -> None:
-    """Raise ``ValueError`` when two keyed tool definitions share the same dedupe key."""
-    seen: set[str] = set()
-    duplicates: set[str] = set()
-    for item in tool_definitions:
-        key = _openai_tool_definition_key(item)
-        if key is None:
-            continue
-        if key in seen:
-            duplicates.add(key)
-        seen.add(key)
-    if duplicates:
-        raise ValueError(f"Duplicate tool names found: {duplicates}")

haystack/components/generators/chat/openai.py

@@ -23,12 +23,10 @@ from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
 from pydantic import BaseModel
 
 from haystack import component, default_from_dict, default_to_dict, logging
-from haystack.components.generators.chat._openai_tool_defs import (
-    _check_duplicate_openai_tool_definition_keys,
-    _merge_openai_tool_definitions,
-)
 from haystack.components.generators.utils import (
+    _check_duplicate_openai_tool_definition_keys,
     _convert_streaming_chunks_to_chat_message,
+    _merge_openai_tool_definitions,
     _normalize_messages,
     _serialize_object,
 )

haystack/components/generators/chat/openai_responses.py

@@ -13,11 +13,12 @@ from openai.types.responses import ParsedResponse, Response, ResponseOutputRefus
 from pydantic import BaseModel
 
 from haystack import component, default_from_dict, default_to_dict, logging
-from haystack.components.generators.chat._openai_tool_defs import (
+from haystack.components.generators.utils import (
     _check_duplicate_openai_tool_definition_keys,
     _merge_openai_tool_definitions,
+    _normalize_messages,
+    _serialize_object,
 )
-from haystack.components.generators.utils import _normalize_messages, _serialize_object
 from haystack.dataclasses import (
     ChatMessage,
     ChatRole,

haystack/components/generators/utils.py

@@ -226,3 +226,58 @@ def _normalize_messages(messages: list[ChatMessage] | str) -> list[ChatMessage]:
     if isinstance(messages, list) and all(isinstance(msg, ChatMessage) for msg in messages):
         return messages
     raise TypeError("Invalid messages type. Expected list[ChatMessage] or str.")
+
+
+def _openai_tool_definition_key(item: dict[str, Any]) -> str | None:
+    """Return a stable dedupe key for an OpenAI-style tool definition, or None if it cannot be keyed."""
+    tool_type = item.get("type")
+    nested = item.get("function")
+    if tool_type == "function" and isinstance(nested, dict):
+        name = nested.get("name")
+        if isinstance(name, str) and name:
+            return f"function:{name}"
+    name = item.get("name")
+    if isinstance(name, str) and name:
+        type_label = tool_type if isinstance(tool_type, str) else "tool"
+        return f"{type_label}:{name}"
+    return None
+
+
+def _merge_openai_tool_definitions(primary: list[dict[str, Any]], extra: Any) -> list[dict[str, Any]]:
+    """
+    Merge OpenAI API tool definition lists.
+
+    Entries from ``primary`` (typically from the ``tools`` run argument) win over ``extra`` (typically from
+    ``generation_kwargs["tools"]``) when they share the same dedupe key.
+    """
+    merged = list(primary)
+    seen = {key for item in primary if (key := _openai_tool_definition_key(item)) is not None}
+    if not isinstance(extra, list):
+        return merged
+    for item in extra:
+        if not isinstance(item, dict):
+            continue
+        key = _openai_tool_definition_key(item)
+        if key is None:
+            merged.append(item)
+            continue
+        if key in seen:
+            continue
+        seen.add(key)
+        merged.append(item)
+    return merged
+
+
+def _check_duplicate_openai_tool_definition_keys(tool_definitions: list[dict[str, Any]]) -> None:
+    """Raise ``ValueError`` when two keyed tool definitions share the same dedupe key."""
+    seen: set[str] = set()
+    duplicates: set[str] = set()
+    for item in tool_definitions:
+        key = _openai_tool_definition_key(item)
+        if key is None:
+            continue
+        if key in seen:
+            duplicates.add(key)
+        seen.add(key)
+    if duplicates:
+        raise ValueError(f"Duplicate tool names found: {duplicates}")