{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "1098497c",
      "metadata": {
        "id": "1098497c"
      },
      "source": [
        "# Building a Cost-Aware Agent with Hooks\n",
        "\n",
        "*Notebook by [Bilge Yücel](https://www.linkedin.com/in/bilge-yucel/)*\n",
        "\n",
        "> 🚀 **Part of [Haystack 3.0 Launch Week](https://haystack.deepset.ai/launch-week/haystack-3)**: five days of new drops (July 20–24).\n",
        "\n",
        "Every call to [`Agent.run()`](https://docs.haystack.deepset.ai/docs/agent) returns [metadata](https://docs.haystack.deepset.ai/docs/agent#run-metadata) alongside the agent's reply. In this cookbook you'll use that metadata, and [Agent hooks](https://docs.haystack.deepset.ai/docs/hooks), to enforce soft and hard budget policies.\n",
        "\n",
        "You'll:\n",
        "\n",
        "1. Build a simple agent with a custom tool and inspect `step_count`, `token_usage`, and `tool_call_counts`.\n",
        "2. Implement a reusable post-run budget policy with soft and hard limits.\n",
        "3. Enforce the same limits *inside* the agent loop with hooks, so an over-budget run can stop before the next LLM call.\n",
        "\n",
        "> **Prerequisite:** An [OpenAI API key](https://platform.openai.com/api-keys).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a487aa5f",
      "metadata": {
        "id": "a487aa5f"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Install the latest Haystack version to be able to use Hooks.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "3e1ef2e0",
      "metadata": {
        "id": "3e1ef2e0"
      },
      "outputs": [],
      "source": [
        "!pip install -q haystack-ai>=3.0.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "8be1d950",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8be1d950",
        "outputId": "a100ab76-346e-4c7a-dc07-6674196cd377"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Enter OpenAI API key: ··········\n"
          ]
        }
      ],
      "source": [
        "from getpass import getpass\n",
        "import os\n",
        "\n",
        "if \"OPENAI_API_KEY\" not in os.environ:\n",
        "    os.environ[\"OPENAI_API_KEY\"] = getpass(\"Enter OpenAI API key: \")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b508ab96",
      "metadata": {
        "id": "b508ab96"
      },
      "source": [
        "## Part 1: A Simple Cost-Aware Agent\n",
        "\n",
        "Start with a minimal agent that has a single tool. After the run, inspect the three metadata fields to see what the agent spent.\n",
        "\n",
        "| Field | What it tells you |\n",
        "|---|---|\n",
        "| `step_count` | How many agent steps (LLM + tool call pairs) were used |\n",
        "| `token_usage` | Aggregated token counts across all LLM calls in the run |\n",
        "| `tool_call_counts` | How many times each tool was invoked |\n",
        "\n",
        "### Define a tool\n",
        "\n",
        "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.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "id": "7de9b20c",
      "metadata": {
        "id": "7de9b20c"
      },
      "outputs": [],
      "source": [
        "from typing import Annotated\n",
        "from haystack.tools import tool\n",
        "\n",
        "@tool\n",
        "def price_lookup(product: Annotated[str, \"Product name to look up\"]) -> str:\n",
        "    \"\"\"Look up a mock product price.\"\"\"\n",
        "    mock_prices = {\"laptop\": \"$999\", \"keyboard\": \"$99\", \"monitor\": \"$349\"}\n",
        "    return mock_prices.get(product.lower(), \"Unknown product\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9fec2dfa",
      "metadata": {
        "id": "9fec2dfa"
      },
      "source": [
        "### Create and run the agent\n",
        "\n",
        "Pass the tool to `Agent` together with a `system_prompt` that tells the LLM when to use it.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "id": "372098a0",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "372098a0",
        "outputId": "3d2e6f97-7ff1-4105-bc6c-00b925075734"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The price of a laptop is $999.\n"
          ]
        }
      ],
      "source": [
        "from haystack.components.agents import Agent\n",
        "from haystack.components.generators.chat import OpenAIChatGenerator\n",
        "from haystack.dataclasses import ChatMessage\n",
        "\n",
        "agent = Agent(\n",
        "    chat_generator=OpenAIChatGenerator(model=\"gpt-4o-mini\"),\n",
        "    tools=[price_lookup],\n",
        "    system_prompt=\"Always call price_lookup before answering product price questions.\",\n",
        ")\n",
        "\n",
        "result = agent.run(messages=[ChatMessage.from_user(\"What's the price of a laptop?\")])\n",
        "print(result[\"last_message\"].text)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "88ae633f",
      "metadata": {
        "id": "88ae633f"
      },
      "source": [
        "### Inspect the run metadata\n",
        "\n",
        "`Agent.run()` returns `step_count`, `token_usage`, and `tool_call_counts` alongside the reply:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "id": "10b9d20d",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "10b9d20d",
        "outputId": "a168e903-ca53-4222-d70f-1c3452ede9ad"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Agent steps:        2\n",
            "Total tokens:       186\n",
            "  Prompt tokens:    161\n",
            "  Completion tokens:25\n",
            "Tool call counts:   {'price_lookup': 1}\n"
          ]
        }
      ],
      "source": [
        "steps = result[\"step_count\"]\n",
        "tokens = result[\"token_usage\"]\n",
        "tool_calls = result[\"tool_call_counts\"]\n",
        "\n",
        "total_tokens = tokens.get(\"total_tokens\", 0)\n",
        "prompt_tokens = tokens.get(\"prompt_tokens\", 0)\n",
        "completion_tokens = tokens.get(\"completion_tokens\", 0)\n",
        "\n",
        "print(f\"Agent steps:        {steps}\")\n",
        "print(f\"Total tokens:       {total_tokens}\")\n",
        "print(f\"  Prompt tokens:    {prompt_tokens}\")\n",
        "print(f\"  Completion tokens:{completion_tokens}\")\n",
        "print(f\"Tool call counts:   {tool_calls}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can already act on these values directly"
      ],
      "metadata": {
        "id": "k3c9SX2nUY2B"
      },
      "id": "k3c9SX2nUY2B"
    },
    {
      "cell_type": "code",
      "source": [
        "if total_tokens > 1200:\n",
        "    print(\"⚠️  Budget warning: consider switching to a smaller model on the next run.\")\n",
        "if steps > 3:\n",
        "    print(\"⚠️  Latency warning: tighten the system prompt or reduce tool hops.\")"
      ],
      "metadata": {
        "id": "CEvM6AHRUWwN"
      },
      "id": "CEvM6AHRUWwN",
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "13d3064e",
      "metadata": {
        "id": "13d3064e"
      },
      "source": [
        "Define structured policy object and a utility function so you can reuse across endpoints.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "id": "77516bea",
      "metadata": {
        "id": "77516bea"
      },
      "outputs": [],
      "source": [
        "from dataclasses import dataclass\n",
        "from typing import Any\n",
        "\n",
        "@dataclass\n",
        "class BudgetPolicy:\n",
        "    soft_token_limit: int = 1200  # log a warning but still return the response\n",
        "    hard_token_limit: int = 2500  # block downstream processing\n",
        "    max_steps: int = 6  # too many loops = poorly-scoped prompt\n",
        "\n",
        "def evaluate_agent_budget(result: dict[str, Any], policy: BudgetPolicy) -> dict[str, Any]:\n",
        "    \"\"\"Return a structured budget decision for an agent run result.\"\"\"\n",
        "    usage = result.get(\"token_usage\") or {}\n",
        "    total_tokens = usage.get(\"total_tokens\", 0)\n",
        "    steps = result.get(\"step_count\", 0)\n",
        "\n",
        "    soft_exceeded = total_tokens > policy.soft_token_limit\n",
        "    hard_exceeded = total_tokens > policy.hard_token_limit\n",
        "    step_exceeded = steps > policy.max_steps\n",
        "\n",
        "    action = \"allow\"\n",
        "    if hard_exceeded or step_exceeded:\n",
        "        action = \"block\"\n",
        "    elif soft_exceeded:\n",
        "        action = \"warn\"\n",
        "\n",
        "    return {\n",
        "        \"action\": action,\n",
        "        \"total_tokens\": total_tokens,\n",
        "        \"steps\": steps,\n",
        "        \"tool_call_counts\": result.get(\"tool_call_counts\", {}),\n",
        "        \"soft_exceeded\": soft_exceeded,\n",
        "        \"hard_exceeded\": hard_exceeded,\n",
        "        \"step_exceeded\": step_exceeded,\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2d8476e9",
      "metadata": {
        "id": "2d8476e9"
      },
      "source": [
        "Evaluate the result from Part 1:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "id": "a993a229",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a993a229",
        "outputId": "a841bbf8-113d-4f2d-cc76-50a855631ade"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "✅ Within budget.\n",
            "{'action': 'allow', 'total_tokens': 186, 'steps': 2, 'tool_call_counts': {'price_lookup': 1}, 'soft_exceeded': False, 'hard_exceeded': False, 'step_exceeded': False}\n"
          ]
        }
      ],
      "source": [
        "policy = BudgetPolicy(soft_token_limit=1200, hard_token_limit=2500, max_steps=6)\n",
        "decision = evaluate_agent_budget(result, policy)\n",
        "\n",
        "if decision[\"action\"] == \"block\":\n",
        "    print(\"🚫 Hard budget exceeded — request blocked.\")\n",
        "elif decision[\"action\"] == \"warn\":\n",
        "    print(\"⚠️  Soft budget exceeded — response returned with a warning.\")\n",
        "else:\n",
        "    print(\"✅ Within budget.\")\n",
        "\n",
        "print(decision)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9959e55f",
      "metadata": {
        "id": "9959e55f"
      },
      "source": [
        "## Part 2: Enforcing Budgets Inside the Agent Loop with Hooks\n",
        "\n",
        "Post-run evaluation catches over-budget runs *after* they finish. To stop a run *before* a costly next LLM call, use [Agent hooks](https://docs.haystack.deepset.ai/docs/hooks).\n",
        "\n",
        "Hooks run at specific points in the agent loop:\n",
        "\n",
        "| Hook | When it runs |\n",
        "|---|---|\n",
        "| `before_run` | Once at the start of a run - ideal to set the initial state |\n",
        "| `before_llm` | Before each LLM call - ideal for budget enforcement |\n",
        "| `before_tool` | Before tool calls - ideal for HITL or data anonymization|\n",
        "| `after_tool` | After tool results are written to state - ideal for modifying the tool results |\n",
        "| `on_exit` | When the agent is about to stop on an exit condition (not when `max_agent_steps` is hit) |\n",
        "| `after_run` | Once at the end of every run — ideal for final reporting |\n",
        "\n",
        "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.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "id": "3a7ae323",
      "metadata": {
        "id": "3a7ae323"
      },
      "outputs": [],
      "source": [
        "from haystack.components.agents import Agent\n",
        "from haystack.components.agents.state import State\n",
        "from haystack.components.generators.chat import OpenAIChatGenerator\n",
        "from haystack.dataclasses import ChatMessage\n",
        "from haystack.hooks import hook\n",
        "\n",
        "policy = BudgetPolicy(soft_token_limit=100, hard_token_limit=500, max_steps=6)\n",
        "\n",
        "@hook\n",
        "def enforce_budget_before_llm(state: State) -> None:\n",
        "    \"\"\"Raise before the next LLM call if the hard budget is already exceeded.\"\"\"\n",
        "    decision = evaluate_agent_budget(state.data, policy)\n",
        "    state.set(\"budget_decision\", decision[\"action\"])\n",
        "    state.set(\"budget_snapshot\", decision)\n",
        "    if decision[\"action\"] == \"block\":\n",
        "        raise RuntimeError(\n",
        "            f\"Hard budget exceeded — stopping before next LLM call. \"\n",
        "            f\"tokens={decision['total_tokens']}, steps={decision['steps']}\"\n",
        "        )\n",
        "    elif decision[\"action\"] == \"warn\":\n",
        "        print(\"Soft budget exceeded - still continuing\")\n",
        "    else:\n",
        "        print(\"Everything within the limit\")\n",
        "\n",
        "@hook\n",
        "def finalize_budget_after_run(state: State) -> None:\n",
        "    \"\"\"Recompute the budget decision once at the end to capture the full run.\"\"\"\n",
        "    final = evaluate_agent_budget(state.data, policy)\n",
        "    state.set(\"budget_snapshot\", final)\n",
        "    state.set(\"budget_decision\", final[\"action\"])\n",
        "\n",
        "@hook\n",
        "def mock_tool_hook(state: State) -> None:\n",
        "    pending = state.data[\"messages\"][-1].tool_calls or []\n",
        "    for tool_call in pending:\n",
        "        print(f\"{tool_call.tool_name} tool is called with {tool_call.arguments} arguments!\")\n",
        "\n",
        "agent_with_hooks_tool = Agent(\n",
        "    chat_generator=OpenAIChatGenerator(model=\"gpt-4o-mini\"),\n",
        "    tools=[price_lookup],\n",
        "    system_prompt=\"Always call price_lookup before answering product price questions.\",\n",
        "    hooks={\n",
        "        \"before_llm\": [enforce_budget_before_llm],\n",
        "        \"before_tool\": [mock_tool_hook],\n",
        "        \"after_run\": [finalize_budget_after_run],\n",
        "    },\n",
        "    state_schema={\n",
        "        \"budget_decision\": {\"type\": str},\n",
        "        \"budget_snapshot\": {\"type\": dict[str, Any]},\n",
        "    },\n",
        ")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5eda4374",
      "metadata": {
        "id": "5eda4374"
      },
      "source": [
        "### Soft budget warning\n",
        "\n",
        "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:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "id": "ee93affc",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ee93affc",
        "outputId": "6bad12fb-f794-4e47-f19f-c7adcac8619f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Everything within the limit\n",
            "price_lookup tool is called with {'product': 'laptop'} arguments!\n",
            "price_lookup tool is called with {'product': 'keyboard'} arguments!\n",
            "price_lookup tool is called with {'product': 'monitor'} arguments!\n",
            "Soft budget exceeded - still continuing\n",
            "Budget decision: warn\n",
            "Budget snapshot: {'action': 'warn', 'total_tokens': 312, 'steps': 2, 'tool_call_counts': {'price_lookup': 3}, 'soft_exceeded': True, 'hard_exceeded': False, 'step_exceeded': False}\n",
            "\n",
            "Agent reply:\n",
            " Here are the prices for the products you requested:\n",
            "\n",
            "- Laptop: **$999**\n",
            "- Keyboard: **$99**\n",
            "- Monitor: **$349**\n"
          ]
        }
      ],
      "source": [
        "try:\n",
        "    result = agent_with_hooks_tool.run(\n",
        "        messages=[ChatMessage.from_user(\"Compare prices for laptop, keyboard, and monitor.\")]\n",
        "    )\n",
        "\n",
        "    print(\"Budget decision:\", result[\"budget_decision\"])\n",
        "    print(\"Budget snapshot:\", result[\"budget_snapshot\"])\n",
        "    print(\"\\nAgent reply:\\n\", result[\"last_message\"].text)\n",
        "except RuntimeError as exc:\n",
        "    print(exc)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "35feedf0",
      "metadata": {
        "id": "35feedf0"
      },
      "source": [
        "### Hit the hard budget\n",
        "\n",
        "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.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "id": "e7d9648f",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e7d9648f",
        "outputId": "952dd188-166d-40d6-9e51-f8328863f9e2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Everything within the limit\n",
            "price_lookup tool is called with {'product': 'laptop'} arguments!\n",
            "price_lookup tool is called with {'product': 'keyboard'} arguments!\n",
            "price_lookup tool is called with {'product': 'monitor'} arguments!\n",
            "Hard budget exceeded — stopping before next LLM call. tokens=1587, steps=1\n"
          ]
        }
      ],
      "source": [
        "# Inflate the prompt so the first LLM call exceeds hard_token_limit=500.\n",
        "long_context = (\n",
        "    \"Consider inventory constraints, regional pricing, warranty options, \"\n",
        "    \"and shipping costs for this purchasing decision. \"\n",
        ") * 80\n",
        "\n",
        "try:\n",
        "    result = agent_with_hooks_tool.run(\n",
        "        messages=[\n",
        "            ChatMessage.from_user(\n",
        "                long_context\n",
        "                + \"Now look up prices for laptop, keyboard, and monitor one by one \"\n",
        "                \"(separate price_lookup calls for each), then compare them.\"\n",
        "            )\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    print(\"Budget decision:\", result[\"budget_decision\"])\n",
        "    print(\"Budget snapshot:\", result[\"budget_snapshot\"])\n",
        "    print(\"\\nAgent reply:\\n\", result[\"last_message\"].text)\n",
        "except RuntimeError as exc:\n",
        "    print(exc)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ESjvbhxuE56d"
      },
      "source": [
        "This pattern gives you both:\n",
        "\n",
        "- **Live enforcement** — the `before_llm` hook stops an over-budget run before another expensive LLM call.\n",
        "- **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`).\n"
      ],
      "id": "ESjvbhxuE56d"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7cYvAIM7E56d"
      },
      "source": [
        "## What's next\n",
        "\n",
        "You've seen how to make a Haystack Agent cost-aware using built-in metadata and hooks:\n",
        "\n",
        "- Read `step_count`, `token_usage`, and `tool_call_counts` from any agent run.\n",
        "- Apply a structured `BudgetPolicy` to allow, warn, or block a response.\n",
        "- Enforce budgets limits with `before_llm`, and report the final decision with `after_run`.\n",
        "\n",
        "Related resources:\n",
        "\n",
        "- [Build a Tool-Calling Agent](https://haystack.deepset.ai/tutorials/43_building_a_tool_calling_agent) — add web search and pipeline tools to your agent.\n",
        "- [Agent Hooks](https://docs.haystack.deepset.ai/docs/hooks) — all hook points, ready-made hooks, and patterns for auditing or continuing a run.\n",
        "- [Agent State](https://docs.haystack.deepset.ai/docs/state) — store custom info alongside budget data in a single run-scoped state object.\n",
        "\n",
        "To stay up to date, [subscribe to the newsletter](https://landing.deepset.ai/haystack-community-updates) or [join Discord](https://discord.gg/Dr63fr9NDS).\n"
      ],
      "id": "7cYvAIM7E56d"
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "pygments_lexer": "ipython3"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}