feat: add context_tokens to internal Agents State (#12102)

Author: sjrl

docs-website/docs/pipeline-components/agents-1/hooks.mdx

@@ -46,6 +46,7 @@ The Agent manages a few state keys that hooks interact with. Like the run-metada
 - `continue_run`: Set by an `on_exit` hook to keep the Agent running.
 - `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.
 
 Hooks can also read the automatically tracked run metadata: `step_count`, `token_usage`, and `tool_call_counts`.
 

docs-website/docs/pipeline-components/agents-1/state.mdx

@@ -71,7 +71,7 @@ If you don't specify a handler, State automatically assigns a default based on t
 The `Agent` manages some state keys itself and rejects them in a user-provided `state_schema` with a `ValueError`:
 
 - The run-metadata keys `step_count`, `token_usage`, `tool_call_counts`, and `exit_reason`, which the Agent populates automatically during a run: tools and hooks can read them from the live `State`, and they are returned in the result dictionary. `exit_reason` reports why the Agent stopped (`"text"`, the name of the tool that triggered a tool exit condition, or `"max_agent_steps"`).
-- The hook-facing keys `continue_run` (set by an `on_exit` hook to keep the Agent running), `tools` (the tools available in the current step, for hooks to inspect), and `hook_context` (the request-scoped resources passed to `Agent.run(hook_context={...})`).
+- The hook-facing keys `continue_run` (set by an `on_exit` hook to keep the Agent running), `tools` (the tools available in the current step, for hooks to inspect), `hook_context` (the request-scoped resources passed to `Agent.run(hook_context={...})`), and `context_tokens` (an approximate count of the tokens currently in the context window, refreshed after each LLM call for hooks to read — for example, to trigger context compaction). Unlike the run-metadata keys, these are not returned in the result dictionary.
 
 If one of your state keys clashes, rename it (for example, `my_token_usage`).
 :::

haystack/components/agents/agent.py

@@ -18,6 +18,7 @@
 )
 from haystack.components.agents.state.state_utils import merge_lists
 from haystack.components.agents.tool_calling import _run_tool, _run_tool_async
+from haystack.components.agents.utils import _record_context_tokens
 from haystack.components.builders import ChatPromptBuilder
 from haystack.components.generators.chat.types import ChatGenerator
 from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict
@@ -82,10 +83,14 @@
 # - `continue_run`: set by an `on_exit` hook to keep the Agent running instead of stopping (re-read each exit attempt).
 # - `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
+#   (e.g. a `before_llm` hook that triggers compaction once the context grows too large). Kept internal rather than
+#   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},
     "tools": {"type": list, "handler": replace_values},
     "hook_context": {"type": dict[str, Any], "handler": replace_values},
+    "context_tokens": {"type": int, "handler": replace_values},
 }
 
 
@@ -938,6 +943,7 @@ def _initialize_fresh_execution(
         state.set("messages", messages)
         state.set("step_count", 0)
         state.set("token_usage", {})
+        state.set("context_tokens", 0)
         state.set("tool_call_counts", {tool.name: 0 for tool in flat_tools})
         state.set("exit_reason", None)
         state.set("continue_run", False)
@@ -1187,6 +1193,7 @@ def _run_step(self, exe_context: _ExecutionContext, agent_span: tracing.Span) ->
             llm_messages = result["replies"]
             exe_context.state.set("messages", llm_messages)
             _record_llm_usage(exe_context.state, llm_messages)
+            _record_context_tokens(exe_context.state, llm_messages)
 
             # Stop on the "no tool call" exit: no tools available, or a plain assistant text reply (see _is_text_exit).
             if not current_tools or _is_text_exit(llm_messages):
@@ -1251,6 +1258,7 @@ async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: trac
             llm_messages = result["replies"]
             exe_context.state.set("messages", llm_messages)
             _record_llm_usage(exe_context.state, llm_messages)
+            _record_context_tokens(exe_context.state, llm_messages)
 
             # Stop on the "no tool call" exit: no tools available, or a plain assistant text reply (see _is_text_exit).
             if not current_tools or _is_text_exit(llm_messages):

haystack/components/agents/utils.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+from haystack.components.agents.state.state import State
+from haystack.dataclasses import ChatMessage
+
+# Input/output token key conventions across chat generators: most report OpenAI-style
+# `prompt_tokens`/`completion_tokens`; OpenAIResponsesChatGenerator reports `input_tokens`/`output_tokens`.
+_INPUT_TOKEN_KEYS = ("prompt_tokens", "input_tokens")
+_OUTPUT_TOKEN_KEYS = ("completion_tokens", "output_tokens")
+
+
+def _first_numeric(usage: dict[str, Any], keys: tuple[str, ...]) -> int:
+    """
+    Return the first numeric value found under `keys` in `usage`, or 0 if none is present.
+
+    :param usage: A ChatMessage `meta["usage"]` payload.
+    :param keys: Candidate keys to check, in priority order.
+    :returns: The first `int`/`float` value (as an `int`), or 0. bool values are skipped (not token counts).
+    """
+    for key in keys:
+        value = usage.get(key)
+        # bool is an int subclass, so exclude it explicitly: True/False is not a token count.
+        if isinstance(value, bool):
+            continue
+        if isinstance(value, (int, float)):
+            return int(value)
+    return 0
+
+
+def _context_tokens_from_usage(usage: dict[str, Any]) -> int:
+    """
+    Sum the input and output tokens reported in a single `meta["usage"]` dict.
+
+    :param usage: A ChatMessage `meta["usage"]` payload.
+    :returns: Input plus output tokens, or 0 if neither key convention is present.
+    """
+    return _first_numeric(usage, _INPUT_TOKEN_KEYS) + _first_numeric(usage, _OUTPUT_TOKEN_KEYS)
+
+
+def _record_context_tokens(state: State, llm_messages: list[ChatMessage]) -> None:
+    """
+    Store the approximate current context-window token count from the latest LLM call.
+
+    A chat-generator call returns a single reply, so only the last message is inspected. Unlike
+    `token_usage`, which accumulates across the run, this value is replaced each call with that reply's
+    prompt-plus-completion tokens. Only writes when usage is reported, so generators that don't surface
+    usage leave the previous value untouched.
+
+    :param state: The Agent's State, used to write the latest `context_tokens` count.
+    :param llm_messages: The ChatMessage objects returned from the latest LLM call.
+    """
+    if not llm_messages:
+        return
+    usage = llm_messages[-1].meta.get("usage")
+    if isinstance(usage, dict):
+        tokens = _context_tokens_from_usage(usage)
+        if tokens:
+            state.set("context_tokens", tokens)

releasenotes/notes/add-agent-context-tokens-831273f057452c71.yaml

@@ -0,0 +1,11 @@
+---
+enhancements:
+  - |
+    The ``Agent`` now tracks an approximate current context-window size in its internal ``State`` under
+    ``context_tokens``, refreshed after every LLM call with that reply's prompt-plus-completion tokens
+    (normalized across the ``prompt_tokens``/``completion_tokens`` and ``input_tokens``/``output_tokens``
+    key conventions). Unlike ``token_usage``, which accumulates across the whole run, ``context_tokens``
+    is replaced each call. Hooks can read it via ``state.get("context_tokens")`` — for example, a
+    ``before_llm`` hook that triggers context compaction once the value crosses a threshold. It is a
+    best-effort snapshot: it is ``0`` when the generator does not report usage, and does not count
+    messages appended after the latest call until the next call refreshes it.

test/components/agents/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0