Creating an LLM-based agent that uses multiple tools#

Introduction#

The previous part of this tutorial series showed how to define external tools so an LLM can retrieve data or perform specialized tasks. Now, let’s see how to integrate these tools into an agent capable of handling multi-step queries in a flexible conversation.

In this tutorial, you’ll create an agent to deploy multiple tools to manage user queries in a conversational workflow. Whenever a user asks for information, the agent decides how to respond. It could answer directly, but more often it might have to engage a “tool.” In this example, the agent relies on two tools from the previous tutorial. To recap, one retrieves a user’s name, position, and company based on an ID, while a second tool searches the Internet to find company information.

Agentic workflow#

The workflow of our LLM-based agent is broken down into several steps:

The agent receives a user query and analyzes its contents.
It checks available tools that might help answer the query.
When a tool is needed, the agent calls the corresponding function, waits for the output, and integrates it into its reasoning.
Finally, it produces a coherent response that it returns to the user.

The corresponding and processing functions remain the same as in the previous part and are now part of a chat session involving the agent. Namely, get_customer_details(customer_id), search_company_info(company_name), and process_tool_calls(tool_calls) are called within the agent’s conversation flow.

Using pre-defined tools#

Since the tools are well-defined, they are included as a JSON file: tools.json. Place the file in a location accessible to code in Dataiku. An option is to use </> > Libraries. If you upload the file in a folder named python, you can access it in any Python code in the project.

"Put your file within </> >Libraries"

Creating the agent#

The agent is built around a function create_chat_session which sets up the LLM with the system prompt and tool definitions. You need to provide context to the LLM using a prompt that explains its role and how to use the tools. The code snippet below shows the essential pieces of the agent: system prompt definition, tool import, and conversation flow management. A function create_chat_session focuses on this setup, making sure the LLM is ready to use the appropriate tools.

Defining a chat session#

def create_chat_session():
    """Create a new chat session with tools and prompt"""
    chat = llm.new_completion()
    
    # Import tools from JSON file
    library = project.get_library()
    tools_file = library.get_file("/python/tools.json")
    tools_str = tools_file.read()
    tools_dict = json.loads(tools_str)

    chat.settings["tools"] = tools_dict["tools"]
    
    SYSTEM_PROMPT = """You are a customer information assistant. You can:
    1. Look up customer information in our database
    2. Search for additional information about companies online

    When asked about a customer:
    1. First look up their basic information
    2. If company information is requested, search for additional details
    3. Combine the information into a coherent, single-paragraph response"""
    
    chat.with_message(SYSTEM_PROMPT, role="system")

Accessing tools#

During the next step, process_query() guides the conversation with the user. By structuring requests in a loop, the agent can make multiple calls to tools if the user’s question is complex. For example, if the question involves a customer’s detailed profile along with their company information, the agent can call one function for the customer’s data and another function for the company’s data before returning a final result.

How a single query is resolved#

def process_query(query):
    """Process a user query using the agent"""
    chat = create_chat_session()

    chat.with_message(query, role="user")

    conversation_log = []
    while True:
        response = chat.execute()

        if not response.tool_calls:
            # Final answer received
            chat.with_message(response.text, role="assistant")
            conversation_log.append(f"Final Answer: {response.text}")
            break

        # Handle tool calls
        chat.with_tool_calls(response.tool_calls, role="assistant")
        tool_call_result = process_tool_calls(response.tool_calls)
        chat.with_tool_output(tool_call_result, tool_call_id=response.tool_calls[0]["id"])

        # Log the step
        tool_name = response.tool_calls[0]["function"]["name"]
        tool_args = response.tool_calls[0]["function"]["arguments"]
        conversation_log.append(f"Tool: {tool_name}\nInput: {tool_args}\nResult: {tool_call_result}\n{'-'*50}")
    

Wrapping up#

Here’s what a conversation with the agent could look like:

Example usage#

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)

