diff --git a/docs-website/docs/pipeline-components/agents-1/agent.mdx b/docs-website/docs/pipeline-components/agents-1/agent.mdx index 837af2000d1..5b936b0e635 100644 --- a/docs-website/docs/pipeline-components/agents-1/agent.mdx +++ b/docs-website/docs/pipeline-components/agents-1/agent.mdx @@ -69,7 +69,7 @@ print(response["tool_call_counts"]) # {"calculator": 1} - `streaming_callback`: A callback invoked for each streamed token. Use the built-in `print_streaming_chunk` for console output. - `max_agent_steps`: Maximum number of LLM + tool call iterations before the agent stops. Defaults to `100`. - `raise_on_tool_invocation_failure`: If `True`, raises an exception when a tool call fails. If `False` (default), the error is passed back to the LLM as a message so it can recover. -- `hooks`: A dict mapping a hook point (`"before_llm"`, `"before_tool"`, `"after_tool"`, `"on_exit"`) to a list of hooks the agent runs at that point. Hooks receive the live `State` and influence the run by mutating it — for example, to build run-time context or require human confirmation of tool calls. See [Hooks](./hooks.mdx) and [Human in the Loop](./human-in-the-loop.mdx). +- `hooks`: A dict mapping a hook point (`"before_run"`, `"before_llm"`, `"before_tool"`, `"after_tool"`, `"on_exit"`, `"after_run"`) to a list of hooks the agent runs at that point. Hooks receive the live `State` and influence the run by mutating it — for example, to build run-time context or require human confirmation of tool calls. See [Hooks](./hooks.mdx) and [Human in the Loop](./human-in-the-loop.mdx). - `tool_concurrency_limit`: Maximum number of tool calls to execute at the same time. Defaults to `4`; set to `1` to disable parallel tool execution. - `tool_streaming_callback_passthrough`: If `True`, passes the streaming callback to tools that accept it. diff --git a/docs-website/docs/pipeline-components/agents-1/hooks.mdx b/docs-website/docs/pipeline-components/agents-1/hooks.mdx index b586078fb84..c0833905e50 100644 --- a/docs-website/docs/pipeline-components/agents-1/hooks.mdx +++ b/docs-website/docs/pipeline-components/agents-1/hooks.mdx @@ -2,12 +2,12 @@ title: "Hooks" id: hooks slug: "/hooks" -description: "Hooks let you run custom logic at defined points of an Agent's run loop — before each LLM call, before and after tool execution, and on exit." +description: "Hooks let you run custom logic at defined points of an Agent's run loop — at the start and end of a run, before each LLM call, before and after tool execution, and on exit." --- # Hooks -Hooks let you run custom logic at defined points of an [`Agent`](./agent.mdx)'s run loop — before each LLM call, before and after tool execution, and on exit. +Hooks let you run custom logic at defined points of an [`Agent`](./agent.mdx)'s run loop — at the start and end of a run, before each LLM call, before and after tool execution, and on exit.
@@ -30,10 +30,12 @@ This enables patterns such as building run-time system context, retrieving memor ### Hook points +- `before_run`: Runs once per run, after the state is initialized and before the first chat-generator call. Use it to rewrite the initial messages or seed state — for example, to turn the user query into a task brief — without re-running on every step like `before_llm` does. - `before_llm`: Runs before each chat-generator call. - `before_tool`: Runs after the model requests tool calls, before any tools run. After these hooks run, the Agent re-reads the current last message from `state.data["messages"]`. If that message contains tool calls, those calls are executed. If it does not, no tools run for that step, no tool-based exit condition is triggered, and the Agent loops back to the next LLM call unless `max_agent_steps` has been reached. - `after_tool`: Runs after tools execute, once their result messages are in `state.data["messages"]`, before the exit-condition check and the next LLM call. Use it to rewrite the freshly produced tool-result messages — for example, to offload, redact, truncate, or summarize results. It does not run on the plain-text exit step. It does still run when a `before_tool` hook removed the pending tool calls: no tools executed on that step, so don't assume the last message is a fresh tool result. -- `on_exit`: Runs when the Agent is about to stop on an exit condition. An `on_exit` hook can keep the Agent running by setting the `continue_run` control flag (`state.set("continue_run", True)`), usually alongside a message telling the model what to do next. `on_exit` hooks run when the Agent stops on an exit condition, but not when it stops because `max_agent_steps` is reached. +- `on_exit`: Runs when the Agent is about to stop on an exit condition. An `on_exit` hook can keep the Agent running by setting the `continue_run` control flag (`state.set("continue_run", True)`), usually alongside a message telling the model what to do next. `on_exit` hooks run when the Agent stops on an exit condition, but not when it stops because `max_agent_steps` is reached — use `after_run` for logic that must run however the run ends. +- `after_run`: Runs once per run, after the step loop has ended and before the Agent builds its return value — regardless of whether the run stopped on an exit condition or because `max_agent_steps` was reached (unlike `on_exit`). Mutations to the state, such as appending a final message, are reflected in the returned `messages` / `last_message` and `state_schema` outputs. Setting `continue_run` here has no effect. Registering a hook under an unknown hook point raises a `ValueError` at construction. A hook class can declare an `allowed_hook_points` attribute listing the hook points it supports; the Agent validates it and fails fast if the hook is registered somewhere it doesn't belong. diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py index 5dc2ed2f988..94595a18e99 100644 --- a/haystack/components/agents/agent.py +++ b/haystack/components/agents/agent.py @@ -23,7 +23,17 @@ from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict from haystack.dataclasses import ChatMessage, ChatRole, StreamingCallbackT, select_streaming_callback from haystack.hooks.invocation import _run_hooks, _run_hooks_async -from haystack.hooks.protocol import AFTER_TOOL, BEFORE_LLM, BEFORE_TOOL, ON_EXIT, VALID_HOOK_POINTS, Hook, HookPoint +from haystack.hooks.protocol import ( + AFTER_RUN, + AFTER_TOOL, + BEFORE_LLM, + BEFORE_RUN, + BEFORE_TOOL, + ON_EXIT, + VALID_HOOK_POINTS, + Hook, + HookPoint, +) from haystack.hooks.utils import ( _deserialize_hooks_dictionary, _serialize_hooks_dictionary, @@ -579,6 +589,9 @@ def __init__( # noqa: PLR0913 :param hooks: A dictionary mapping a hook point to a list of hooks the Agent runs at that point. Each hook receives the live `State` and influences the run by mutating it in place; hooks for a hook point run in list order. Valid hook points are: + - "before_run": Runs once per run, after the state is initialized and before the first chat-generator + call. Use it to rewrite the initial messages or seed state (e.g. turn the user query into a task + brief) without re-running on every step like "before_llm" does. - "before_llm": Runs before each chat-generator call. - "before_tool": Runs after the model requests tool calls, before any tools run. After these hooks run, the Agent re-reads the current last message from `state.data["messages"]`. If that message contains tool @@ -592,6 +605,11 @@ def __init__( # noqa: PLR0913 Agent running by setting the `continue_run` control flag (`state.set("continue_run", True)`), usually alongside a message telling the model what to do next. "on_exit" hooks run when the Agent stops on an exit condition, but not when it stops because `max_agent_steps` is reached. + - "after_run": Runs once per run, after the step loop has ended and before the Agent builds its return + value — regardless of whether the run stopped on an exit condition or because `max_agent_steps` was + reached (unlike "on_exit"). Mutations to the state (e.g. appending a final message) are reflected in + the returned `messages` / `last_message` and `state_schema` outputs. Setting `continue_run` here has + no effect. :raises TypeError: If the chat_generator does not support tools parameter in its run method. :raises ValueError: If any `user_prompt` variable overlaps with the `state_schema` or `run` method parameters, if a hook is registered under an unknown hook point, or if a hook is registered under a hook point it does @@ -1030,6 +1048,7 @@ def run( with self._create_agent_span(exe_context.tools) as span: span.set_content_tag("haystack.agent.input", agent_inputs) + _run_hooks(self.hooks, BEFORE_RUN, exe_context.state) while exe_context.counter < self.max_agent_steps: if not self._run_step(exe_context, span): break @@ -1038,6 +1057,7 @@ def run( "Agent reached maximum agent steps of {max_agent_steps}, stopping.", max_agent_steps=self.max_agent_steps, ) + _run_hooks(self.hooks, AFTER_RUN, exe_context.state) result = _public_outputs(exe_context.state) if msgs := result.get("messages"): result["last_message"] = msgs[-1] @@ -1102,6 +1122,7 @@ async def run_async( with self._create_agent_span(exe_context.tools) as span: span.set_content_tag("haystack.agent.input", agent_inputs) + await _run_hooks_async(self.hooks, BEFORE_RUN, exe_context.state) while exe_context.counter < self.max_agent_steps: if not await self._run_step_async(exe_context, span): break @@ -1110,6 +1131,7 @@ async def run_async( "Agent reached maximum agent steps of {max_agent_steps}, stopping.", max_agent_steps=self.max_agent_steps, ) + await _run_hooks_async(self.hooks, AFTER_RUN, exe_context.state) result = _public_outputs(exe_context.state) if msgs := result.get("messages"): result["last_message"] = msgs[-1] diff --git a/haystack/hooks/__init__.py b/haystack/hooks/__init__.py index d85f7b795de..1a57253837c 100644 --- a/haystack/hooks/__init__.py +++ b/haystack/hooks/__init__.py @@ -8,15 +8,27 @@ from lazy_imports import LazyImporter _import_structure = { - "protocol": ["Hook", "HookPoint", "BEFORE_LLM", "BEFORE_TOOL", "AFTER_TOOL", "ON_EXIT", "VALID_HOOK_POINTS"], + "protocol": [ + "Hook", + "HookPoint", + "BEFORE_RUN", + "BEFORE_LLM", + "BEFORE_TOOL", + "AFTER_TOOL", + "ON_EXIT", + "AFTER_RUN", + "VALID_HOOK_POINTS", + ], "from_function": ["FunctionHook", "hook"], } if TYPE_CHECKING: from .from_function import FunctionHook as FunctionHook from .from_function import hook as hook + from .protocol import AFTER_RUN as AFTER_RUN from .protocol import AFTER_TOOL as AFTER_TOOL from .protocol import BEFORE_LLM as BEFORE_LLM + from .protocol import BEFORE_RUN as BEFORE_RUN from .protocol import BEFORE_TOOL as BEFORE_TOOL from .protocol import ON_EXIT as ON_EXIT from .protocol import VALID_HOOK_POINTS as VALID_HOOK_POINTS diff --git a/haystack/hooks/protocol.py b/haystack/hooks/protocol.py index 61302951810..1527c67f058 100644 --- a/haystack/hooks/protocol.py +++ b/haystack/hooks/protocol.py @@ -7,12 +7,14 @@ from haystack.components.agents.state.state import State # Points in the Agent's run loop at which hooks can be registered. -HookPoint = Literal["before_llm", "before_tool", "after_tool", "on_exit"] +HookPoint = Literal["before_run", "before_llm", "before_tool", "after_tool", "on_exit", "after_run"] +BEFORE_RUN: HookPoint = "before_run" BEFORE_LLM: HookPoint = "before_llm" BEFORE_TOOL: HookPoint = "before_tool" AFTER_TOOL: HookPoint = "after_tool" ON_EXIT: HookPoint = "on_exit" +AFTER_RUN: HookPoint = "after_run" VALID_HOOK_POINTS: tuple[HookPoint, ...] = get_args(HookPoint) diff --git a/releasenotes/notes/add-before-run-after-run-hook-points-79609e8eda11fb6e.yaml b/releasenotes/notes/add-before-run-after-run-hook-points-79609e8eda11fb6e.yaml new file mode 100644 index 00000000000..54b41fc2ace --- /dev/null +++ b/releasenotes/notes/add-before-run-after-run-hook-points-79609e8eda11fb6e.yaml @@ -0,0 +1,11 @@ +--- +features: + - | + Added ``before_run`` and ``after_run`` hook points to the ``Agent``. Hooks registered under ``before_run`` run + exactly once per run, after the state is initialized and before the first LLM call — use them to rewrite the + initial messages or seed state without re-running on every step like ``before_llm`` does. Hooks registered under + ``after_run`` run exactly once per run, after the step loop has ended and before the Agent builds its return + value, regardless of whether the run stopped on an exit condition or because ``max_agent_steps`` was reached + (unlike ``on_exit``); mutations to the state are reflected in the returned ``messages``, ``last_message``, and + ``state_schema`` outputs. Register them like any other hook: + ``hooks={"before_run": [my_hook], "after_run": [my_other_hook]}``. diff --git a/test/components/agents/test_agent_hooks.py b/test/components/agents/test_agent_hooks.py index 342ac0c3105..314be7c2800 100644 --- a/test/components/agents/test_agent_hooks.py +++ b/test/components/agents/test_agent_hooks.py @@ -53,6 +53,33 @@ def record_b(state: State) -> None: state.set("trace", ["b"]) +@hook +def record_before_run(state: State) -> None: + state.set("trace", ["before_run"]) + + +@hook +def record_on_exit(state: State) -> None: + state.set("trace", ["on_exit"]) + + +@hook +def record_after_run(state: State) -> None: + state.set("trace", ["after_run"]) + + +@hook +def rewrite_query_into_brief(state: State) -> None: + # Realistic before_run use case: turn the user query into a task brief, once, before the first LLM call. + state.set("messages", [ChatMessage.from_user("BRIEF")], handler_override=replace_values) + + +@hook +def write_report(state: State) -> None: + # Realistic after_run use case: produce a final artifact from the collected state, however the run ended. + state.set("messages", [ChatMessage.from_assistant("REPORT")]) + + @hook def build_context(state: State) -> None: if state.get("step_count") == 0: @@ -176,6 +203,114 @@ def test_continue_run_is_not_exposed_as_output(self): assert "continue_run" not in result +class TestBeforeRunHook: + def test_runs_once_across_multi_step_run(self): + agent = _agent( + MockChatGenerator(), + tools=[save], + state_schema={"trace": {"type": list}}, + hooks={"before_run": [record_before_run]}, + ) + agent.chat_generator.run = MagicMock( + side_effect=[ + {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]}, + {"replies": [ChatMessage.from_assistant("done")]}, + ] + ) + result = agent.run(messages=[ChatMessage.from_user("hi")]) + # Unlike before_llm, before_run fires once per run, not once per step. + assert result["trace"] == ["before_run"] + assert agent.chat_generator.run.call_count == 2 + + def test_rewrites_messages_before_first_llm_call(self): + agent = _agent(MockChatGenerator(), hooks={"before_run": [rewrite_query_into_brief]}) + agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]}) + result = agent.run(messages=[ChatMessage.from_user("original query")]) + first_call_messages = agent.chat_generator.run.call_args_list[0].kwargs["messages"] + assert [m.text for m in first_call_messages] == ["BRIEF"] + assert [m.text for m in result["messages"]] == ["BRIEF", "done"] + + def test_hooks_run_in_list_order(self): + agent = _agent( + MockChatGenerator(), state_schema={"trace": {"type": list}}, hooks={"before_run": [record_a, record_b]} + ) + agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]}) + result = agent.run(messages=[ChatMessage.from_user("hi")]) + assert result["trace"] == ["a", "b"] + + +class TestAfterRunHook: + def test_runs_once_on_text_exit(self): + agent = _agent( + MockChatGenerator(), state_schema={"trace": {"type": list}}, hooks={"after_run": [record_after_run]} + ) + agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]}) + result = agent.run(messages=[ChatMessage.from_user("hi")]) + assert result["trace"] == ["after_run"] + assert agent.chat_generator.run.call_count == 1 + + def test_runs_on_tool_based_exit(self): + agent = _agent( + MockChatGenerator(), + tools=[final_answer], + exit_conditions=["final_answer"], + state_schema={"trace": {"type": list}}, + hooks={"after_run": [record_after_run]}, + ) + agent.chat_generator.run = MagicMock( + return_value={ + "replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("final_answer", {"answer": "42"})])] + } + ) + result = agent.run(messages=[ChatMessage.from_user("hi")]) + assert result["trace"] == ["after_run"] + + def test_runs_when_max_agent_steps_is_exhausted(self): + # The motivating use case: unlike on_exit, after_run still fires when the step budget runs out, + # so a report-writing hook cannot be skipped. + agent = _agent(MockChatGenerator(), tools=[save], max_agent_steps=2, hooks={"after_run": [write_report]}) + agent.chat_generator.run = MagicMock( + return_value={"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]} + ) + result = agent.run(messages=[ChatMessage.from_user("hi")]) + assert agent.chat_generator.run.call_count == 2 + assert result["last_message"].text == "REPORT" + + def test_mutations_are_reflected_in_the_returned_result(self): + agent = _agent(MockChatGenerator(), hooks={"after_run": [write_report]}) + agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]}) + result = agent.run(messages=[ChatMessage.from_user("hi")]) + assert [m.text for m in result["messages"]] == ["hi", "done", "REPORT"] + assert result["last_message"].text == "REPORT" + + def test_continue_run_has_no_effect(self): + agent = _agent(MockChatGenerator(), max_agent_steps=5, hooks={"after_run": [always_continue]}) + agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]}) + agent.run(messages=[ChatMessage.from_user("hi")]) + # The step loop is already over when after_run fires; setting continue_run cannot restart it. + assert agent.chat_generator.run.call_count == 1 + + def test_on_exit_runs_before_after_run(self): + agent = _agent( + MockChatGenerator(), + state_schema={"trace": {"type": list}}, + hooks={"on_exit": [record_on_exit], "after_run": [record_after_run]}, + ) + agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]}) + result = agent.run(messages=[ChatMessage.from_user("hi")]) + assert result["trace"] == ["on_exit", "after_run"] + + def test_allowed_hook_points_is_enforced_for_new_points(self): + class BeforeRunOnlyHook: + allowed_hook_points = ("before_run",) + + def run(self, state: State) -> None: + pass + + with pytest.raises(ValueError): + Agent(chat_generator=MockChatGenerator(), hooks={"after_run": [BeforeRunOnlyHook()]}) + + class TestBeforeLlmHook: def test_step_0_only_context_injected_once_across_loops(self): agent = _agent(MockChatGenerator(), tools=[save], hooks={"before_llm": [build_context]}) @@ -405,12 +540,19 @@ def test_to_dict_from_dict_roundtrip(self, monkeypatch): agent = Agent( chat_generator=OpenAIChatGenerator(), tools=[save], - hooks={"before_llm": [build_context], "on_exit": [require_save]}, + hooks={ + "before_run": [rewrite_query_into_brief], + "before_llm": [build_context], + "on_exit": [require_save], + "after_run": [write_report], + }, ) restored = Agent.from_dict(agent.to_dict()) - assert set(restored.hooks) == {"before_llm", "on_exit"} + assert set(restored.hooks) == {"before_run", "before_llm", "on_exit", "after_run"} + assert restored.hooks["before_run"][0].function is rewrite_query_into_brief.function assert restored.hooks["before_llm"][0].function is build_context.function assert restored.hooks["on_exit"][0].function is require_save.function + assert restored.hooks["after_run"][0].function is write_report.function class TestAgentHooksAsync: @@ -440,6 +582,32 @@ async def test_after_tool_hook_rewrites_results(self): "OFFLOADED" ] + @pytest.mark.asyncio + async def test_sync_before_run_and_after_run_hooks_run_in_async_run(self): + agent = _agent( + MockChatGenerator(), + state_schema={"trace": {"type": list}}, + hooks={"before_run": [record_before_run], "after_run": [record_after_run]}, + ) + agent.chat_generator.run_async = AsyncMock(return_value={"replies": [ChatMessage.from_assistant("done")]}) + result = await agent.run_async(messages=[ChatMessage.from_user("hi")]) + assert result["trace"] == ["before_run", "after_run"] + + @pytest.mark.asyncio + async def test_async_after_run_hook_runs_when_max_agent_steps_is_exhausted(self): + async def write_report_async(state: State) -> None: + state.set("messages", [ChatMessage.from_assistant("REPORT")]) + + agent = _agent( + MockChatGenerator(), tools=[save], max_agent_steps=2, hooks={"after_run": [hook(write_report_async)]} + ) + agent.chat_generator.run_async = AsyncMock( + return_value={"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]} + ) + result = await agent.run_async(messages=[ChatMessage.from_user("hi")]) + assert agent.chat_generator.run_async.call_count == 2 + assert result["last_message"].text == "REPORT" + @pytest.mark.asyncio async def test_async_on_exit_hook_forces_extra_step(self): async def require_save_async(state: State) -> None: