GenAI: Select a tool/agent based on a user query

Photo by Dan Cristian Pădureț on Unsplash

Article about how to select a tool to handle a user query.

The fast and basic way: using classification based on semantic similarity

First we start by classifying a query

This example is based on a HuggingFace pipeline for zero-shot classification (https://huggingface.co/facebook/bart-large-mnli)

We will start by building a dictionary of some query example and the expected tool name to handle them. (Here there are two tools: ingestion and weather)

First example

In the result we can see there is 99.40% of probability the query “I want to learn about RAG” has the same meaning than “I want to learn about something”

Second example

In the second example the results are not as good, the closest query example for “What will be the weather for tomorrow at Paris?” is “I want to learn about something” with a 40.15% of probability. However we an also consider that the total probability of the query not being “I want to learn about something” and corresponding to the examples for the weather tool is 59.85%, which is higher

Then we can map the classification result with the tool-name

To have the name of the tool to use, we have to retrieve it from the values of the previously declared tools dictionary.

We should also take into account that multiple classification can be proposed for the same

The classification to tool-name with the results from the second query example

The “first_item” variable is a tuple with the tool_name as first value and the probability it is the good one as the second value.

The full example code

from transformers import pipeline

tools = {
    "I want to learn about something": "ingestion", 
    "I want to know the weather for a date and a place": "weather",
    "What will be the weather for a date and a place": "weather"
}

pipe = pipeline(model="facebook/bart-large-mnli")

result = pipe(
    "What will be the weather for tomorrow at Paris?",
    candidate_labels=list(tools.keys()),
)

print(result)

tools_result = {}

labels = result["labels"]
scores = result["scores"]

for label_index, label in enumerate(labels):
    tool_name = tools[label]
    tools_result[tool_name] = tools_result.get(tool_name, 0) + scores[label_index]
    
tools_result = dict(sorted(tools_result.items(), key=lambda item: item[1], reverse=True))

print(tools_result)

first_item = next(iter(tools_result.items()))

print(first_item)

The pros and cons

The pros:

• it is easy to implement it,
• it does not need access to internet to run (except for downloading the model for the first call) / no call costs,
• there is no limit to the number of tool/classes you can use (but it will result in a longer execution time)

The cons:

• it is best suited in the case the user already knows the kind of query that can be used with the tool, which limit the usability,
• it cannot extract informations from the query or rephrase it (it must be done later in a next step).

A bit more complex way, using a Large Language Model to generate an answer

Both LlamaIndex and LangChain provide ways to select a tool using a Large Language Model

Selecting a single tool with LlamaIndex LLMSingleSelector

Example with first query

Example with second query

As you can see in the examples, you can get both the index of the tool to use and a reason what this tool is proposed to be used.

from llama_index.tools import ToolMetadata
from llama_index.selectors import LLMSingleSelector

ingestion_tool = ToolMetadata(name="ingestion", description="Useful to learn about a subject")
weather_tool = ToolMetadata(name="weather", description="Useful to get the weather")

os.environ["OPENAI_API_KEY"] = "sk-..."

selector = LLMSingleSelector.from_defaults()

result = selector.select([ingestion_tool, weather_tool], "What will be the weather for tomorrow at Paris?")

print(result)

Other LLM-Based tool selection mechanisms (LlamaIndex)

from llama_index.selectors import LLMMultiSelector

To select multiple tools from the query, same than LLMSingleSelector but with multiple results possibles

from llama_index.selectors import EmbeddingSingleSelector

This LlamaIndex selector is similar to the HuggingFace based method proposed at the start of this article, it is based on text semantic similarity.

LangChain also has possibilities for tool calling

Choosing between multiple tools | 🦜️🔗 LangChain

I won’t enter in the detail there, as this article is most focused on what is possible than in exploring all the available mechanisms to select one/many tool.

For the LangGraph part, you can take a look at:

🦜🕸️LangGraph | 🦜️🔗 LangChain

Last solution we will propose, training a sentence classification model

Update on September 28, 2024.

I have talked about how to fine tune a model for text classification in the following article:

(Hugging Face) — Text classification, going further than the tutorial

The idea is quite the same than using a zero-shot classification model, but this time we train the model with examples for each class in order for him to be able to accurately find out which class a text belongs to.

The pros and cons

The pros are the same than using a zero-shot classification model and with a better accuracy.

The cons are the same than using a zero-shot classification model with the time spent for training the model and the need to have a dataset prepared to do so.

Some use cases

Action based cases

The tool that is selected can perform an actual action/task or series of tasks, like ingesting new articles in a knowledge base.

(Which I already talked about in this previous article)

AI: dynamic task scheduling in the context of advanced chatbots — a case study

Retrieval-Augmented Generation linked cases

• select the tool/retriever most suited to handle the user query,
• from a user query and a list of tools description, ask a LLM to generate a list of tool-name + sub-question pairs, it can help answer a complex query by dividing it in different sub-question and answering them before synthesizing a final answer.
• ReACT: perform another Retrieval-Augmented Generation step if a LLM call with the previous result determine the question isn’t fully answered.
• … and probably even more… !

---

Thanks for reading till the end!

Feel free to clap, comment or contact me if you want!