CONTENT = "Give me all the information you can find about customer with id fdouetteau"
print(process_query(CONTENT))

# Tool: get_customer_info
# Input: {"customer_id":"fdouetteau"}
# Result: The customer's name is "Florian Douetteau", holding the position "CEO" at the company named Dataiku
# --------------------------------------------------
# Tool: get_company_info
# Input: {"company_name":"Dataiku"}
# Result: Information found about Dataiku: Our Story. Dataiku is the leading platform for Everyday AI · 
# Leadership and Team · Over 1000 people work hard to ensure the quality of our product and company as ...
# --------------------------------------------------
# Final Answer: The customer, Florian Douetteau, is the CEO of Dataiku, a leading platform for Everyday AI. 
# Dataiku employs over 1,000 people dedicated to ensuring the quality of their products and services 
# in the AI domain.

Now, you have an LLM-based agent that can handle multi-tool calling.
The system prompt and tool definitions provide a flexible and iterable way of guiding the LLM.
The chat session could be designed to make as many tool calls as needed to retrieve the required information or limit tool usage, according to the user’s needs.

In the next part, you’ll learn how to surface this agent in a browser-based web application.

conversation.py

Longer code block with full script#

import dataiku
from typing import List, Dict
import json
from dataiku import SQLExecutor2
from duckduckgo_search import DDGS

def get_customer_details(customer_id):
    """Get customer information from database"""
    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):
    """Process tool calls and return results"""
    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"])

def create_chat_session():
    """Create a new chat session with tools and prompt"""
    chat = llm.new_completion()
    
    # Import tools from JSON file
    library = project.get_library()
    tools_file = library.get_file("/python/tools.json")
    tools_str = tools_file.read()
    tools_dict = json.loads(tools_str)

    chat.settings["tools"] = tools_dict["tools"]
    
    SYSTEM_PROMPT = """You are a customer information assistant. You can:
    1. Look up customer information in our database
    2. Search for additional information about companies online

    When asked about a customer:
    1. First look up their basic information
    2. If company information is requested, search for additional details
    3. Combine the information into a coherent, single-paragraph response"""
    
    chat.with_message(SYSTEM_PROMPT, role="system")
    return chat

def process_query(query):
    """Process a user query using the agent"""
    chat = create_chat_session()

    chat.with_message(query, role="user")

    conversation_log = []
    while True:
        response = chat.execute()

        if not response.tool_calls:
            # Final answer received
            chat.with_message(response.text, role="assistant")
            conversation_log.append(f"Final Answer: {response.text}")
            break

        # Handle tool calls
        chat.with_tool_calls(response.tool_calls, role="assistant")
        tool_call_result = process_tool_calls(response.tool_calls)
        chat.with_tool_output(tool_call_result, tool_call_id=response.tool_calls[0]["id"])

        # Log the step
        tool_name = response.tool_calls[0]["function"]["name"]
        tool_args = response.tool_calls[0]["function"]["arguments"]
        conversation_log.append(f"Tool: {tool_name}\nInput: {tool_args}\nResult: {tool_call_result}\n{'-'*50}")
    
    return "\n".join(conversation_log)

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)

CONTENT = "Give me all the information you can find about customer with id fdouetteau"
print(process_query(CONTENT))

# Tool: get_customer_info
# Input: {"customer_id":"fdouetteau"}
# Result: The customer's name is "Florian Douetteau", holding the position "CEO" at the company named Dataiku
# --------------------------------------------------
# Tool: get_company_info
# Input: {"company_name":"Dataiku"}
# Result: Information found about Dataiku: Our Story. Dataiku is the leading platform for Everyday AI · 
# Leadership and Team · Over 1000 people work hard to ensure the quality of our product and company as ...
# --------------------------------------------------
# Final Answer: The customer, Florian Douetteau, is the CEO of Dataiku, a leading platform for Everyday AI. 
# Dataiku employs over 1,000 people dedicated to ensuring the quality of their products and services 
# in the AI domain.