ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

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

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目前是实验性experimentalAPI其接口可能在任意版本中变更不遵循常规弃用流程。工作原理在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_llmTokenBudgetHook的构造函数只有两个参数均为关键字参数max_total_tokens必填累积用量达到该值后停止 run。传小于 1 的值会抛出ValueErroradd_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_generatorOpenAIChatGenerator(modelgpt-5-mini), tools[search], system_promptYou are a research assistant. Search one aspect at a time before answering., hooks{before_llm: [TokenBudgetHook(max_total_tokens3_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_exceededmax_total_tokens3_000是文档特意选的低阈值目的是让研究任务在 Agent 完成报告之前就被停止。如果你的任务预期用量更高把阈值调到对应数值即可例如 release notes 中的示例使用max_total_tokens100_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_tokens100后断言 chat generator 只被调用了 2 次且exit_reason token_budget_exceeded同时验证了total_tokens、prompt_tokens/completion_tokens、input_tokens/output_tokens三种用量报告格式都能正确触发停止。可选为预算停止补一条说明消息预算触发的停止发生时最后一条消息可能是一条工具结果而不是最终回答。如果希望结果里带一条解释性收尾消息设置add_final_messageTrueTokenBudgetHook(max_total_tokens3_000, add_final_messageTrue)此时会追加一条固定文案的 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/。【免费下载链接】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创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表