Controlling Agent Reasoning Loop with Return Direct Tools¶

All tools have an option for return_direct -- if this is set to True, and the associated tool is called (without any other tools being called), the agent reasoning loop is ended and the tool output is returned directly.

This can be useful for speeding up response times when you know the tool output is good enough, to avoid the agent re-writing the response, and for ending the reasoning loop.

This notebook walks through a notebook where an agent needs to gather information from a user in order to make a restaurant booking.

In [ ]:

!pip install llama-index-core llama-index-llms-anthropic

!pip install llama-index-core llama-index-llms-anthropic

In [ ]:

import os

os.environ["ANTHROPIC_API_KEY"] = "sk-ant-..."

import os os.environ["ANTHROPIC_API_KEY"] = "sk-ant-..."

Tools setup¶

In [ ]:

from typing import Optional

from llama_index.core.tools import FunctionTool
from llama_index.core.bridge.pydantic import BaseModel

# we will store booking under random IDs
bookings = {}


# we will represent and track the state of a booking as a Pydantic model
class Booking(BaseModel):
    name: Optional[str] = None
    email: Optional[str] = None
    phone: Optional[str] = None
    date: Optional[str] = None
    time: Optional[str] = None


def get_booking_state(user_id: str) -> str:
    """Get the current state of a booking for a given booking ID."""
    try:
        return str(bookings[user_id].dict())
    except:
        return f"Booking ID {user_id} not found"


def update_booking(user_id: str, property: str, value: str) -> str:
    """Update a property of a booking for a given booking ID. Only enter details that are explicitly provided."""
    booking = bookings[user_id]
    setattr(booking, property, value)
    return f"Booking ID {user_id} updated with {property} = {value}"


def create_booking(user_id: str) -> str:
    """Create a new booking and return the booking ID."""
    bookings[user_id] = Booking()
    return "Booking created, but not yet confirmed. Please provide your name, email, phone, date, and time."


def confirm_booking(user_id: str) -> str:
    """Confirm a booking for a given booking ID."""
    booking = bookings[user_id]

    if booking.name is None:
        raise ValueError("Please provide your name.")

    if booking.email is None:
        raise ValueError("Please provide your email.")

    if booking.phone is None:
        raise ValueError("Please provide your phone number.")

    if booking.date is None:
        raise ValueError("Please provide the date of your booking.")

    if booking.time is None:
        raise ValueError("Please provide the time of your booking.")

    return f"Booking ID {user_id} confirmed!"


# create tools for each function
get_booking_state_tool = FunctionTool.from_defaults(fn=get_booking_state)
update_booking_tool = FunctionTool.from_defaults(fn=update_booking)
create_booking_tool = FunctionTool.from_defaults(
    fn=create_booking, return_direct=True
)
confirm_booking_tool = FunctionTool.from_defaults(
    fn=confirm_booking, return_direct=True
)

from typing import Optional from llama_index.core.tools import FunctionTool from llama_index.core.bridge.pydantic import BaseModel # we will store booking under random IDs bookings = {} # we will represent and track the state of a booking as a Pydantic model class Booking(BaseModel): name: Optional[str] = None email: Optional[str] = None phone: Optional[str] = None date: Optional[str] = None time: Optional[str] = None def get_booking_state(user_id: str) -> str: """Get the current state of a booking for a given booking ID.""" try: return str(bookings[user_id].dict()) except: return f"Booking ID {user_id} not found" def update_booking(user_id: str, property: str, value: str) -> str: """Update a property of a booking for a given booking ID. Only enter details that are explicitly provided.""" booking = bookings[user_id] setattr(booking, property, value) return f"Booking ID {user_id} updated with {property} = {value}" def create_booking(user_id: str) -> str: """Create a new booking and return the booking ID.""" bookings[user_id] = Booking() return "Booking created, but not yet confirmed. Please provide your name, email, phone, date, and time." def confirm_booking(user_id: str) -> str: """Confirm a booking for a given booking ID.""" booking = bookings[user_id] if booking.name is None: raise ValueError("Please provide your name.") if booking.email is None: raise ValueError("Please provide your email.") if booking.phone is None: raise ValueError("Please provide your phone number.") if booking.date is None: raise ValueError("Please provide the date of your booking.") if booking.time is None: raise ValueError("Please provide the time of your booking.") return f"Booking ID {user_id} confirmed!" # create tools for each function get_booking_state_tool = FunctionTool.from_defaults(fn=get_booking_state) update_booking_tool = FunctionTool.from_defaults(fn=update_booking) create_booking_tool = FunctionTool.from_defaults( fn=create_booking, return_direct=True ) confirm_booking_tool = FunctionTool.from_defaults( fn=confirm_booking, return_direct=True )

A user has walked in! Let's help them make a booking¶

In [ ]:

