概览
处理大型数据集时,直接通过模型传递原始数据会导致以下问题:减慢响应速度、增加 token 费用并引入推理错误。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
返回示例可减少所有未来对话轮次的输入令牌数量,从而降低费用并缩短响应时间。它还可以防止大型查询结果挤占上下文窗口,从而损害模型遵循指令的能力。
确保代理提示明确说明工具会按设计返回样本。如果没有此上下文,模型可能会因部分数据而感到困惑,浪费计算和推理令牌来尝试协调缺失的记录,或启动冗余的工具调用来获取其余数据。
示例:处理大型结果集
- 用户询问“目前延误比率超过 2 的路线有哪些?”
- 代理针对
historical_travel_time的查询返回了数千行。 after_tool_callback会将完整结果保存在会话状态中,并仅返回一个小样本以及总行数。- 模型会根据该样本进行推理,然后撰写摘要。
- 应用或下游工具直接从会话状态读取完整的数据集,以呈现完整的表格或导出供用户使用。
要点总结
- 避免在提示中包含大型载荷:将大型数据集存储在会话状态中,而不是通过提示历史记录来回传递。
- 截断大型工具输出:使用
after_tool_callback向模型返回小样本,同时将完整结果保存在会话状态中。 - 告知模型有关样本的信息:添加提示规则,说明结果是抽样的,以防止模型混淆、浪费推理令牌和进行冗余的工具调用。
后续步骤
- ADK 回调参考文档:阅读官方 ADK 回调指南,详细了解模型级、智能体级和工具级拦截。
贡献者
Nathaniel Thomas | 软件工程实习生,Google Maps Platform