Haystack Agent 如何用 TokenBudgetHook 在 token 用量达到阈值时停止执行-育师

Haystack Agent 如何用 TokenBudgetHook 在 token 用量达到阈值时停止执行

【免费下载链接】haystackOpen-source AI orchestration framework for building context-engineered, production-ready LLM applications. Design modular pipelines and agent workflows with explicit control over retrieval, routing, memory, and generation. Built for scalable agents, RAG, multimodal applications, semantic search, and conversational systems.项目地址: https://gitcode.com/GitHub_Trending/ha/haystack

如果你用 Haystack 的Agent组件跑带工具调用的研究、检索类任务，会遇到一个控制问题：一次 run 可能连续调用多次 LLM 和工具，token 花费不受限制。TokenBudgetHook就是为此提供的现成 hook：把它注册到 Agent 的before_llmhook point 上，Agent 在每次 LLM 调用前检查累积的token_usage，一旦达到你配置的阈值，本次 run 会在下一次 LLM 调用之前结束，exit_reason被置为"token_budget_exceeded"，此前收集到的消息全部保留。完成配置后，你可以用一行print(result["exit_reason"])核对预算是否真的生效。

适用前提：使用支持工具调用的 Chat Generator 构建Agent（示例使用OpenAIChatGenerator），依赖包为haystack-ai。注意TokenBudgetHook目前是实验性（experimental）API，其接口可能在任意版本中变更，不遵循常规弃用流程。

工作原理：在`before_llm`处比较累积用量

TokenBudgetHook是 Token Budget 文档描述的一个 hook 用例，它基于 Hooks 机制工作：

hook 必须注册在before_llmhook point 下，在每次 chat generator 调用之前执行；
hook 读取 Agent 状态中自动累积的运行元数据token_usage（该值跨整个 run 累加，而不是某一次调用的用量），与max_total_tokens比较；
当用量达到或超过阈值时，hook 通过stop_run状态键请求停止 run。stop_run会在下一次 LLM 调用前被读取，并直接作为exit_reason输出。

两点边界需要清楚：

预算只覆盖 Agent 自身 chat generator 回复产生的用量。工具内部发起的 LLM 调用（例如某个工具内部调用了自己的模型）不计入这个预算。
检查发生在调用之前，所以把总用量推过阈值的那次调用已经完成。最终的实际用量可能超出max_total_tokens大约一次 LLM 调用的成本。

token_usage的读取在源码中做了兼容处理：优先读total_tokens，否则按已知的 input/output 命名约定累加求和，因此 OpenAI 风格（prompt_tokens/completion_tokens）和其他风格的用量报告都能被识别。

配置：把 hook 注册到`before_llm`

TokenBudgetHook的构造函数只有两个参数（均为关键字参数）：

max_total_tokens（必填）：累积用量达到该值后停止 run。传小于 1 的值会抛出ValueError；
add_final_message（默认False）：预算触发的停止发生时，追加一条说明停止原因的 assistant 消息。

从haystack.hooks.budget导入后，通过Agent的hooks参数注册。下面的完整示例来自官方文档，包含一个模拟搜索的@tool占位实现（文档标注它会调用真实的搜索 API，示例里只是随机重复事实字符串）：

import random from typing import Annotated from haystack.components.agents import Agent from haystack.components.generators.chat import OpenAIChatGenerator from haystack.dataclasses import ChatMessage from haystack.hooks.budget import TokenBudgetHook from haystack.tools import tool FACTS = [ "Capybaras are the largest living rodents, weighing up to 65 kg. ", "Capybaras are highly social and live in groups of ten to twenty. ", "Capybaras are excellent swimmers and can stay underwater for five minutes. ", "Capybaras are famously relaxed and often share space with birds and monkeys. ", ] @tool def search(query: Annotated[str, "The search query"]) -> str: """Search the web.""" # Placeholder: would call a real search API # Repeat the result to simulate a longer search response return random.choice(FACTS) * 20 agent = Agent( chat_generator=OpenAIChatGenerator(model="gpt-5-mini"), tools=[search], system_prompt="You are a research assistant. Search one aspect at a time before answering.", hooks={"before_llm": [TokenBudgetHook(max_total_tokens=3_000)]}, ) agent.warm_up() result = agent.run( messages=[ ChatMessage.from_user( "Research capybaras: size, social life, swimming and temperament." ) ] ) print(result["exit_reason"]) # >> token_budget_exceeded