from llama_index.llms.anthropic import Anthropic
from llama_index.core.llms import ChatMessage
from llama_index.core.agent import FunctionCallingAgent

llm = Anthropic(model="claude-3-sonnet-20240229", temperature=0.1)

user = "user123"
prefix_messages = [
    ChatMessage(
        role="system",
        content=(
            f"You are now connected to the booking system and helping {user} with making a booking. "
            "Only enter details that the user has explicitly provided. "
            "Do not make up any details."
        ),
    )
]

agent = FunctionCallingAgent.from_tools(
    tools=[
        get_booking_state_tool,
        update_booking_tool,
        create_booking_tool,
        confirm_booking_tool,
    ],
    llm=llm,
    prefix_messages=prefix_messages,
    max_function_calls=10,
    allow_parallel_tool_calls=False,
    verbose=True,
)

from llama_index.llms.anthropic import Anthropic from llama_index.core.llms import ChatMessage from llama_index.core.agent import FunctionCallingAgent llm = Anthropic(model="claude-3-sonnet-20240229", temperature=0.1) user = "user123" prefix_messages = [ ChatMessage( role="system", content=( f"You are now connected to the booking system and helping {user} with making a booking. " "Only enter details that the user has explicitly provided. " "Do not make up any details." ), ) ] agent = FunctionCallingAgent.from_tools( tools=[ get_booking_state_tool, update_booking_tool, create_booking_tool, confirm_booking_tool, ], llm=llm, prefix_messages=prefix_messages, max_function_calls=10, allow_parallel_tool_calls=False, verbose=True, )

In [ ]:

response = agent.chat("Hello! I would like to make a booking, around 5pm?")

response = agent.chat("Hello! I would like to make a booking, around 5pm?")

Added user message to memory: Hello! I would like to make a booking, around 5pm?
=== LLM Response ===
Okay, let's create a new booking for you. To do that, I'll invoke the create_booking tool with your user ID:
=== Calling Function ===
Calling function: create_booking with args: {"user_id": "user123"}
=== Function Output ===
Booking created, but not yet confirmed. Please provide your name, email, phone, date, and time.

In [ ]:

print(str(response))

print(str(response))

Booking created, but not yet confirmed. Please provide your name, email, phone, date, and time.

Perfect, we can see the function output was retruned directly, with no modification or final LLM call!

In [ ]:

response = agent.chat("Sure! My name is Logan, and my email is [email protected]")

response = agent.chat("Sure! My name is Logan, and my email is [email protected]")

Added user message to memory: Sure! My name is Logan, and my email is [email protected]
=== LLM Response ===
Got it, thanks for providing your name and email. Let me update the booking with those details:
=== Calling Function ===
Calling function: update_booking with args: {"user_id": "user123", "property": "name", "value": "Logan"}
=== Function Output ===
Booking ID user123 updated with name = Logan
=== Calling Function ===
Calling function: update_booking with args: {"user_id": "user123", "property": "email", "value": "[email protected]"}
=== Function Output ===
Booking ID user123 updated with email = [email protected]
=== LLM Response ===
I still need your phone number, date, and preferred time for the booking. Please provide those details.

In [ ]:

print(str(response))

print(str(response))

assistant: I still need your phone number, date, and preferred time for the booking. Please provide those details.

In [ ]:

response = agent.chat(
    "Right! My phone number is 1234567890, the date of the booking is April 5, at 5pm."
)

response = agent.chat( "Right! My phone number is 1234567890, the date of the booking is April 5, at 5pm." )

Added user message to memory: Right! My phone number is 1234567890, the date of the booking is April 5, at 5pm.
=== LLM Response ===
Great, thank you for providing the remaining details. Let me update the booking:
=== Calling Function ===
Calling function: update_booking with args: {"user_id": "user123", "property": "phone", "value": "1234567890"}
=== Function Output ===
Booking ID user123 updated with phone = 1234567890
=== Calling Function ===
Calling function: update_booking with args: {"user_id": "user123", "property": "date", "value": "2023-04-05"}
=== Function Output ===
Booking ID user123 updated with date = 2023-04-05
=== Calling Function ===
Calling function: update_booking with args: {"user_id": "user123", "property": "time", "value": "17:00"}
=== Function Output ===
Booking ID user123 updated with time = 17:00
=== LLM Response ===
Your booking for April 5th at 5pm is now created with all the provided details. To confirm it, I'll invoke:
=== Calling Function ===
Calling function: confirm_booking with args: {"user_id": "user123"}
=== Function Output ===
Booking ID user123 confirmed!

In [ ]:

print(str(response))

print(str(response))

Booking ID user123 confirmed!

In [ ]:

print(bookings["user123"])

print(bookings["user123"])

name='Logan' email='[email protected]' phone='1234567890' date='2023-04-05' time='17:00'