Defining and using tools with the LLM Mesh#
Introduction#
Large Language Models (LLMs) are incredibly versatile in understanding and generating human-like text. Yet they still have notable limitations. When LLMs struggle with operations requiring precise mathematical calculations or updated data, external tools can complement them. Tools are predefined functions that an LLM can call during a conversation to solve specific problems, like performing calculations or querying databases.
In this first part of a series of tutorials, you’ll learn how to integrate tools into your workflows using the LLM Mesh, a centralized and governed interface for accessing models from multiple providers. You’ll define tools, implement tool calls, and see how the LLM interacts with these tools.
Defining Tools#
First, let’s define a tool in the context of an LLM. Tools are functions invoked during a conversation to perform a predefined task. Like functions, they accept specific parameters and return a response via a process outside the LLM workflow, which is then used as part of the conversation.
Tools can be defined using a JSON schema, specifying parameters and their types. This schema helps the LLM understand what kind of input the tool requires and what output it can expect. JSON schemas are also helpful by providing clear, human-readable descriptions and metadata.
In this tutorial, you will define two tools:
Customer Information Tool - retrieves customer details from a database
Company Information Tool - searches the internet for company information
tools = [
{
"type": "function",
"function": {
"name": "get_customer_info",
"description": "Get customer details from the database given their ID",
"parameters": {
"type": "object",
"properties": {
"customer_id": {
"type": "string",
"description": "The unique identifier for the customer",
},
},
"required": ["customer_id"],
},
}
},
{
"type": "function",
"function": {
"name": "get_company_info",
"description": "Get company information from internet search",
"parameters": {
"type": "object",
"properties": {
"company_name": {
"type": "string",
"description": "Name of the company to search for",
},
},
"required": ["company_name"],
},
}
}
]
Adding the tool with an LLM Mesh workflow#
You’ll create a client to interact with a specific LLM Mesh connection. Next, you’ll start a new chat, which is basically a completion task for the model based on user and tool inputs. To that chat, you’ll add tools for retrieving customer data and fetching company information.
As with any LLM workflow, providing context to the model is essential. This context defines the model’s role as well as what it can expect from the user, the tool and how to interpret the inputs it receives. You’ll specify that the LLM acts as a helpful assistant with access to customer and company information.
chat = llm.new_completion()
chat.settings["tools"] = tools
CONTEXT = '''
You are a helpful assistant with access to customer and company information.
You have two tools available:
- get_customer_info: retrieves customer details from our database
- get_company_info: searches the internet for company information
Use these tools to provide comprehensive responses about customers and their companies.
'''
CONTENT = 'Who are you and what is your purpose?'
chat.with_message(CONTEXT, role="system")
chat.with_message(CONTENT, role="user")
response = chat.execute()
response.text
# I am an AI assistant designed to help you gather information about customers
# and companies. My main functions include retrieving customer details from our
# database and searching the internet for company information. If you have
# specific questions or tasks related to customers or companies, feel free to
# ask!
chat.with_message(response.text, role="assistant")
Chatting with the LLM to use the tool#
Once the model shows it understands these questions, you can ask it to retrieve customer information. At this point, the LLM decides whether it needs to call the tool to help with the user’s request. Once it does, the model calls the tool as a response. The parameters needed to retrieve the information are extracted from the conversation, i.e., from the user’s statement.
customer_id = "fdouetteau"
CONTENT = f"The customer's id is {customer_id}"
chat.with_message(CONTENT, role="user")
response = chat.execute()
tool_calls = response.tool_calls
tool_calls
# [{'type': 'function',
# 'function': {'name': 'get_customer_info',
# 'arguments': '{"customer_id":"fdouetteau"}'},
# 'id': 'call_cQECgVOCgU7OLLb5mrBOZrg5'}]
Then, you’ll need to use the extracted request to call the tool function. Let’s
implement a simple Python function (process_tool_calls) that calls other
functions that retrieve customer information (get_customer_info) or search
for the company online (search_company_info) based on the tool that was called
and using the parameters provided. The tool_call_result is added back to the
conversation.
def get_customer_details(customer_id):
dataset = dataiku.Dataset("pro_customers_sql")
table_name = dataset.get_location_info().get('info', {}).get('table')
executor = SQLExecutor2(dataset=dataset)
customer_id = customer_id.replace("'", "\\'")
query_reader = executor.query_to_iter(
f"""SELECT name, job, company FROM "{table_name}" WHERE id = '{customer_id}'""")
for (name, job, company) in query_reader.iter_tuples():
return f"The customer's name is \"{name}\", holding the position \"{job}\" at the company named {company}"
return f"No information can be found about the customer {customer_id}"
def search_company_info(company_name):
"""Search for company information online"""
with DDGS() as ddgs:
results = list(ddgs.text(f"{company_name} (company)", max_results=1))
if results:
return f"Information found about {company_name}: {results[0]['body']}"
return f"No information found about {company_name}"
def process_tool_calls(tool_calls):
tool_name = tool_calls[0]["function"]["name"]
llm_args = json.loads(tool_calls[0]["function"]["arguments"])
if tool_name == "get_customer_info":
return get_customer_details(llm_args["customer_id"])
elif tool_name == "get_company_info":
return search_company_info(llm_args["company_name"])
Each time the tool is called, the tool call, and the function output are logged in the conversation history. This helps the LLM record what was executed and integrate the tool workflow into the conversation.
chat.with_tool_calls(tool_calls, role="assistant")
tool_call_result = process_tool_calls(tool_calls)
chat.with_tool_output(tool_call_result, tool_call_id=tool_calls[0]["id"])
Checking the results#
To verify whether and how the LLM used the tools, you can look at the
history of the chat. To access the entire history, you can use
chat.cq["messages"].
Let’s take a closer look at a possible outcome to see how tool calls are structured.
# Chat history
from pprint import pprint
pprint(chat.cq["messages"], indent=2, width=80)
# [ { 'content': '\n'
# ' You are a helpful assistant with access to customer and '
# 'company information.\n'
# ' You have two tools available:\n'
# ' - get_customer_info: retrieves customer details from our '
# 'database\n'
# ' - get_company_info: searches the internet for company '
# 'information\n'
# ' Use these tools to provide comprehensive responses about '
# 'customers and their companies.\n',
# 'role': 'system'},
# {'content': 'Who are you and what is your purpose?', 'role': 'user'},
# { 'content': 'I am an AI assistant designed to help you retrieve information '
# 'about customers and companies. My purpose is to provide '
# 'comprehensive and accurate responses based on the data '
# 'available in our database and from online resources. Whether '
# "you need customer details or company information, I'm here to "
# 'assist you. How can I help you today?',
# 'role': 'assistant'},
# {'content': "The customer's id is fdouetteau", 'role': 'user'},
# { 'role': 'assistant',
# 'toolCalls': [ { 'function': { 'arguments': '{"customer_id":"fdouetteau"}',
# 'name': 'get_customer_info'},
# 'id': 'call_OoXImpZSMNYEqe5w1QT8eFUy',
# 'type': 'function'}]},
# { 'role': 'tool',
# 'toolOutputs': [ { 'callId': 'call_OoXImpZSMNYEqe5w1QT8eFUy',
# 'output': 'The customer\'s name is "Florian Douetteau", '
# 'holding the position "CEO" at the company '
# 'named Dataiku'}]},
# { 'content': 'Find more information about the company from a search.',
# 'role': 'user'},
# { 'role': 'assistant',
# 'toolCalls': [ { 'function': { 'arguments': '{"company_name":"Dataiku"}',
# 'name': 'get_company_info'},
# 'id': 'call_9KerL9juzQMJG3s8FN31yOdo',
# 'type': 'function'}]},
# { 'role': 'tool',
# 'toolOutputs': [ { 'callId': 'call_9KerL9juzQMJG3s8FN31yOdo',
# 'output': 'Information found about Dataiku: Dataiku is '
# 'the leading platform for Everyday AI ... '
# "We're pioneering “Everyday AI,” helping "
# 'everyone in an organization — from technical '
# 'teams to business leaders\xa0...'}]}]
Wrapping Up#
This tutorial, teaches you how to use tool calls via the LLM Mesh. By defining and implementing tools, the LLM seamlessly integrates additional functionality like querying databases or performing computations. The LLM Mesh also manages context and message history.
When building robust, extensible workflows, this approach might come in handy. Plus, with the LLM Mesh, you can worry less about manual message tracking or complex integrations, especially when handling multiple models.
The next tutorial in this series covers how to create an LLM-based agent that uses the multiple tools defined in this tutorial.
chat.py
import dataiku
import json
from dataiku import SQLExecutor2
from duckduckgo_search import DDGS
PROJECT = "" # DSS project key goes here
client = dataiku.api_client()
project = client.get_project(PROJECT)
LLM_ID = "" # LLM ID for the LLM Mesh connection + model goes here
llm = project.get_llm(LLM_ID)
chat = llm.new_completion()
tools = [
{
"type": "function",
"function": {
"name": "get_customer_info",
"description": "Get customer details from the database given their ID",
"parameters": {
"type": "object",
"properties": {
"customer_id": {
"type": "string",
"description": "The unique identifier for the customer",
},
},
"required": ["customer_id"],
},
}
},
{
"type": "function",
"function": {
"name": "get_company_info",
"description": "Get company information from internet search",
"parameters": {
"type": "object",
"properties": {
"company_name": {
"type": "string",
"description": "Name of the company to search for",
},
},
"required": ["company_name"],
},
}
}
]
chat.settings["tools"] = tools
CONTEXT = '''
You are a helpful assistant with access to customer and company information.
You have two tools available:
- get_customer_info: retrieves customer details from our database
- get_company_info: searches the internet for company information
Use these tools to provide comprehensive responses about customers and their companies.
'''
CONTENT = 'Who are you and what is your purpose?'
chat.with_message(CONTEXT, role="system")
chat.with_message(CONTENT, role="user")
response = chat.execute()
response.text
# I am an AI assistant designed to help you gather information about customers
# and companies. My main functions include retrieving customer details from our
# database and searching the internet for company information. If you have
# specific questions or tasks related to customers or companies, feel free to
# ask!
chat.with_message(response.text, role="assistant")
customer_id = "fdouetteau"
CONTENT = f"The customer's id is {customer_id}"
chat.with_message(CONTENT, role="user")
response = chat.execute()
tool_calls = response.tool_calls
tool_calls
# [{'type': 'function',
# 'function': {'name': 'get_customer_info',
# 'arguments': '{"customer_id":"fdouetteau"}'},
# 'id': 'call_cQECgVOCgU7OLLb5mrBOZrg5'}]
chat.with_tool_calls(tool_calls, role="assistant")
def get_customer_details(customer_id):
dataset = dataiku.Dataset("pro_customers_sql")
table_name = dataset.get_location_info().get('info', {}).get('table')
executor = SQLExecutor2(dataset=dataset)
customer_id = customer_id.replace("'", "\\'")
query_reader = executor.query_to_iter(
f"""SELECT name, job, company FROM "{table_name}" WHERE id = '{customer_id}'""")
for (name, job, company) in query_reader.iter_tuples():
return f"The customer's name is \"{name}\", holding the position \"{job}\" at the company named {company}"
return f"No information can be found about the customer {customer_id}"
def search_company_info(company_name):
"""Search for company information online"""
with DDGS() as ddgs:
results = list(ddgs.text(f"{company_name} (company)", max_results=1))
if results:
return f"Information found about {company_name}: {results[0]['body']}"
return f"No information found about {company_name}"
def process_tool_calls(tool_calls):
tool_name = tool_calls[0]["function"]["name"]
llm_args = json.loads(tool_calls[0]["function"]["arguments"])
if tool_name == "get_customer_info":
return get_customer_details(llm_args["customer_id"])
elif tool_name == "get_company_info":
return search_company_info(llm_args["company_name"])
tool_call_result = process_tool_calls(tool_calls)
chat.with_tool_output(tool_call_result, tool_call_id=tool_calls[0]["id"])
# Continue the conversation
CONTENT = "Find more information about the company from a search."
chat.with_message(CONTENT, role="user")
response = chat.execute()
tool_calls = response.tool_calls
tool_calls
# [{'type': 'function',
# 'function': {'name': 'get_company_info',
# 'arguments': '{"company_name":"Dataiku"}'},
# 'id': 'call_4lg3yspLrMdJvISnL2aBtzfn'}]
chat.with_tool_calls(tool_calls, role="assistant")
tool_call_result = process_tool_calls(tool_calls)
chat.with_tool_output(tool_call_result, tool_call_id=tool_calls[0]["id"])
# Chat history
from pprint import pprint
pprint(chat.cq["messages"], indent=2, width=80)
# [ { 'content': '\n'
# ' You are a helpful assistant with access to customer and '
# 'company information.\n'
# ' You have two tools available:\n'
# ' - get_customer_info: retrieves customer details from our '
# 'database\n'
# ' - get_company_info: searches the internet for company '
# 'information\n'
# ' Use these tools to provide comprehensive responses about '
# 'customers and their companies.\n',
# 'role': 'system'},
# {'content': 'Who are you and what is your purpose?', 'role': 'user'},
# { 'content': 'I am an AI assistant designed to help you retrieve information '
# 'about customers and companies. My purpose is to provide '
# 'comprehensive and accurate responses based on the data '
# 'available in our database and from online resources. Whether '
# "you need customer details or company information, I'm here to "
# 'assist you. How can I help you today?',
# 'role': 'assistant'},
# {'content': "The customer's id is fdouetteau", 'role': 'user'},
# { 'role': 'assistant',
# 'toolCalls': [ { 'function': { 'arguments': '{"customer_id":"fdouetteau"}',
# 'name': 'get_customer_info'},
# 'id': 'call_OoXImpZSMNYEqe5w1QT8eFUy',
# 'type': 'function'}]},
# { 'role': 'tool',
# 'toolOutputs': [ { 'callId': 'call_OoXImpZSMNYEqe5w1QT8eFUy',
# 'output': 'The customer\'s name is "Florian Douetteau", '
# 'holding the position "CEO" at the company '
# 'named Dataiku'}]},
# { 'content': 'Find more information about the company from a search.',
# 'role': 'user'},
# { 'role': 'assistant',
# 'toolCalls': [ { 'function': { 'arguments': '{"company_name":"Dataiku"}',
# 'name': 'get_company_info'},
# 'id': 'call_9KerL9juzQMJG3s8FN31yOdo',
# 'type': 'function'}]},
# { 'role': 'tool',
# 'toolOutputs': [ { 'callId': 'call_9KerL9juzQMJG3s8FN31yOdo',
# 'output': 'Information found about Dataiku: Dataiku is '
# 'the leading platform for Everyday AI ... '
# "We're pioneering “Everyday AI,” helping "
# 'everyone in an organization — from technical '
# 'teams to business leaders\xa0...'}]}]