max_total_tokens=3_000是文档特意选的低阈值，目的是让研究任务在 Agent 完成报告之前就被停止。如果你的任务预期用量更高，把阈值调到对应数值即可（例如 release notes 中的示例使用max_total_tokens=100_000）。

验证预算是否生效

Agent.run()返回的字典包含messages、last_message、step_count、token_usage、tool_call_counts、exit_reason等运行元数据。判断预算是否按预期工作：

主判据：result["exit_reason"]是否为"token_budget_exceeded"。上面的文档示例输出即为token_budget_exceeded（文档示例）；
消息保留：预算停止不会丢数据，run 到此为止收集的消息仍可通过result["messages"]访问。文档说明 Agent 是"中途停止研究、保留已收集的消息"；
对照场景：文档指出，把阈值调高后 Agent 可以完成报告并返回"text"作为exit_reason。因此exit_reason同时能区分"正常完成"和"预算耗尽"两种结局，方便下游用ConditionalRouter之类的组件做路由。
日志：触发停止时 hook 会输出一条 WARNING 日志，格式为Agent reached its token budget of {max_total_tokens} ({total_tokens} used); requesting a stop.，方便在运行日志中确认停止时机（{...}为格式化占位符，实际日志会代入具体数字，见源码）。

仓库中的单元测试展示了更细的验证方式：用MockChatGenerator模拟每次回复各消耗 60 token，配置max_total_tokens=100后断言 chat generator 只被调用了 2 次且exit_reason == "token_budget_exceeded"；同时验证了total_tokens、prompt_tokens/completion_tokens、input_tokens/output_tokens三种用量报告格式都能正确触发停止。

可选：为预算停止补一条说明消息

预算触发的停止发生时，最后一条消息可能是一条工具结果而不是最终回答。如果希望结果里带一条解释性收尾消息，设置add_final_message=True：

TokenBudgetHook(max_total_tokens=3_000, add_final_message=True)

此时会追加一条固定文案的 assistant 消息（The Agent stopped because the token budget was exceeded.），并成为last_message。

如果文案需要自定义，或者要同时处理max_agent_steps等其他停止原因，文档给出的路径是写一个after_runhook——它无论 run 因何结束都会执行：

from haystack.components.agents.state import State from haystack.dataclasses import ChatMessage from haystack.hooks import hook @hook def explain_stop(state: State) -> None: if state.get("exit_reason") == "token_budget_exceeded": state.set( "messages", [ChatMessage.from_assistant("I ran out of budget before finishing.")], )

把这个 hook 注册到hooks={"after_run": [explain_stop]}即可与TokenBudgetHook并存。

限制与注意

TokenBudgetHook是实验性 API，文档明确提示其 API 可能在任意版本中变更，不遵循常规弃用策略；生产使用需留意版本升级时的行为变化。
预算不含工具和外部 hook 发起的 LLM 调用，如果你的成本主要来自工具内部调用，这个 hook 管不到那部分。
最终实际用量可能超出阈值一次 LLM 调用的成本，这是检查时机决定的，不是 bug。
max_total_tokens必须大于等于 1，否则构造时抛ValueError。
hook 实现了to_dict/from_dict，随 Agent 序列化保存后可以正常反序列化，适合放在 YAML 定义的管线配置里。

相关文档：Token Budget、Hooks、Agent，源码位于 haystack/hooks/budget/。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

Haystack Agent 如何用 TokenBudgetHook 在 token 用量达到阈值时停止执行