Building a Cost-Aware Agent with Hooks

_{Last Updated:
July 21, 2026}

🚀 Part of Haystack 3.0 Launch Week: five days of new drops (July 20–24).

Every call to Agent.run() returns metadata alongside the agent’s reply. In this cookbook you’ll use that metadata, and Agent hooks, to enforce soft and hard budget policies.

You’ll:

Build a simple agent with a custom tool and inspect step_count, token_usage, and tool_call_counts.
Implement a reusable post-run budget policy with soft and hard limits.
Enforce the same limits inside the agent loop with hooks, so an over-budget run can stop before the next LLM call.

Prerequisite: An OpenAI API key.

Setup

Install the latest Haystack version to be able to use Hooks.

!pip install -q haystack-ai>=3.0.0

from getpass import getpass
import os

if "OPENAI_API_KEY" not in os.environ:
    os.environ["OPENAI_API_KEY"] = getpass("Enter OpenAI API key: ")

Enter OpenAI API key: ··········

Part 1: A Simple Cost-Aware Agent

Start with a minimal agent that has a single tool. After the run, inspect the three metadata fields to see what the agent spent.

Field	What it tells you
`step_count`	How many agent steps (LLM + tool call pairs) were used
`token_usage`	Aggregated token counts across all LLM calls in the run
`tool_call_counts`	How many times each tool was invoked

Define a tool

Use the @tool decorator to turn any Python function into an agent-callable tool. The docstring and type annotations generate the tool schema the LLM sees.

from typing import Annotated
from haystack.tools import tool

@tool
def price_lookup(product: Annotated[str, "Product name to look up"]) -> str:
    """Look up a mock product price."""
    mock_prices = {"laptop": "$999", "keyboard": "$99", "monitor": "$349"}
    return mock_prices.get(product.lower(), "Unknown product")

Create and run the agent

Pass the tool to Agent together with a system_prompt that tells the LLM when to use it.

from haystack.components.agents import Agent
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.dataclasses import ChatMessage

agent = Agent(
    chat_generator=OpenAIChatGenerator(model="gpt-4o-mini"),
    tools=[price_lookup],
    system_prompt="Always call price_lookup before answering product price questions.",
)

result = agent.run(messages=[ChatMessage.from_user("What's the price of a laptop?")])
print(result["last_message"].text)

The price of a laptop is $999.

Inspect the run metadata

Agent.run() returns step_count, token_usage, and tool_call_counts alongside the reply:

steps = result["step_count"]
tokens = result["token_usage"]
tool_calls = result["tool_call_counts"]

total_tokens = tokens.get("total_tokens", 0)
prompt_tokens = tokens.get("prompt_tokens", 0)
completion_tokens = tokens.get("completion_tokens", 0)

print(f"Agent steps:        {steps}")
print(f"Total tokens:       {total_tokens}")
print(f"  Prompt tokens:    {prompt_tokens}")
print(f"  Completion tokens:{completion_tokens}")
print(f"Tool call counts:   {tool_calls}")

Agent steps:        2
Total tokens:       186
  Prompt tokens:    161
  Completion tokens:25
Tool call counts:   {'price_lookup': 1}

You can already act on these values directly

if total_tokens > 1200:
    print("⚠️  Budget warning: consider switching to a smaller model on the next run.")
if steps > 3:
    print("⚠️  Latency warning: tighten the system prompt or reduce tool hops.")

Define structured policy object and a utility function so you can reuse across endpoints.

from dataclasses import dataclass
from typing import Any

@dataclass
class BudgetPolicy:
    soft_token_limit: int = 1200  # log a warning but still return the response
    hard_token_limit: int = 2500  # block downstream processing
    max_steps: int = 6  # too many loops = poorly-scoped prompt

def evaluate_agent_budget(result: dict[str, Any], policy: BudgetPolicy) -> dict[str, Any]:
    """Return a structured budget decision for an agent run result."""
    usage = result.get("token_usage") or {}
    total_tokens = usage.get("total_tokens", 0)
    steps = result.get("step_count", 0)

    soft_exceeded = total_tokens > policy.soft_token_limit
    hard_exceeded = total_tokens > policy.hard_token_limit
    step_exceeded = steps > policy.max_steps

    action = "allow"
    if hard_exceeded or step_exceeded:
        action = "block"
    elif soft_exceeded:
        action = "warn"

    return {
        "action": action,
        "total_tokens": total_tokens,
        "steps": steps,
        "tool_call_counts": result.get("tool_call_counts", {}),
        "soft_exceeded": soft_exceeded,
        "hard_exceeded": hard_exceeded,
        "step_exceeded": step_exceeded,
    }

Evaluate the result from Part 1:

policy = BudgetPolicy(soft_token_limit=1200, hard_token_limit=2500, max_steps=6)
decision = evaluate_agent_budget(result, policy)

if decision["action"] == "block":
    print("🚫 Hard budget exceeded — request blocked.")
elif decision["action"] == "warn":
    print("⚠️  Soft budget exceeded — response returned with a warning.")
else:
    print("✅ Within budget.")

print(decision)

✅ Within budget.
{'action': 'allow', 'total_tokens': 186, 'steps': 2, 'tool_call_counts': {'price_lookup': 1}, 'soft_exceeded': False, 'hard_exceeded': False, 'step_exceeded': False}

Part 2: Enforcing Budgets Inside the Agent Loop with Hooks

Post-run evaluation catches over-budget runs after they finish. To stop a run before a costly next LLM call, use Agent hooks.

Hooks run at specific points in the agent loop:

Hook	When it runs
`before_run`	Once at the start of a run - ideal to set the initial state
`before_llm`	Before each LLM call - ideal for budget enforcement
`before_tool`	Before tool calls - ideal for HITL or data anonymization
`after_tool`	After tool results are written to state - ideal for modifying the tool results
`on_exit`	When the agent is about to stop on an exit condition (not when `max_agent_steps` is hit)
`after_run`	Once at the end of every run — ideal for final reporting

Each hook receives a State object with read/write access to the run’s accumulated metadata. Lower the soft and hard token limits to see how the agent behaves when the budget is exceeded.

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

policy = BudgetPolicy(soft_token_limit=100, hard_token_limit=500, max_steps=6)

@hook
def enforce_budget_before_llm(state: State) -> None:
    """Raise before the next LLM call if the hard budget is already exceeded."""
    decision = evaluate_agent_budget(state.data, policy)
    state.set("budget_decision", decision["action"])
    state.set("budget_snapshot", decision)
    if decision["action"] == "block":
        raise RuntimeError(
            f"Hard budget exceeded — stopping before next LLM call. "
            f"tokens={decision['total_tokens']}, steps={decision['steps']}"
        )
    elif decision["action"] == "warn":
        print("Soft budget exceeded - still continuing")
    else:
        print("Everything within the limit")

@hook
def finalize_budget_after_run(state: State) -> None:
    """Recompute the budget decision once at the end to capture the full run."""
    final = evaluate_agent_budget(state.data, policy)
    state.set("budget_snapshot", final)
    state.set("budget_decision", final["action"])

@hook
def mock_tool_hook(state: State) -> None:
    pending = state.data["messages"][-1].tool_calls or []
    for tool_call in pending:
        print(f"{tool_call.tool_name} tool is called with {tool_call.arguments} arguments!")

agent_with_hooks_tool = Agent(
    chat_generator=OpenAIChatGenerator(model="gpt-4o-mini"),
    tools=[price_lookup],
    system_prompt="Always call price_lookup before answering product price questions.",
    hooks={
        "before_llm": [enforce_budget_before_llm],
        "before_tool": [mock_tool_hook],
        "after_run": [finalize_budget_after_run],
    },
    state_schema={
        "budget_decision": {"type": str},
        "budget_snapshot": {"type": dict[str, Any]},
    },
)

Soft budget warning

This shorter query typically stays under the hard limit but crosses the soft limit (100 tokens), so the hook prints a warning and the run still completes:

try:
    result = agent_with_hooks_tool.run(
        messages=[ChatMessage.from_user("Compare prices for laptop, keyboard, and monitor.")]
    )

    print("Budget decision:", result["budget_decision"])
    print("Budget snapshot:", result["budget_snapshot"])
    print("\nAgent reply:\n", result["last_message"].text)
except RuntimeError as exc:
    print(exc)

Everything within the limit
price_lookup tool is called with {'product': 'laptop'} arguments!
price_lookup tool is called with {'product': 'keyboard'} arguments!
price_lookup tool is called with {'product': 'monitor'} arguments!
Soft budget exceeded - still continuing
Budget decision: warn
Budget snapshot: {'action': 'warn', 'total_tokens': 312, 'steps': 2, 'tool_call_counts': {'price_lookup': 3}, 'soft_exceeded': True, 'hard_exceeded': False, 'step_exceeded': False}

Agent reply:
 Here are the prices for the products you requested:

- Laptop: **$999**
- Keyboard: **$99**
- Monitor: **$349**

Hit the hard budget

The previous query only crossed the soft limit (100 tokens) and finished cleanly. To trigger a hard stop, send a much longer message so the first LLM call alone pushes token_usage past hard_token_limit (500). The agent still needs another LLM call after the tool results and before_llm raises before that call runs.

# Inflate the prompt so the first LLM call exceeds hard_token_limit=500.
long_context = (
    "Consider inventory constraints, regional pricing, warranty options, "
    "and shipping costs for this purchasing decision. "
) * 80

try:
    result = agent_with_hooks_tool.run(
        messages=[
            ChatMessage.from_user(
                long_context
                + "Now look up prices for laptop, keyboard, and monitor one by one "
                "(separate price_lookup calls for each), then compare them."
            )
        ]
    )

    print("Budget decision:", result["budget_decision"])
    print("Budget snapshot:", result["budget_snapshot"])
    print("\nAgent reply:\n", result["last_message"].text)
except RuntimeError as exc:
    print(exc)

Everything within the limit
price_lookup tool is called with {'product': 'laptop'} arguments!
price_lookup tool is called with {'product': 'keyboard'} arguments!
price_lookup tool is called with {'product': 'monitor'} arguments!
Hard budget exceeded — stopping before next LLM call. tokens=1587, steps=1

This pattern gives you both:

Live enforcement — the before_llm hook stops an over-budget run before another expensive LLM call.
Final reporting — the after_run hook ensures the output always includes an up-to-date budget_snapshot, including when the run ends because max_agent_steps was reached (unlike on_exit).

What’s next

You’ve seen how to make a Haystack Agent cost-aware using built-in metadata and hooks:

Read step_count, token_usage, and tool_call_counts from any agent run.
Apply a structured BudgetPolicy to allow, warn, or block a response.
Enforce budgets limits with before_llm, and report the final decision with after_run.

Related resources:

Build a Tool-Calling Agent — add web search and pipeline tools to your agent.
Agent Hooks — all hook points, ready-made hooks, and patterns for auditing or continuing a run.
Agent State — store custom info alongside budget data in a single run-scoped state object.

To stay up to date, subscribe to the newsletter or join Discord.