tooluniverse.tool_finder_embedding 模块

class tooluniverse.tool_finder_embedding.ToolFinderEmbedding[源代码]

基类：BaseTool

一种工具查找模型，利用 RAG（检索增强生成）通过语义相似度搜索，根据用户查询查找相关工具。

该类利用句子转换器对工具描述进行编码，并通过基于嵌入的相似度匹配，找到与给定查询最相关的工具。

rag_model_name

用于生成嵌入的句子转换器模型名称

类型:

str

rag_model

已加载的句子转换器模型

类型:

SentenceTransformer

tool_desc_embedding

工具描述的缓存嵌入向量

类型:

torch.Tensor

tool_name

可用工具名称列表

类型:

列表

tool_embedding_path

缓存工具嵌入文件的路径

类型:

str

special_tools_name

排除结果的特殊工具列表

类型:

列表

tooluniverse

参考包含所有工具的工具库

__init__(tool_config, tooluniverse)[源代码]

使用配置和RAG模型初始化ToolFinderEmbedding。

参数:

tool_config (dict) – 工具的配置字典

load_rag_model()[源代码]

加载用于基于RAG工具检索的句子转换器模型。

配置模型的适当序列长度和分词器设置，以实现工具描述编码的最佳性能。

The model is automatically moved to GPU if available for faster inference.

抛出:

ImportError – If sentence-transformers is not installed.

load_tool_desc_embedding(tooluniverse, include_names=None, exclude_names=None, include_categories=None, exclude_categories=None)[源代码]

从工具库中加载或生成工具描述的嵌入数据。

此方法会从磁盘加载缓存的嵌入，或者通过编码所有工具描述生成新的嵌入。嵌入会被缓存到磁盘，以加快后续加载速度。在生成嵌入后，会适当清理内存以避免出现内存溢出（OOM）问题。

参数:

  • tooluniverse – ToolUniverse 实例，包含所有可用工具

  • include_names (list, optional) – 特定工具名称包括

  • exclude_names (list, optional) – 要排除的工具名称

  • include_categories (list, optional) – 工具类别包括

  • exclude_categories (list, optional) – 要排除的工具类别

rag_infer(query, top_k=5)[源代码]

执行RAG推理以查找与给定查询最相关的工具。

使用查询嵌入与预先计算的工具嵌入之间的语义相似度来识别最相关的工具。

参数:

  • query (str) – 用户查询或所需功能描述

  • top_k (int, optional) – 返回的顶级工具数量。默认值为5。

返回值

list: List of top-k tool names ranked by relevance to the query

抛出:
find_tools(message=None, picked_tool_names=None, rag_num=5, return_call_result=False, categories=None)[源代码]

根据消息或预先选择的工具名称查找相关工具。

该方法通过RAG推理根据消息查找工具，或处理预先选择的工具名称列表。它会过滤特殊工具，并返回适用于智能体工作流的工具提示。

参数:

  • message (str, optional) – 查询消息以查找工具。若 picked_tool_names 为 None，则为必填项。

  • picked_tool_names (list, optional) – 预先选择的工具名称以进行处理。若消息为 None，则为必填项。

  • rag_num (int, optional) – 过滤后返回的工具数量。默认值为5。

  • return_call_result (bool, optional) – 如果为 True，则返回提示词和工具名称。默认值为 False。

  • categories (list, optional) – 可用于筛选的工具类别列表。目前尚未针对基于嵌入的搜索实现该功能。

返回值
str or tuple:

  • 如果 return_call_result 为 False：工具以格式化字符串形式提示

  • 如果 return_call_result 为 True：返回 (tool_prompts, tool_names) 元组

抛出:

AssertionError – 如果 message 和 picked_tool_names 均为 None

run(arguments)[源代码]

按照标准工具接口，使用给定参数运行工具查找器。

这是使用 ToolFinderEmbedding 作为标准工具的主要入口点。它从参数字典中提取参数，并委托给 find_tools() 函数。

参数:

arguments (dict) – 字典包含： - description（字符串，可选）：查询消息，用于查找工具（对应于 ‘message’） - limit（整数，可选）：返回的工具数量（对应于 ‘rag_num’），默认为5 - picked_tool_names（列表，可选）：预先选择的工具名称列表，用于处理 - return_call_result（布尔值，可选）：是否返回提示和名称，默认为False - categories（列表，可选）：用于筛选的工具类别列表