tooluniverse.tool_finder_embedding module

class tooluniverse.tool_finder_embedding.ToolFinderEmbedding[source]

Bases: BaseTool

A tool finder model that uses RAG (Retrieval-Augmented Generation) to find relevant tools based on user queries using semantic similarity search.

This class leverages sentence transformers to encode tool descriptions and find the most relevant tools for a given query through embedding-based similarity matching.

rag_model_name

Name of the sentence transformer model for embeddings

Type:

str

rag_model

The loaded sentence transformer model

Type:

SentenceTransformer

tool_desc_embedding

Cached embeddings of tool descriptions

Type:

torch.Tensor

tool_name

List of available tool names

Type:

list

tool_embedding_path

Path to cached tool embeddings file

Type:

str

special_tools_name

List of special tools to exclude from results

Type:

list

tooluniverse

Reference to the tool universe containing all tools

__init__(tool_config, tooluniverse)[source]

Initialize the ToolFinderEmbedding with configuration and RAG model.

Parameters:

tool_config (dict) – Configuration dictionary for the tool

load_rag_model()[source]

Load the sentence transformer model for RAG-based tool retrieval.

Configures the model with appropriate sequence length and tokenizer settings for optimal performance in tool description encoding.

load_tool_desc_embedding(tooluniverse, include_names=None, exclude_names=None, include_categories=None, exclude_categories=None)[source]

Load or generate embeddings for tool descriptions from the tool universe.

This method either loads cached embeddings from disk or generates new ones by encoding all tool descriptions. Embeddings are cached to disk for faster subsequent loads. Memory is properly cleaned up after embedding generation to avoid OOM issues.

Parameters:

  • tooluniverse – ToolUniverse instance containing all available tools

  • include_names (list, optional) – Specific tool names to include

  • exclude_names (list, optional) – Tool names to exclude

  • include_categories (list, optional) – Tool categories to include

  • exclude_categories (list, optional) – Tool categories to exclude

rag_infer(query, top_k=5)[source]

Perform RAG inference to find the most relevant tools for a given query.

Uses semantic similarity between the query embedding and pre-computed tool embeddings to identify the most relevant tools.

Parameters:

  • query (str) – User query or description of desired functionality

  • top_k (int, optional) – Number of top tools to return. Defaults to 5.

Returns:

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

Return type:

list

Raises:

SystemExit – If tool_desc_embedding is not loaded

find_tools(message=None, picked_tool_names=None, rag_num=5, return_call_result=False, categories=None)[source]

Find relevant tools based on a message or pre-selected tool names.

This method either uses RAG inference to find tools based on a message or processes a list of pre-selected tool names. It filters out special tools and returns tool prompts suitable for use in agent workflows.

Parameters:

  • message (str, optional) – Query message to find tools for. Required if picked_tool_names is None.

  • picked_tool_names (list, optional) – Pre-selected tool names to process. Required if message is None.

  • rag_num (int, optional) – Number of tools to return after filtering. Defaults to 5.

  • return_call_result (bool, optional) – If True, returns both prompts and tool names. Defaults to False.

  • categories (list, optional) – List of tool categories to filter by. Currently not implemented for embedding-based search.

Returns:

  • If return_call_result is False: Tool prompts as a formatted string

  • If return_call_result is True: Tuple of (tool_prompts, tool_names)

Return type:

str or tuple

Raises:

AssertionError – If both message and picked_tool_names are None

run(arguments)[source]

Run the tool finder with given arguments following the standard tool interface.

This is the main entry point for using ToolFinderEmbedding as a standard tool. It extracts parameters from the arguments dictionary and delegates to find_tools().

Parameters:

arguments (dict) – Dictionary containing: - description (str, optional): Query message to find tools for (maps to ‘message’) - limit (int, optional): Number of tools to return (maps to ‘rag_num’). Defaults to 5. - picked_tool_names (list, optional): Pre-selected tool names to process - return_call_result (bool, optional): Whether to return both prompts and names. Defaults to False. - categories (list, optional): List of tool categories to filter by