概览
处理大型数据集时,直接将原始数据传递给模型会导致一些问题:响应速度变慢、词元费用增加,以及推理错误。RMI 数据可能包含数百万行,因此让智能体查询无限量的数据并将其全部提取以进行分析,会消耗大量 token 并导致上下文窗口膨胀,从而降低智能体在后续每个轮次的性能。
例如,对 RMI historical_travel_time 表的查询可能会返回数千个路段行。如果将完整的结果集放回对话记录中,后续的每个对话轮次都会变慢,成本会增加,并且模型会分心,无法合成最终答案。
仅靠提示指令无法解决这些问题。唯一可靠的方法是拦截工具调用,并在模型的上下文窗口之外路由大型数据载荷。
ADK 工具回调
ADK 提供回调钩子来拦截工具执行。如需处理大型结果集,请使用 after_tool_callback:
before_tool_callback(tool, args, tool_context):在模型选择工具之后、工具执行之前运行。after_tool_callback(tool, args, tool_context, tool_response):在工具运行完毕后、结果返回模型之前运行。使用它来修改工具回答。
回调会收到一个 tool_context 对象,该对象有权访问会话 state。这样,您就可以在对话轮次之间将大型数据集存储在后台内存中,而无需将其放入模型的提示中。
# root_agent = llm_agent.Agent(
# ...,
# after_tool_callback=stash_and_truncate_results,
# )
执行后截断和暂存
目标:向模型展示一小部分查询结果样本,以便其做出决策,同时保存完整数据集以供最终输出。
对于 RMI 智能体,after_tool_callback 会拦截 execute_sql 结果,将会话状态中的完整行列表保存,并返回一个简短示例以及总行数:
def stash_and_truncate_results(tool, args, tool_context, tool_response):
if tool.name == "execute_sql" and "rows" in tool_response:
rows = tool_response["rows"]
tool_context.state["last_sql_result"] = rows # keep everything
tool_response["total_rows_fetched"] = len(rows)
tool_response["rows"] = rows[:SAMPLE_LIMIT] # show a sample
tool_response["message"] = (
f"Showing {SAMPLE_LIMIT} of {len(rows)} rows. "
"Full result stashed in background state."
)
return tool_response
返回样本可减少后续所有轮次的输入 token,从而降低费用并缩短回答时间。它还可以防止大型查询结果挤占上下文窗口,从而影响模型遵循指令的能力。
确保代理提示明确说明工具会按设计返回示例。如果没有这些上下文,模型可能会因数据不完整而感到困惑,从而浪费计算和推理 token 来尝试协调缺失的记录,或发起冗余的工具调用来提取其余数据。
示例:处理大型结果集
- 用户询问:“目前哪些路线的延误率超过 2?”
- 代理针对
historical_travel_time的查询会返回数千行。 after_tool_callback会将完整结果保存在会话状态中,并仅返回一小部分样本以及总行数。- 模型会根据该样本进行推理,以撰写摘要。
- 应用或下游工具直接从会话状态读取完整数据集,以便为用户呈现完整表格或导出数据。
要点总结
- 避免在提示中包含大型载荷:将大型数据集存储在会话状态中,而不是通过提示历史记录来回传递。
- 截断大型工具输出:使用
after_tool_callback向模型返回少量样本,同时将完整结果保存在会话状态中。 - 告知模型有关样本的信息:添加提示规则,说明结果是抽样得出的,以防止模型混淆、浪费推理 token 和进行冗余工具调用。
后续步骤
- ADK 回调参考文档:阅读官方 ADK 回调指南,详细了解模型、智能体和工具级拦截。
贡献者
Nathaniel Thomas | Google Maps Platform 软件工程实习生