Using tool calls 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 tutorial, you’ll see 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.

Prerequisites#

Dataiku >= 13.2
LLM Mesh connection to a provider that supports tool calls (tested with OpenAI GPT-4o-mini model; you can find a compatible list here)
Project permissions for “Read project content” and “Write project content”
An SQL dataset pro_customers_sql (a CSV file can be downloaded here) in the flow containing customer data with customer_id and some other columns as headers.

Tip

You could run the code interactively within a Jupyter notebook since this script does not take inputs and outputs, such as a dataset. That might also help with debugging since model outputs are not always guaranteed.

Defining a tool#

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 are 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 type of output it can expect. JSON schemas are also helpful in guiding the LLM when interacting with these tools by providing a clear description.

In this tutorial, we will define two tools: one to retrieve customer information from a database and another to fetch company information from an internet search.

JSON schema for tool definition#

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 information and fetching company information.

As with any LLM workflow, it is important to provide context to the model. 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 settings and context#

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 asks, 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.

Chatting with the LLM to extract parameters#

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'}]

You’ll need to use the extracted request to then call the tool function. Let’s implement a simple Python function (process_tool_calls) that calls other functions that retrieves customer information (get_customer_info) or searches for the company online (search_company_info) based the tool that was called and using the parameters provided. The tool_call_result is added back to the conversation.

Function to retrieve customer information#

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):
    results = DDGS().answers(company_name + " (company)")
    result = "Information found about " + company_name + ": " + results[0]["text"] + "\n" \
        if len(results) > 0 and "text" in results[0] \
        else None
    if not result:
        results = DDGS().answers(company_name)
        result = "Information found about " + company_name + ": " + results[0]["text"] + "\n" \
            if len(results) > 0 and "text" in results[0] \
            else "No information can be found about the company " + company_name
    return result

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.

Recording tool call and output#

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 made use of the tools, you can look at the history of the chat. You could also define a separate function to print the conversation in a more readable format. This function iterates through all messages, including any tool calls or outputs.

Let’s take a closer look at a possible outcome to see how tool calls are structured. Check how roles and responses, including those from tools, are integrated into the conversation.

💻 : 
  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.


👤 : Who are you and what is your purpose?

🤖 : I am an AI assistant designed to provide information and support regarding
customers and companies. My purpose is to assist you by retrieving customer details
from our database or by gathering company information from online resources.
If you have any specific questions or need information about a customer or a company, 
feel free to ask!

👤 : The customer's id is fdouetteau

   🤖 Assistant: Let me look that up
   🔧 Tool call (call_9qCXkF7NAt515MBrjycLuVTz): Calling get_customer_info with {"customer_id":"fdouetteau"}
   🛠️ Tool response:  The customer's name is "Florian Douetteau", holding the position "CEO" at the company named Dataiku


👤 : Find more information about the company from a search.

   🤖 Assistant: Let me look that up
   🔧 Tool call (call_rRlHPcIcbce18eH8Zmhak1vc): Calling get_company_info with {"company_name":"Dataiku"}
   🛠️ Tool response:  Information found about Dataiku: Dataiku is an American artificial 
   intelligence and machine learning company which was founded in 2013. In December 2019, 
   Dataiku announced that CapitalG—the late-stage growth venture capital fund financed by 
   Alphabet Inc.—joined Dataiku as an investor and that it had achieved unicorn status. As 
   of 2021, Dataiku is valued at $4.6 billion. Dataiku currently employs more than 1,000 
   people worldwide between offices in New York, Denver, Washington DC, Los Angeles, Paris, 
   London, Munich, Frankfurt, Sydney, Singapore, Tokyo, and Dubai.

Wrapping Up#

In this tutorial, you learned 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.

This approach might come in handy, when building robust, extensible workflows. Plus, you can worry less about manual message tracking or complex integrations, especially when you need to handle multiple models.

chat.py

Longer code block with full script#

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):
    results = DDGS().answers(company_name + " (company)")
    result = "Information found about " + company_name + ": " + results[0]["text"] + "\n" \
        if len(results) > 0 and "text" in results[0] \
        else None
    if not result:
        results = DDGS().answers(company_name)
        result = "Information found about " + company_name + ": " + results[0]["text"] + "\n" \
            if len(results) > 0 and "text" in results[0] \
            else "No information can be found about the company " + company_name
    return result

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"])

def pretty_print_conversation(messages):
    ROLE_EMOJIS = {
        "system": "💻",
        "user": "👤", 
        "assistant": "🤖",
        "tool": "🛠️"
    }
    
    for message in messages:
        role = message["role"]
        content = message.get("content")
        tool_calls = message.get("toolCalls")
            
        # Print main message content first if it exists
        if content and content != "No content":
            print(f"{ROLE_EMOJIS.get(role, '❓')} : {content}\n")
        
        # Then print any tool calls
        if tool_calls:
            for tc in tool_calls:
                print(f"   🤖 Assistant: Let me look that up")
                print(f"   🔧 Tool call ({tc['id']}): Calling {tc['function']['name']} with {tc['function']['arguments']}")
        elif role == "tool":
            print(f"   🛠️ Tool response:  {str(message['toolOutputs'][0]['output'])}\n\n")

pretty_print_conversation(chat.cq["messages"])

# 💻 : 
#   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.


# 👤 : Who are you and what is your purpose?

# 🤖 : I am an AI assistant designed to provide information and support regarding customers and companies. My purpose is to assist you by retrieving customer details from our database or by gathering company information from online resources. If you have any specific questions or need information about a customer or a company, feel free to ask!

# 👤 : The customer's id is fdouetteau

#    🤖 Assistant: Let me look that up
#    🔧 Tool call (call_9qCXkF7NAt515MBrjycLuVTz): Calling get_customer_info with {"customer_id":"fdouetteau"}
#    🛠️ Tool response:  The customer's name is "Florian Douetteau", holding the position "CEO" at the company named Dataiku


# 👤 : Find more information about the company from a search.

#    🤖 Assistant: Let me look that up
#    🔧 Tool call (call_rRlHPcIcbce18eH8Zmhak1vc): Calling get_company_info with {"company_name":"Dataiku"}
#    🛠️ Tool response:  Information found about Dataiku: Dataiku is an American artificial intelligence and machine learning company which was founded in 2013. In December 2019, Dataiku announced that CapitalG—the late-stage growth venture capital fund financed by Alphabet Inc.—joined Dataiku as an investor and that it had achieved unicorn status. As of 2021, Dataiku is valued at $4.6 billion. Dataiku currently employs more than 1,000 people worldwide between offices in New York, Denver, Washington DC, Los Angeles, Paris, London, Munich, Frankfurt, Sydney, Singapore, Tokyo, and Dubai.