Tool Retriever
Agent在实际生产应用中往往涉及到的工具数量较多,如果把所用的工具全部添加至Agent会产生如下问题:
- 占用大量输入token。
- 和问题无关的工具太多,影响模型的判断。
通过Tool Retriever可以解决上述问题,其原理是在Agent运行前,先从所有可用的工具中选择与问题最相关的工具,再交给Agent去处理,示例如下:
- 定义一个Tool Retriever:
final List<Tool> toolList = Arrays.asList(new MetricQuery(), new ReserveMeeting(), new ReserveMeetingRoom(), new StuffQuery()); // 新增InMemoryToolProvider,添加工具集 final InMemoryToolProvider inMemoryToolProvider = new InMemoryToolProvider(); inMemoryToolProvider.add(toolList); // 初始化CSSToolRetriever final CSSToolRetriever cssToolRetriever = new CSSToolRetriever(inMemoryToolProvider, VectorStoreConfig.builder() .indexName(TestConstant.CSS_TOOL_RETRIEVER_INDEX) .vectorFields(Arrays.asList("name", "description")) .build());
定义一个ToolRetriever包含ToolProvider和向量数据库配置2个参数。其中,ToolProvider的作用为根据工具检索的结果组装工具。
上述例子使用了一个简单的InMemoryToolProvider,InMemoryToolProvider的原理为将完整的工具存入内存,再根据工具检索的结果(toolId)将其从内存中取出。一般来说,ToolProvider将由用户自定义,将在后续示例中说明。
此外,上述例子使用的向量数据库配置指定索引名称,以及使用name和description作为向量化字段,因此工具入库时,会将工具的name和description进行向量化,并在后续的检索中生效。
注意,上述toolList中包含的工具在SDK中并不存在,需要替换成实际的工具。
- 向ToolRetriever中添加工具:
// 添加工具 cssToolRetriever.addTools(toolList);
工具添加后,会存储在向量库的索引中,并将指定的字段向量化。
- 从ToolRetriever中查找工具:
// 查找工具 List<Tool> result = cssToolRetriever.search("预订会议室", 2);
返回的result中,包含与预订会议室最相关的工具。搜索支持topK和阈值2个参数,例如上例指定topK=2,则最多返回2个工具。
- 从ToolRetriever中删除工具:
// 删除工具 cssToolRetriever.remove(Collections.singletonList("a_tool_id"));
以上为一个比较基础的用法,在实际使用过程中会有更加灵活的场景,可以通过自定义ToolProvider的方式解决。
- 自定义ToolProvider:
// 初始化CSSToolRetriever,使用ToolProviderWithMetadata作为ToolProvider final CSSToolRetriever cssToolRetriever = new CSSToolRetriever(new ToolProviderWithMetadata(), VectorStoreConfig.builder() .indexName(TestConstant.CSS_TOOL_RETRIEVER_INDEX) .vectorFields(Arrays.asList("name", "description")) .build());
其中,ToolProviderWithMetadata为自定义ToolProvider:private static class ToolProviderWithMetadata implements ToolProvider { @Override public List<Tool> provide(List<RetrievedTool> retrievedTools, String query) { doSomeFilter(retrievedTools, query); return retrievedTools.stream().map((Function<RetrievedTool, Tool>) retrievedTool -> { final DynamicTool tool = new DynamicTool(); tool.setToolId(retrievedTool.getToolId()); final Map<String, Object> toolMetadata = retrievedTool.getToolMetadata(); tool.setToolDesc(toolMetadata.get("description").toString()); tool.setToolPrinciple(toolMetadata.get("description").toString()); // 从plugin_schema和openapi_schema构建工具参数信息 tool.setInputDesc("会议开始结束时间,会议室"); tool.setOutputDesc("会议预订结果"); tool.setInputSchema( "{\"type\":\"object\",\"properties\":{\"meetingRoom\":{\"type\":\"string\",\"description\":\"会议室\"},\"start\":{\"type\":\"string\",\"description\":\"会议开始时间,格式为HH:mm\"},\"end\":{\"type\":\"string\",\"description\":\"会议结束时间,格式为HH:mm\"}},\"required\":[\"meetingRoom\",\"start\",\"end\"]}"); tool.setOutputSchema(""); tool.setFunction(s -> { log.info("do the open api call plugin_schema={} openapi_schema={}", toolMetadata.get("plugin_schema"), toolMetadata.get("openapi_schema")); return null; }); return tool; }).collect(Collectors.toList()); } private void doSomeFilter(List<RetrievedTool> retrievedTools, String query) { // do some filter here log.info("{} {}", query, retrievedTools); } }
其中,toolProvider中实现了provide接口,可以利用工具检索的返回动态构建出工具列表,同时也可以加一些后处理工作,如根据黑白名单做工具的过滤。
- 与上述的toolProvide呼应,在向toolRetriever中添加工具时,可以添加任意的元数据,用于在tooProvider中把工具组装出来:
// 构造工具元数据 Map<String, Object> toolMetaData = new HashMap<>(); toolMetaData.put("name_for_human", "预订会议室"); toolMetaData.put("name", "reserve_meeting_room"); toolMetaData.put("description_for_human", "预订会议室"); toolMetaData.put("description", "预订会议室,请在需要预订会议室时调用此工具"); toolMetaData.put("status", "on"); toolMetaData.put("plugin_type", "API"); toolMetaData.put("plugin_schema", "this is a plugin schema"); toolMetaData.put("openapi_schema", "this is a openapi schema"); final ToolMetadata toolMetadata = new ToolMetadata(); toolMetadata.setToolId("reserve_meeting_room"); toolMetadata.setToolMetadata(toolMetaData); // 工具管理面添加工具到toolRetriever,这里实际可以添加若干个工具 cssToolRetriever.addToolsFromMetadata(Collections.singletonList(toolMetadata)); // 运行时检索工具,并添加到Agent执行 final List<Tool> toolList = cssToolRetriever.search("预订会议室", 1, 0.8f);
工具的检索与之前的用法一致。
- 将Tool Retriever集成在Agent中的完整示例如下:
// 工具集 final List<Tool> toolList = Arrays.asList(new MetricQuery(), new ReserveMeeting(), new ReserveMeetingRoom(), new StuffQuery()); // 新增InMemoryToolProvider,添加工具集 final InMemoryToolProvider inMemoryToolProvider = new InMemoryToolProvider(); inMemoryToolProvider.add(toolList); // 初始化CSSToolRetriever final VectorStoreConfig vectorStoreConfig = VectorStoreConfig.builder() .indexName(TestConstant.CSS_TOOL_RETRIEVER_INDEX) .vectorFields(Arrays.asList("name", "description")) .build(); final CSSToolRetriever cssToolRetriever = new CSSToolRetriever(inMemoryToolProvider, vectorStoreConfig); // 添加工具 cssToolRetriever.addTools(toolList); // 添加多轮改写 cssToolRetriever .setQueryPreprocessor(messages -> new ConversationRewriteSkill(LLMs.of(LLMs.PANGU)).rewrite(messages)); // 为Agent添加ToolRetriever Agent agent = new ReactPanguAgent( new YundaoLLM(LLMConfig.builder().llmModuleConfig(LLMModuleConfig.builder().build()).build())); agent.setToolRetriever(cssToolRetriever); // 多轮对话调用 List<ConversationMessage> messages = new ArrayList<>(); messages.add(ConversationMessage.builder().role(Role.USER).content("定个2点的会议").build()); messages.add( ConversationMessage.builder().role(Role.ASSISTANT).content("请问您的会议预计何时结束?另外,您是需要预订线上会议还是实体会议室?").build()); messages.add(ConversationMessage.builder().role(Role.USER).content("4点结束,线上会议").build()); agent.run(messages); // 删除工具 cssToolRetriever.remove(toolList.stream().map(Tool::getToolId).collect(Collectors.toList()));
有两个变化值得关注,一是为ToolRetriever添加了一个queryPreprocessor,它的作用为对用户输入的多轮对话进行改写,会将改写后的结果作为工具检索的输入,这里使用了系统内置的ConversationRewriteSkill,它的作用为将多轮对话改写为单轮。二是在创建一个Agent后,调用了setToolRetriever方法为其添加了一个ToolRetriever,这样Agent所使用的工具会根据用户的对话动态的选择。