diff --git a/docs-website/docs/pipeline-components/agents-1/hooks.mdx b/docs-website/docs/pipeline-components/agents-1/hooks.mdx index 29a047fb81..d1d51fa3a1 100644 --- a/docs-website/docs/pipeline-components/agents-1/hooks.mdx +++ b/docs-website/docs/pipeline-components/agents-1/hooks.mdx @@ -44,6 +44,7 @@ Registering a hook under an unknown hook point raises a `ValueError` at construc The Agent manages a few state keys that hooks interact with. Like the run-metadata keys (`step_count`, `token_usage`, `tool_call_counts`), they are reserved — using any of them in your own `state_schema` raises a `ValueError`. See [State](./state.mdx#schema-definition) for the full list: - `continue_run`: Set by an `on_exit` hook to keep the Agent running. +- `stop_run`: Set by any hook to stop the run, read before each LLM call and used as the `exit_reason`. - `tools`: The tools available in the current step, for hooks to inspect. - `hook_context`: Request-scoped resources passed to `Agent.run(hook_context={...})` / `run_async(hook_context={...})`. Hooks read it with `state.data["hook_context"]` or `state.data.get("hook_context")` — use it for per-request resources such as a user ID, a WebSocket, or a database client. Avoid the plain `state.get("hook_context")` here: `State.get` returns a deep copy of the value, which often fails for the kinds of resources stored in this dict (such as a WebSocket or a database client). - `context_tokens`: An approximate count of the tokens currently in the context window, refreshed after each LLM call with that reply's prompt-plus-completion tokens (read it with `state.get("context_tokens")`). Unlike `token_usage`, which accumulates over the run, this is replaced on every call. It's `0` until the first reply that reports usage and doesn't count messages appended after the latest call. A `before_llm` hook can read it to trigger context compaction once it crosses a threshold. diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py index 7dc4dd583a..9ba87512bd 100644 --- a/haystack/components/agents/agent.py +++ b/haystack/components/agents/agent.py @@ -82,6 +82,7 @@ # Internal state keys the Agent manages for run control and hooks. Like run-metadata keys they are reserved and cannot # be redefined by users, but unlike them they are NOT exposed as Agent inputs or outputs (purely internal state): # - `continue_run`: set by an `on_exit` hook to keep the Agent running instead of stopping (re-read each exit attempt). +# - `stop_run`: set by a hook to stop the run, read before each LLM call and used as the `exit_reason`. # - `tools`: the flattened tools available in the current step, so a hook can inspect them (e.g. HITL confirmation). # - `hook_context`: per-run request-scoped resources passed to `run`/`run_async` for hooks to read. # - `context_tokens`: approximate current context-window size, refreshed after each LLM call, for hooks to read @@ -89,6 +90,7 @@ # exposed as an output because it is a best-effort snapshot; see `_record_context_tokens`. _INTERNAL_STATE_KEYS: dict[str, dict[str, Any]] = { "continue_run": {"type": bool, "handler": replace_values}, + "stop_run": {"type": str, "handler": replace_values}, "tools": {"type": list, "handler": replace_values}, "hook_context": {"type": dict[str, Any], "handler": replace_values}, "context_tokens": {"type": int, "handler": replace_values}, @@ -838,9 +840,10 @@ def run( `meta["usage"]`. - "tool_call_counts": Mapping of tool name to the number of times that tool was invoked. - "exit_reason": Why the Agent stopped, useful for routing the output downstream (e.g. with a - `ConditionalRouter`). One of: `"text"` (the model returned a reply with no tool calls), the name of + `ConditionalRouter`). This is `"text"` when the model returned a reply with no tool calls, the name of the tool that satisfied a tool exit condition (in which case `last_message` is that tool's result), - or `"max_agent_steps"` (the Agent hit `max_agent_steps` before meeting an exit condition). + `"max_agent_steps"` when the Agent hit `max_agent_steps`, or a custom reason a hook supplied through + the `stop_run` state key. - Any additional keys defined in the `state_schema`. """ agent_inputs = {"messages": messages, "streaming_callback": streaming_callback, **kwargs} @@ -922,9 +925,10 @@ async def run_async( `meta["usage"]`. - "tool_call_counts": Mapping of tool name to the number of times that tool was invoked. - "exit_reason": Why the Agent stopped, useful for routing the output downstream (e.g. with a - `ConditionalRouter`). One of: `"text"` (the model returned a reply with no tool calls), the name of + `ConditionalRouter`). This is `"text"` when the model returned a reply with no tool calls, the name of the tool that satisfied a tool exit condition (in which case `last_message` is that tool's result), - or `"max_agent_steps"` (the Agent hit `max_agent_steps` before meeting an exit condition). + `"max_agent_steps"` when the Agent hit `max_agent_steps`, or a custom reason a hook supplied through + the `stop_run` state key. - Any additional keys defined in the `state_schema`. """ agent_inputs = {"messages": messages, "streaming_callback": streaming_callback, **kwargs} @@ -978,6 +982,10 @@ def _run_step(self, exe_context: _ExecutionContext, agent_span: tracing.Span) -> exe_context.state.set("tools", current_tools, handler_override=replace_values) _run_hooks(hooks=self.hooks, hook_point=BEFORE_LLM, state=exe_context.state) + # A hook requested a stop: end the run at the step boundary, before spending another LLM call. + if (reason := exe_context.state.data.get("stop_run")) is not None: + exe_context.state.set("exit_reason", reason) + return False chat_generator_inputs = { "messages": exe_context.state.data["messages"], **exe_context.chat_generator_inputs, @@ -1041,6 +1049,10 @@ async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: trac exe_context.state.set("tools", current_tools, handler_override=replace_values) await _run_hooks_async(hooks=self.hooks, hook_point=BEFORE_LLM, state=exe_context.state) + # A hook requested a stop: end the run at the step boundary, before spending another LLM call. + if (reason := exe_context.state.data.get("stop_run")) is not None: + exe_context.state.set("exit_reason", reason) + return False chat_generator_inputs = { "messages": exe_context.state.data["messages"], **exe_context.chat_generator_inputs, diff --git a/haystack/hooks/budget/__init__.py b/haystack/hooks/budget/__init__.py new file mode 100644 index 0000000000..12c7426f39 --- /dev/null +++ b/haystack/hooks/budget/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import sys +from typing import TYPE_CHECKING + +from lazy_imports import LazyImporter + +_import_structure = {"hooks": ["TokenBudgetHook"]} + +if TYPE_CHECKING: + from .hooks import TokenBudgetHook as TokenBudgetHook +else: + sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure) diff --git a/haystack/hooks/budget/hooks.py b/haystack/hooks/budget/hooks.py new file mode 100644 index 0000000000..8a5a3225d1 --- /dev/null +++ b/haystack/hooks/budget/hooks.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack import logging +from haystack.components.agents.state.state import State +from haystack.components.agents.utils import _INPUT_TOKEN_KEYS, _OUTPUT_TOKEN_KEYS, _first_numeric +from haystack.core.serialization import default_from_dict, default_to_dict +from haystack.dataclasses import ChatMessage +from haystack.utils.experimental import _experimental + +logger = logging.getLogger(__name__) + + +_FINAL_MESSAGE_TEXT = "The Agent stopped because the token budget was exceeded." + + +@_experimental +class TokenBudgetHook: + """ + Stop an Agent run when its token usage reaches a configured budget. + + The hook runs at the `before_llm` hook point and checks the cumulative token usage recorded in the Agent state. + When the budget is reached, the run ends before the next LLM call with the exit reason `"token_budget_exceeded"`. + + Only calls made by the Agent's chat generator contribute to `token_usage`; calls made by tools or other hooks are + not included. + + + ```python + from haystack.components.agents import Agent + from haystack.components.generators.chat import OpenAIChatGenerator + from haystack.hooks.budget import TokenBudgetHook + + agent = Agent( + chat_generator=OpenAIChatGenerator(), + tools=[web_search], + hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100_000)]}, + ) + + result = agent.run(messages=[...]) + ``` + """ + + allowed_hook_points = ("before_llm",) + + def __init__(self, *, max_total_tokens: int, add_final_message: bool = False) -> None: + """ + Create a token budget hook. + + :param max_total_tokens: Maximum cumulative token usage before the Agent is stopped. + :param add_final_message: Whether to append an assistant message explaining why the Agent stopped. + :raises ValueError: If `max_total_tokens` is less than 1. + """ + if max_total_tokens < 1: + raise ValueError(f"`max_total_tokens` must be a positive number of tokens, got {max_total_tokens}.") + self.max_total_tokens = max_total_tokens + self.add_final_message = add_final_message + + def run(self, state: State) -> None: + """ + Stop the Agent if its cumulative token usage has reached the budget. + + :param state: Agent state containing the cumulative token usage. + """ + usage = state.data.get("token_usage") or {} + # Not every chat generator reports `total_tokens`, so fall back to summing the input and output keys across + # the known naming conventions. + total_tokens = _first_numeric(usage, ("total_tokens",)) + if not total_tokens: + total_tokens = _first_numeric(usage, _INPUT_TOKEN_KEYS) + _first_numeric(usage, _OUTPUT_TOKEN_KEYS) + if total_tokens >= self.max_total_tokens: + logger.warning( + "Agent reached its token budget of {max_total_tokens} ({total_tokens} used); requesting a stop.", + max_total_tokens=self.max_total_tokens, + total_tokens=total_tokens, + ) + state.set("stop_run", "token_budget_exceeded") + if self.add_final_message: + state.set("messages", [ChatMessage.from_assistant(_FINAL_MESSAGE_TEXT)]) + + def to_dict(self) -> dict[str, Any]: + """ + Serialize this hook to a dictionary. + + :returns: Serialized representation of the hook. + """ + return default_to_dict(self, max_total_tokens=self.max_total_tokens, add_final_message=self.add_final_message) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "TokenBudgetHook": + """ + Create a hook from its serialized representation. + + :param data: Serialized hook data. + :returns: The deserialized hook. + """ + return default_from_dict(cls, data=data) diff --git a/pydoc/hooks_api.yml b/pydoc/hooks_api.yml index fef10aad17..9f35e24193 100644 --- a/pydoc/hooks_api.yml +++ b/pydoc/hooks_api.yml @@ -1,6 +1,6 @@ loaders: - search_path: [../haystack/hooks] - modules: ["protocol", "from_function", "compaction/hooks", "compaction/sliding_window", + modules: ["protocol", "from_function", "budget/hooks", "compaction/hooks", "compaction/sliding_window", "compaction/summarization", "compaction/tool_result_pruning", "compaction/types/protocol", "human_in_the_loop/dataclasses", "human_in_the_loop/hooks", "human_in_the_loop/policies", "human_in_the_loop/strategies", "human_in_the_loop/user_interfaces", diff --git a/releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml b/releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml new file mode 100644 index 0000000000..6dc0ef44b0 --- /dev/null +++ b/releasenotes/notes/agent-stop-run-token-budget-hook-dafd7694063c5791.yaml @@ -0,0 +1,19 @@ +--- +features: + - | + Added ``stop_run``, a control flag that lets an Agent hook end a run cleanly instead of raising an error. The + hook sets it to a reason of your choice (``state.set("stop_run", "my_reason")``), and the Agent stops without + another model call and reports that reason in ``exit_reason``. + - | + Added ``TokenBudgetHook`` (experimental), a ready-made hook that caps how many tokens an Agent run may spend. + The run ends with ``exit_reason`` set to ``"token_budget_exceeded"``, keeping the messages collected so far. + + .. code-block:: python + + from haystack.hooks.budget import TokenBudgetHook + + agent = Agent( + chat_generator=OpenAIChatGenerator(model="gpt-5.4-nano"), + tools=[search], + hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100_000)]}, + ) diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py index 568c0486c0..7c368b0d6c 100644 --- a/test/components/agents/test_agent.py +++ b/test/components/agents/test_agent.py @@ -221,6 +221,7 @@ def test_state_schema_resolution(self, weather_tool): "tool_call_counts": {"type": dict[str, int], "handler": replace_values}, "exit_reason": {"type": str, "handler": replace_values}, "continue_run": {"type": bool, "handler": replace_values}, + "stop_run": {"type": str, "handler": replace_values}, "tools": {"type": list, "handler": replace_values}, "hook_context": {"type": dict[str, Any], "handler": replace_values}, "context_tokens": {"type": int, "handler": replace_values}, @@ -248,7 +249,7 @@ def test_output_types(self, weather_tool, component_tool, monkeypatch): assert internal_key not in agent.__haystack_output__._sockets_dict def test_reserved_state_schema_keys_raise(self, weather_tool): - for reserved in ("step_count", "token_usage", "context_tokens", "tool_call_counts", "exit_reason"): + for reserved in ("step_count", "token_usage", "context_tokens", "tool_call_counts", "exit_reason", "stop_run"): with pytest.raises(ValueError, match="reserved for Agent internal state"): Agent( chat_generator=MockChatGenerator("Hello"), @@ -1227,6 +1228,7 @@ def test_tracing_span_run(self, spying_tracer, weather_tool): "type": "bool", "handler": "haystack.components.agents.state.state_utils.replace_values", }, + "stop_run": {"type": "str", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "tools": {"type": "list", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "hook_context": { "type": "dict[str, typing.Any]", diff --git a/test/components/agents/test_agent_hooks.py b/test/components/agents/test_agent_hooks.py index 1ff08dd506..2d2037a8b4 100644 --- a/test/components/agents/test_agent_hooks.py +++ b/test/components/agents/test_agent_hooks.py @@ -129,6 +129,11 @@ def always_continue(state: State) -> None: state.set("continue_run", True) +@hook +def stop_for_budget(state: State) -> None: + state.set("stop_run", "budget_exceeded") + + @hook def critique(state: State) -> None: # Push back on the first final answer to force one more loop. @@ -527,6 +532,120 @@ def audit(state: State) -> None: assert agent.chat_generator.run.call_count == 1 +class TestStopRun: + def test_before_llm_stop_skips_the_current_call(self): + agent = _agent(MockChatGenerator(), tools=[save], hooks={"before_llm": [stop_for_budget]}) + 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 == 0 + assert result["step_count"] == 0 + assert result["tool_call_counts"]["save"] == 0 + assert result["exit_reason"] == "budget_exceeded" + + @pytest.mark.parametrize("hook_point", ["before_tool", "after_tool"]) + def test_mid_step_stop_skips_the_next_call(self, hook_point): + agent = _agent(MockChatGenerator(), tools=[save], hooks={hook_point: [stop_for_budget]}) + 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 == 1 + assert result["step_count"] == 1 + assert result["tool_call_counts"]["save"] == 1 + assert result["last_message"].tool_call_result is not None + assert result["exit_reason"] == "budget_exceeded" + assert "stop_run" not in result + + def test_after_run_stop_is_ignored(self): + agent = _agent(MockChatGenerator(), hooks={"after_run": [stop_for_budget]}) + agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]}) + result = agent.run(messages=[ChatMessage.from_user("hi")]) + assert result["exit_reason"] == "text" + + @pytest.mark.parametrize("hook_point", ["after_tool", "on_exit"]) + def test_exit_reason_is_the_natural_exit_in_the_same_step(self, hook_point): + agent = _agent( + MockChatGenerator(), + tools=[final_answer], + exit_conditions=["final_answer"], + hooks={hook_point: [stop_for_budget]}, + ) + agent.chat_generator.run = MagicMock( + return_value={ + "replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("final_answer", {"answer": "a"})])] + } + ) + result = agent.run(messages=[ChatMessage.from_user("q")]) + assert result["exit_reason"] == "final_answer" + + def test_exit_reason_is_the_stop_when_continue_run_cancels_the_exit(self): + agent = _agent( + MockChatGenerator(), + tools=[final_answer], + exit_conditions=["final_answer"], + hooks={"after_tool": [stop_for_budget], "on_exit": [always_continue]}, + ) + agent.chat_generator.run = MagicMock( + return_value={ + "replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("final_answer", {"answer": "a"})])] + } + ) + result = agent.run(messages=[ChatMessage.from_user("q")]) + assert agent.chat_generator.run.call_count == 1 + assert result["exit_reason"] == "budget_exceeded" + assert result["last_message"].text == "keep going" + + def test_exit_reason_is_max_steps_on_the_last_step(self): + agent = _agent(MockChatGenerator(), tools=[save], max_agent_steps=1, hooks={"after_tool": [stop_for_budget]}) + 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 result["exit_reason"] == "max_agent_steps" + + def test_later_before_llm_hooks_still_run(self): + fired = [] + + def record(state: State) -> None: + fired.append(1) + + agent = _agent(MockChatGenerator(), hooks={"before_llm": [stop_for_budget, hook(record)]}) + agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]}) + result = agent.run(messages=[ChatMessage.from_user("hi")]) + assert fired == [1] + assert agent.chat_generator.run.call_count == 0 + assert result["exit_reason"] == "budget_exceeded" + + def test_after_run_hooks_still_run(self): + agent = _agent( + MockChatGenerator(), tools=[save], hooks={"after_tool": [stop_for_budget], "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 result["exit_reason"] == "budget_exceeded" + assert result["last_message"].text == "REPORT" + + def test_on_exit_hooks_do_not_run(self): + fired = [] + + def record(state: State) -> None: + fired.append(1) + + agent = _agent( + MockChatGenerator(), tools=[save], hooks={"after_tool": [stop_for_budget], "on_exit": [hook(record)]} + ) + 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 result["exit_reason"] == "budget_exceeded" + assert fired == [] + + class TestHookReuseAcrossHookPoints: def test_same_hook_under_two_hook_points(self): counter = [] @@ -578,6 +697,17 @@ def record(state: State) -> None: await agent.run_async(messages=[ChatMessage.from_user("hi")]) assert fired == [1] + @pytest.mark.asyncio + async def test_mid_step_stop_skips_the_next_call(self): + agent = _agent(MockChatGenerator(), tools=[save], hooks={"after_tool": [stop_for_budget]}) + 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.await_count == 1 + assert result["exit_reason"] == "budget_exceeded" + assert result["last_message"].tool_call_result is not None + @pytest.mark.asyncio async def test_after_tool_hook_rewrites_results(self): agent = _agent(MockChatGenerator(), tools=[save], hooks={"after_tool": [offload_tool_results]}) diff --git a/test/hooks/budget/__init__.py b/test/hooks/budget/__init__.py new file mode 100644 index 0000000000..c1764a6e03 --- /dev/null +++ b/test/hooks/budget/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/test/hooks/budget/test_hooks.py b/test/hooks/budget/test_hooks.py new file mode 100644 index 0000000000..473da1bbe4 --- /dev/null +++ b/test/hooks/budget/test_hooks.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import logging +from typing import Annotated, Any +from unittest.mock import MagicMock + +import pytest + +from haystack.components.agents import Agent +from haystack.components.agents.state import State +from haystack.components.generators.chat import MockChatGenerator +from haystack.dataclasses import ChatMessage, ChatRole, ToolCall +from haystack.hooks.budget import TokenBudgetHook +from haystack.hooks.budget.hooks import _FINAL_MESSAGE_TEXT +from haystack.tools import tool + +pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") + + +@tool +def fetch(topic: Annotated[str, "the topic to fetch"]) -> str: + """Fetch a document about a topic.""" + return "DATA" + + +def _fetch_reply(total_tokens: int) -> dict: + message = ChatMessage.from_assistant( + tool_calls=[ToolCall("fetch", {"topic": "x"})], meta={"usage": {"total_tokens": total_tokens}} + ) + return {"replies": [message]} + + +def _state(usage: dict) -> State: + schema = { + "token_usage": {"type": dict[str, Any]}, + "stop_run": {"type": str}, + "messages": {"type": list[ChatMessage]}, + } + return State(schema=schema, data={"token_usage": usage}) + + +class TestTokenBudgetHook: + @pytest.mark.parametrize( + "usage", + [ + {"total_tokens": 100}, + {"prompt_tokens": 60, "completion_tokens": 40}, + {"input_tokens": 60, "output_tokens": 40}, + ], + ids=["total_tokens", "openai-style", "anthropic-style"], + ) + def test_stops_when_usage_reaches_the_budget(self, usage, caplog): + state = _state(usage) + with caplog.at_level(logging.WARNING): + TokenBudgetHook(max_total_tokens=100).run(state) + assert state.data["stop_run"] == "token_budget_exceeded" + assert state.data.get("messages") is None + assert "token budget of 100 (100 used)" in caplog.text + + @pytest.mark.parametrize("usage", [{"total_tokens": 99}, {}], ids=["under-budget", "no-usage-reported"]) + def test_does_not_stop_below_the_budget(self, usage): + state = _state(usage) + TokenBudgetHook(max_total_tokens=100).run(state) + assert state.data.get("stop_run") is None + + def test_adds_a_final_message(self): + state = _state({"total_tokens": 100}) + TokenBudgetHook(max_total_tokens=100, add_final_message=True).run(state) + assert state.data["messages"][-1].text == _FINAL_MESSAGE_TEXT + assert state.data["messages"][-1].is_from(ChatRole.ASSISTANT) + + def test_non_positive_budget_raises(self): + with pytest.raises(ValueError, match="max_total_tokens"): + TokenBudgetHook(max_total_tokens=0) + + def test_to_dict_from_dict_roundtrip(self): + hook = TokenBudgetHook(max_total_tokens=5000, add_final_message=True) + restored = TokenBudgetHook.from_dict(hook.to_dict()) + assert restored.max_total_tokens == 5000 + assert restored.add_final_message is True + + def test_stops_an_agent_run_when_the_budget_is_spent(self): + agent = Agent( + chat_generator=MockChatGenerator(), + tools=[fetch], + hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100)]}, + ) + agent.warm_up() + agent.chat_generator.run = MagicMock( + side_effect=[_fetch_reply(60), _fetch_reply(60), {"replies": [ChatMessage.from_assistant("done")]}] + ) + result = agent.run(messages=[ChatMessage.from_user("hi")]) + assert agent.chat_generator.run.call_count == 2 + assert result["tool_call_counts"]["fetch"] == 2 + assert result["exit_reason"] == "token_budget_exceeded" + + def test_stops_a_text_only_loop_kept_alive_by_continue_run(self): + class KeepIterating: + def run(self, state: State) -> None: + state.set("continue_run", True) + + agent = Agent( + chat_generator=MockChatGenerator(), + hooks={"before_llm": [TokenBudgetHook(max_total_tokens=100)], "on_exit": [KeepIterating()]}, + ) + agent.warm_up() + agent.chat_generator.run = MagicMock( + return_value={"replies": [ChatMessage.from_assistant("draft", meta={"usage": {"total_tokens": 60}})]} + ) + result = agent.run(messages=[ChatMessage.from_user("hi")]) + assert agent.chat_generator.run.call_count == 2 + assert result["exit_reason"] == "token_budget_exceeded"