diff --git a/docs-website/docs/pipeline-components/agents-1/hooks.mdx b/docs-website/docs/pipeline-components/agents-1/hooks.mdx index 0550490776..918674ac40 100644 --- a/docs-website/docs/pipeline-components/agents-1/hooks.mdx +++ b/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`. diff --git a/docs-website/docs/pipeline-components/agents-1/state.mdx b/docs-website/docs/pipeline-components/agents-1/state.mdx index 191b337f03..ad82d9ed07 100644 --- a/docs-website/docs/pipeline-components/agents-1/state.mdx +++ b/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`). ::: diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py index 8245f5e5a9..4361b89f8b 100644 --- a/haystack/components/agents/agent.py +++ b/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): diff --git a/haystack/components/agents/utils.py b/haystack/components/agents/utils.py new file mode 100644 index 0000000000..1542ae26d7 --- /dev/null +++ b/haystack/components/agents/utils.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# 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) diff --git a/releasenotes/notes/add-agent-context-tokens-831273f057452c71.yaml b/releasenotes/notes/add-agent-context-tokens-831273f057452c71.yaml new file mode 100644 index 0000000000..c57433e30c --- /dev/null +++ b/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. diff --git a/test/components/agents/__init__.py b/test/components/agents/__init__.py new file mode 100644 index 0000000000..c1764a6e03 --- /dev/null +++ b/test/components/agents/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py index 7665d5fdf8..afca9ca070 100644 --- a/test/components/agents/test_agent.py +++ b/test/components/agents/test_agent.py @@ -225,10 +225,14 @@ def test_output_types(self, weather_tool, component_tool, monkeypatch): "tool_call_counts": OutputSocket(name="tool_call_counts", type=dict[str, int], receivers=[]), "exit_reason": OutputSocket(name="exit_reason", type=str, receivers=[]), } - # Check that the internal-state keys are not set up as input sockets + # Check that the run-metadata keys are not set up as input sockets assert {"step_count", "token_usage", "tool_call_counts", "exit_reason"}.isdisjoint( agent.__haystack_input__._sockets_dict.keys() ) + # Internal-only state keys (those that are not also run parameters) are exposed as neither inputs nor outputs. + for internal_key in ("continue_run", "context_tokens"): + assert internal_key not in agent.__haystack_input__._sockets_dict + assert internal_key not in agent.__haystack_output__._sockets_dict def test_to_dict(self, weather_tool, component_tool, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "fake-key") @@ -276,7 +280,7 @@ def test_to_dict(self, weather_tool, component_tool, monkeypatch): "properties": {"location": {"type": "string"}}, "required": ["location"], }, - "function": "test_agent.weather_function", + "function": "agents.test_agent.weather_function", "async_function": None, "outputs_to_string": None, "inputs_from_state": None, @@ -360,7 +364,7 @@ def test_to_dict_with_toolset(self, monkeypatch, weather_tool): "properties": {"location": {"type": "string"}}, "required": ["location"], }, - "function": "test_agent.weather_function", + "function": "agents.test_agent.weather_function", "async_function": None, "outputs_to_string": None, "inputs_from_state": None, @@ -431,7 +435,7 @@ def test_from_dict(self, monkeypatch): "properties": {"location": {"type": "string"}}, "required": ["location"], }, - "function": "test_agent.weather_function", + "function": "agents.test_agent.weather_function", "async_function": None, "outputs_to_string": None, "inputs_from_state": None, @@ -487,6 +491,7 @@ def test_from_dict(self, monkeypatch): "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}, } assert agent.tool_concurrency_limit == 5 assert agent.tool_streaming_callback_passthrough is True @@ -526,7 +531,7 @@ def test_from_dict_with_toolset(self, monkeypatch): "properties": {"location": {"type": "string"}}, "required": ["location"], }, - "function": "test_agent.weather_function", + "function": "agents.test_agent.weather_function", "async_function": None, "outputs_to_string": None, "inputs_from_state": None, @@ -598,6 +603,7 @@ def test_from_dict_state_schema_none(self, monkeypatch): "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}, } def test_serde(self, weather_tool, component_tool, monkeypatch): @@ -619,7 +625,7 @@ def test_serde(self, weather_tool, component_tool, monkeypatch): init_parameters["chat_generator"]["type"] == "haystack.components.generators.chat.openai.OpenAIChatGenerator" ) - assert init_parameters["streaming_callback"] == "test_agent.sync_streaming_callback" + assert init_parameters["streaming_callback"] == "agents.test_agent.sync_streaming_callback" assert init_parameters["tools"][0]["data"]["function"] == serialize_callable(weather_function) assert ( init_parameters["tools"][1]["data"]["component"]["type"] @@ -643,6 +649,7 @@ def test_serde(self, weather_tool, component_tool, monkeypatch): "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}, } assert deserialized_agent.streaming_callback is sync_streaming_callback @@ -1162,7 +1169,7 @@ async def test_run_async_with_sync_streaming_callback_warns(self, weather_tool, def test_reserved_state_schema_keys_raise(self, monkeypatch, weather_tool): monkeypatch.setenv("OPENAI_API_KEY", "fake-key") - for reserved in ("step_count", "token_usage", "tool_call_counts", "exit_reason"): + for reserved in ("step_count", "token_usage", "context_tokens", "tool_call_counts", "exit_reason"): with pytest.raises(ValueError, match="reserved for Agent internal state"): Agent( chat_generator=OpenAIChatGenerator(), tools=[weather_tool], state_schema={reserved: {"type": int}} @@ -1248,6 +1255,64 @@ def test_metadata_outputs_show_defaults_when_no_data(self, weather_tool): assert result["tool_call_counts"] == {"weather_tool": 0} +class TestAgentContextTokens: + """The Agent refreshes the internal `context_tokens` after each LLM call so a hook can read the current + context-window size (e.g. to trigger compaction). It is a per-call snapshot, not accumulated.""" + + def test_before_llm_hook_reads_refreshed_context_tokens(self, weather_tool): + # Step 1: a tool call (prompt 10 + completion 5 = 15). Step 2: the final text answer. + first_step = [ + _assistant_with_usage( + tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})], + usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + ) + ] + second_step = [_assistant_with_usage("Done.", usage={"prompt_tokens": 20, "completion_tokens": 8})] + + seen: list[int] = [] + + @hook + def capture(state: State) -> None: + seen.append(state.get("context_tokens")) + + agent = Agent( + chat_generator=MockChatGenerator(first_step + second_step), + tools=[weather_tool], + hooks={"before_llm": [capture]}, + ) + agent.run([ChatMessage.from_user("Weather in Berlin?")]) + + # Before the first call there is no usage yet (0); before the second call it reflects the first call (15), + # confirming the recorder runs in the loop and the value is refreshed per call rather than accumulated. + assert seen == [0, 15] + + +class TestAgentContextTokensAsync: + @pytest.mark.asyncio + async def test_before_llm_hook_reads_refreshed_context_tokens_async(self, weather_tool): + first_step = [ + _assistant_with_usage( + tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})], + usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + ) + ] + second_step = [_assistant_with_usage("Done.", usage={"prompt_tokens": 20, "completion_tokens": 8})] + + seen: list[int] = [] + + @hook + def capture(state: State) -> None: + seen.append(state.get("context_tokens")) + + agent = Agent( + chat_generator=MockChatGenerator(first_step + second_step), + tools=[weather_tool], + hooks={"before_llm": [capture]}, + ) + await agent.run_async([ChatMessage.from_user("Weather in Berlin?")]) + assert seen == [0, 15] + + class TestAgentExitReason: """The Agent reports why it stopped via the `exit_reason` output, so downstream code can route its output.""" @@ -1417,7 +1482,7 @@ def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool): assert set(llm_tags) == {"haystack.agent.step.llm.input", "haystack.agent.step.llm.output"} assert ( llm_tags["haystack.agent.step.llm.input"] - == '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]}' # noqa: E501 + == '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "agents.test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]}' # noqa: E501 ) assert ( llm_tags["haystack.agent.step.llm.output"] @@ -1432,9 +1497,9 @@ def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool): _, run_tags = agent_spans[2] assert run_tags == { "haystack.agent.max_steps": 100, - "haystack.agent.tools": '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]', # noqa: E501 + "haystack.agent.tools": '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "agents.test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]', # noqa: E501 "haystack.agent.exit_conditions": '["text"]', - "haystack.agent.state_schema": '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}, "step_count": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "token_usage": {"type": "dict[str, typing.Any]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "tool_call_counts": {"type": "dict[str, int]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "exit_reason": {"type": "str", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "continue_run": {"type": "bool", "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]", "handler": "haystack.components.agents.state.state_utils.replace_values"}}', # noqa: E501 + "haystack.agent.state_schema": '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}, "step_count": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "token_usage": {"type": "dict[str, typing.Any]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "tool_call_counts": {"type": "dict[str, int]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "exit_reason": {"type": "str", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "continue_run": {"type": "bool", "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]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "context_tokens": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"}}', # noqa: E501 "haystack.agent.input": '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}', # noqa: E501 "haystack.agent.output": '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}, {"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello"}]}], "step_count": 1, "token_usage": {}, "tool_call_counts": {"weather_tool": 0}, "exit_reason": "text", "last_message": {"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello"}]}}', # noqa: E501 "haystack.agent.steps_taken": 1, @@ -1586,7 +1651,7 @@ async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_t assert set(llm_tags) == {"haystack.agent.step.llm.input", "haystack.agent.step.llm.output"} assert ( llm_tags["haystack.agent.step.llm.input"] - == '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]}' # noqa: E501 + == '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "agents.test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]}' # noqa: E501 ) assert ( llm_tags["haystack.agent.step.llm.output"] @@ -1599,9 +1664,9 @@ async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_t _, run_tags = agent_spans[2] assert run_tags == { "haystack.agent.max_steps": 100, - "haystack.agent.tools": '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]', # noqa: E501 + "haystack.agent.tools": '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "agents.test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null, "async_function": null}}]', # noqa: E501 "haystack.agent.exit_conditions": '["text"]', - "haystack.agent.state_schema": '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}, "step_count": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "token_usage": {"type": "dict[str, typing.Any]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "tool_call_counts": {"type": "dict[str, int]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "exit_reason": {"type": "str", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "continue_run": {"type": "bool", "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]", "handler": "haystack.components.agents.state.state_utils.replace_values"}}', # noqa: E501 + "haystack.agent.state_schema": '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}, "step_count": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "token_usage": {"type": "dict[str, typing.Any]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "tool_call_counts": {"type": "dict[str, int]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "exit_reason": {"type": "str", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "continue_run": {"type": "bool", "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]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "context_tokens": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"}}', # noqa: E501 "haystack.agent.input": '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}', # noqa: E501 "haystack.agent.output": '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}, {"role": "assistant", "meta": {"model": "mock-model", "index": 0, "finish_reason": "stop", "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}}, "name": null, "content": [{"text": "Hello"}]}], "step_count": 1, "token_usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, "tool_call_counts": {"weather_tool": 0}, "exit_reason": "text", "last_message": {"role": "assistant", "meta": {"model": "mock-model", "index": 0, "finish_reason": "stop", "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}}, "name": null, "content": [{"text": "Hello"}]}}', # noqa: E501 "haystack.agent.steps_taken": 1, diff --git a/test/components/agents/test_agent_hitl.py b/test/components/agents/test_agent_hitl.py index cae533a193..90037a339d 100644 --- a/test/components/agents/test_agent_hitl.py +++ b/test/components/agents/test_agent_hitl.py @@ -103,7 +103,7 @@ def test_to_dict(self, tools, confirmation_hook, monkeypatch): "required": ["a", "b"], "type": "object", }, - "function": "test_agent_hitl.addition_tool", + "function": "agents.test_agent_hitl.addition_tool", "async_function": None, "outputs_to_string": None, "inputs_from_state": None, diff --git a/test/components/agents/test_state_class.py b/test/components/agents/test_state_class.py index f98bd53def..66d0d7caf2 100644 --- a/test/components/agents/test_state_class.py +++ b/test/components/agents/test_state_class.py @@ -351,7 +351,7 @@ def test_schema_to_dict(self, basic_schema): def test_schema_to_dict_with_handlers(self, complex_schema): expected_dict = { - "numbers": {"type": "list", "handler": "test_state_class.numbers_handler"}, + "numbers": {"type": "list", "handler": "agents.test_state_class.numbers_handler"}, "metadata": {"type": "dict"}, "name": {"type": "str"}, } @@ -365,7 +365,7 @@ def test_schema_from_dict(self, basic_schema): def test_schema_from_dict_with_handlers(self, complex_schema): schema_dict = { - "numbers": {"type": "list", "handler": "test_state_class.numbers_handler"}, + "numbers": {"type": "list", "handler": "agents.test_state_class.numbers_handler"}, "metadata": {"type": "dict"}, "name": {"type": "str"}, } diff --git a/test/components/agents/test_utils.py b/test/components/agents/test_utils.py new file mode 100644 index 0000000000..71c448e758 --- /dev/null +++ b/test/components/agents/test_utils.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from haystack.components.agents.state import State, replace_values +from haystack.components.agents.utils import _context_tokens_from_usage, _record_context_tokens +from haystack.dataclasses import ChatMessage + + +class TestContextTokensFromUsage: + """`_context_tokens_from_usage` normalizes real provider `meta["usage"]` shapes to input + output tokens.""" + + # OpenAI Chat Completions (core repo): reasoning_tokens is a subset of completion_tokens (64 of 74), so + # context_tokens includes reasoning. + OPENAI_CHAT_USAGE = { + "completion_tokens": 74, + "prompt_tokens": 19, + "total_tokens": 93, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 64, + "rejected_prediction_tokens": 0, + }, + "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, + } + # OpenAI Responses API (core repo). + OPENAI_RESPONSES_USAGE = { + "input_tokens": 19, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 58, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 77, + } + # Anthropic chat generator (integration): no total_tokens, and thinking_tokens is a subset of completion_tokens + # (57 of 63). + ANTHROPIC_USAGE = { + "cache_creation": {"ephemeral_1h_input_tokens": 0, "ephemeral_5m_input_tokens": 0}, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "inference_geo": "not_available", + "output_tokens_details": {"thinking_tokens": 57}, + "server_tool_use": None, + "service_tier": "standard", + "prompt_tokens": 49, + "completion_tokens": 63, + } + # Amazon Bedrock chat generator (integration). + BEDROCK_USAGE = { + "prompt_tokens": 20, + "completion_tokens": 5, + "total_tokens": 25, + "cache_read_input_tokens": 0, + "cache_write_input_tokens": 0, + "cache_details": {}, + } + # Cohere chat generator (integration). + COHERE_USAGE = {"prompt_tokens": 15.0, "completion_tokens": 3.0} + # Mistral chat generator (integration). + MISTRAL_USAGE = { + "prompt_tokens": 30, + "total_tokens": 34, + "completion_tokens": 4, + "prompt_tokens_details": {"cached_tokens": 0}, + } + # Nvidia chat generator (integration). + NVIDIA_USAGE = { + "completion_tokens": 2, + "prompt_tokens": 48, + "total_tokens": 50, + "completion_tokens_details": None, + "prompt_tokens_details": {"audio_tokens": None, "cached_tokens": 16}, + } + # Google GenAI chat generator (integration): thoughts_token_count (281) is NOT part of completion_tokens; + # total_tokens (300) includes it, so context_tokens (19) excludes thoughts. + GOOGLE_GENAI_USAGE = { + "prompt_tokens": 16, + "completion_tokens": 3, + "total_tokens": 300, + "thoughts_token_count": 281, + "prompt_token_count": 16, + "candidates_token_count": 3, + "total_token_count": 300, + "prompt_tokens_details": [{"modality": "TEXT", "token_count": 16}], + } + + @pytest.mark.parametrize( + "usage, expected", + [ + (OPENAI_CHAT_USAGE, 93), + (OPENAI_RESPONSES_USAGE, 77), + (ANTHROPIC_USAGE, 112), + (BEDROCK_USAGE, 25), + (COHERE_USAGE, 18), + (MISTRAL_USAGE, 34), + (NVIDIA_USAGE, 50), + (GOOGLE_GENAI_USAGE, 19), # 16 + 3, deliberately not the 300 total (which includes 281 thoughts tokens) + ({"prompt_tokens": 10}, 10), # only one side reported -> partial count + ({"completion_tokens": 7}, 7), + ({}, 0), # no usage reported + ({"foo": 5, "bar": 9}, 0), # no recognized keys + ], + ) + def test_normalizes_provider_shapes(self, usage, expected): + assert _context_tokens_from_usage(usage) == expected + + def test_bool_values_are_not_counted_as_tokens(self): + # bool is an int subclass; True/False under a token key must be skipped, not summed. + assert _context_tokens_from_usage({"prompt_tokens": True, "completion_tokens": 5}) == 5 + + +class TestRecordContextTokens: + """`_record_context_tokens` replaces the value with the latest reply's input+output, only when usage is reported.""" + + def _state(self) -> State: + state = State(schema={"context_tokens": {"type": int, "handler": replace_values}}) + state.set("context_tokens", 0) + return state + + def test_records_latest_reply_usage_replacing_previous_value(self): + state = self._state() + state.set("context_tokens", 999) + _record_context_tokens( + state, [ChatMessage.from_assistant("Hi", meta={"usage": {"prompt_tokens": 12, "completion_tokens": 3}})] + ) + assert state.get("context_tokens") == 15 + + def test_no_messages_leaves_value_untouched(self): + state = self._state() + state.set("context_tokens", 42) + _record_context_tokens(state, []) + assert state.get("context_tokens") == 42 + + def test_missing_or_empty_usage_leaves_value_untouched(self): + state = self._state() + _record_context_tokens(state, [ChatMessage.from_assistant("no usage here")]) + _record_context_tokens(state, [ChatMessage.from_assistant("empty", meta={"usage": {}})]) + assert state.get("context_tokens") == 0