feat: Add Agent Hooks

[sjrl]

Related Issues

• fixes https://github.com/deepset-ai/haystack-private/issues/416

Proposed Changes:

Added lifecycle hooks to the Agent. Pass hooks as a dictionary mapping a lifecycle event to a list of hooks the Agent runs at that point: before_llm (before each chat-generator call), before_tool (after the model requests tool calls, before they run), and on_exit (when the Agent is about to stop). Each hook receives the live State and influences the run by mutating it in place; hooks for an event run in list order, and the same hook can be registered under multiple events.

A hook is any object with a run(state) method (optionally run_async(state)); use the @hook decorator to build one from a function. This enables patterns such as building run-time system context, retrieving memories, and requiring a tool to be called before finishing. An on_exit hook can keep the Agent running by appending a message that is no longer a valid exit, since the exit condition is re-evaluated after the hooks run.

The example below registers one hook at each placement to show what each can do:

from datetime import datetime, timezone
from typing import Annotated

from haystack.components.agents import Agent
from haystack.components.agents.state import State, replace_values
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.dataclasses import ChatMessage
from haystack.hooks import hook
from haystack.tools import tool


@tool
def search(query: Annotated[str, "The search query"]) -> str:
    """Search the web."""
    # Placeholder: would call a real search API
    return "Fusion startups reported net-energy-gain milestones this year."


@hook
def build_context(state: State) -> None:
    # before_llm: build run-time system context once, before the first model call.
    if state.get("step_count") == 0:
        now = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")
        system = ChatMessage.from_system(f"You are a research assistant. The current time is {now}.")
        state.set("messages", [system, *state.data["messages"]], handler_override=replace_values)


@hook
def audit_tool_calls(state: State) -> None:
    # before_tool: see which tools the model is about to run.
    pending = state.data["messages"][-1].tool_calls
    print(f"about to run: {[tc.tool_name for tc in pending]}")


@hook
def require_search(state: State) -> None:
    # on_exit: keep going until the agent has actually searched.
    if state.get("tool_call_counts", {}).get("search", 0) == 0:
        state.set("messages", [ChatMessage.from_system("Search before answering.")])


agent = Agent(
    chat_generator=OpenAIChatGenerator(model="gpt-5.4-nano"),
    tools=[search],
    hooks={
        "before_llm": [build_context],
        "before_tool": [audit_tool_calls],
        "on_exit": [require_search],
    },
)

result = agent.run(messages=[ChatMessage.from_user("What are the latest developments in fusion energy?")])
print(result["last_message"].text)

Class-based hooks may also implement the optional lifecycle methods warm_up / warm_up_async and close / close_async. The Agent calls them from its own warm_up / warm_up_async and close / close_async, so a hook can defer opening clients or reading credentials until warm-up and release them on close. For example, an on_exit hook can grade the Agent's answer with its own LLM and ask it to improve a weak answer before finishing, warming the judge's client during the Agent's warm-up:

class GradeFinalAnswer:
    """Grade the Agent's answer with an LLM and ask it to improve a weak answer before finishing."""

    def __init__(self, model: str = "gpt-5.4-nano"):
        self._judge = OpenAIChatGenerator(model=model)

    def warm_up(self) -> None:
        # Warm up the judge's own client during the Agent's warm-up.
        self._judge.warm_up()

    def run(self, state: State) -> None:
        answer = state.data["messages"][-1].text or ""
        verdict = self._judge.run(
            messages=[ChatMessage.from_user(f"Reply with only PASS or FAIL. Is this answer complete?\n\n{answer}")]
        )["replies"][0].text or ""
        if "FAIL" in verdict.upper():
            state.set("messages", [ChatMessage.from_user("Your answer was incomplete. Please improve it.")])


agent = Agent(chat_generator=OpenAIChatGenerator(model="gpt-5.4-nano"), hooks={"on_exit": [GradeFinalAnswer()]})
result = agent.run(messages=[ChatMessage.from_user("Explain how vaccines work.")])
print(result["last_message"].text)

How did you test it?

New tests

Notes for the reviewer

Checklist

• I have read the contributors guidelines and the code of conduct.
• I have updated the related issue with new insights and changes.
• I have added unit tests and updated the docstrings.
• I've used one of the conventional commit types for my PR title: fix:, feat:, build:, chore:, ci:, docs:, style:, refactor:, perf:, test: and added ! in case the PR includes breaking changes.
• I have documented my code.
• I have added a release note file, following the contributors guidelines.
• I have run pre-commit hooks and fixed any issue.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

1 Skipped Deployment

Project	Deployment	Actions	Updated (UTC)
haystack-docs	Ignored	Preview	Jun 25, 2026 2:43pm

[github-actions]

Coverage report

Click to see where and how coverage changed

File	Statements	Missing	Coverage	Coverage (new stmts)	Lines missing
haystack/components/agents
agent.py					149, 1199, 1249
haystack/hooks
__init__.py
from_function.py
invocation.py
protocol.py
utils.py
haystack/utils
callable_serialization.py					121
Project Total

This report was generated by python-coverage-comment-action

[anakin87]

Took a first look and left some comments
Overall, I like the implementation

[anakin87]

I see we refer to lifecycle events here and there. Since this might be confused with the warm_up/close topic (also in docs), could we find another option?
I am thinking of something like agent event, hook point, ..

[sjrl]

Done in 1c847d1

I've decided to go with hook point. Will mull over it over the weekend to see if I still like it or something else :D

[anakin87]

the exit part is the most difficult one and I am not sure I got it right (we might want to do something to simplify the code; not this function but the whole handling).

My impression is that this logic assumes that hooks append only, while they might also rewrite history.
One idea could be documenting this or implementing a different mechanism (a continue_run State key?).

[sjrl]

My impression is that this logic assumes that hooks append only, while they might also rewrite history.

Ahh right good catch. An explicit key would make it easier and more declarative. And I was already thinking about adding other flow control keys like this like a should_exit: bool if a hook (i.e. a BudgetHook) decides that the Agent should stop since a budget has been exceeded. Or if we were to add an after_llm hook adding something like a rerun_llm: bool key could be useful if the user wants to trigger the LLM again if the reply was unsatisfactory for some reason.

So I'd probably lean towards adding the continue_run key. Maybe nest it under a flow_control or loop_control State entry so it's easier for us to add more keys like this in the future. WDYT?

[anakin87]

Ok to add it if it makes the behavior simpler and more robust.

[anakin87]

@sjrl I was wondering about Tool result offloading (hooks can work for this?), then I found https://github.com/deepset-ai/haystack-private/issues/405#issuecomment-4797946676 🙂
Feel free to explore a bit if hooks would work well for this use case (can also be done later)

Also, since I'll be off next week, feel free to merge this PR once Julian approves it.