Build your own OpenAI Agent

With the new OpenAI API that supports function calling, it’s never been easier to build your own agent!

In this notebook tutorial, we showcase how to write your own OpenAI agent in under 50 lines of code! It is minimal, yet feature complete (with ability to carry on a conversation and use tools).

Initial Setup

Let’s start by importing some simple building blocks.

The main thing we need is:

  1. the OpenAI API (using our own llama_index LLM class)

  2. a place to keep conversation history

  3. a definition for tools that our agent can use.

import json
from typing import Sequence, List

from llama_index.llms import OpenAI, ChatMessage
from llama_index.tools import BaseTool, FunctionTool

import nest_asyncio

nest_asyncio.apply()

Let’s define some very simple calculator tools for our agent.

def multiply(a: int, b: int) -> int:
    """Multiple two integers and returns the result integer"""
    return a * b


multiply_tool = FunctionTool.from_defaults(fn=multiply)

def add(a: int, b: int) -> int:
    """Add two integers and returns the result integer"""
    return a + b


add_tool = FunctionTool.from_defaults(fn=add)

Agent Definition

Now, we define our agent that’s capable of holding a conversation and calling tools in under 50 lines of code.

The meat of the agent logic is in the chat method. At a high-level, there are 3 steps:

  1. Call OpenAI to decide which tool (if any) to call and with what arguments.

  2. Call the tool with the arguments to obtain an output

  3. Call OpenAI to synthesize a response from the conversation context and the tool output.

The reset method simply resets the conversation context, so we can start another conversation.

class YourOpenAIAgent:
    def __init__(
        self,
        tools: Sequence[BaseTool] = [],
        llm: OpenAI = OpenAI(temperature=0, model="gpt-3.5-turbo-0613"),
        chat_history: List[ChatMessage] = [],
    ) -> None:
        self._llm = llm
        self._tools = {tool.metadata.name: tool for tool in tools}
        self._chat_history = chat_history

    def reset(self) -> None:
        self._chat_history = []

    def chat(self, message: str) -> str:
        chat_history = self._chat_history
        chat_history.append(ChatMessage(role="user", content=message))
        functions = [
            tool.metadata.to_openai_function() for _, tool in self._tools.items()
        ]

        ai_message = self._llm.chat(chat_history, functions=functions).message
        chat_history.append(ai_message)

        function_call = ai_message.additional_kwargs.get("function_call", None)
        if function_call is not None:
            function_message = self._call_function(function_call)
            chat_history.append(function_message)
            ai_message = self._llm.chat(chat_history).message
            chat_history.append(ai_message)

        return ai_message.content

    def _call_function(self, function_call: dict) -> ChatMessage:
        tool = self._tools[function_call["name"]]
        output = tool(**json.loads(function_call["arguments"]))
        return ChatMessage(
            name=function_call["name"],
            content=str(output),
            role="function",
            additional_kwargs={"name": function_call["name"]},
        )

Let’s Try It Out!

agent = YourOpenAIAgent(tools=[multiply_tool, add_tool])

agent.chat("Hi")

'Hello! How can I assist you today?'

agent.chat("What is 2123 * 215123")

'The product of 2123 multiplied by 215123 is 456,706,129.'

Our (Slightly Better) OpenAIAgent Implementation

We provide a (slightly better) OpenAIAgent implementation in LlamaIndex, which you can directly use as follows.

In comparison to the simplified version above:

  • it implements the BaseChatEngine and BaseQueryEngine interface, so you can more seamlessly use it in the LlamaIndex framework.

  • it supports multiple function calls per conversation turn

  • it supports streaming

  • it supports async endpoints

  • it supports callback and tracing

from llama_index.agent import OpenAIAgent
from llama_index.llms import OpenAI

llm = OpenAI(model="gpt-3.5-turbo-0613")
agent = OpenAIAgent.from_tools([multiply_tool, add_tool], llm=llm, verbose=True)

Chat

response = agent.chat("What is (121 * 3) + 42?")
print(str(response))

=== Calling Function ===
Calling function: multiply with args: {
  "a": 121,
  "b": 3
}
Got output: 363
========================
=== Calling Function ===
Calling function: add with args: {
  "a": 363,
  "b": 42
}
Got output: 405
========================
(121 * 3) + 42 is equal to 405.

# inspect sources
print(response.sources)

[ToolOutput(content='363', tool_name='multiply', raw_input={'args': (), 'kwargs': {'a': 121, 'b': 3}}, raw_output=363), ToolOutput(content='405', tool_name='add', raw_input={'args': (), 'kwargs': {'a': 363, 'b': 42}}, raw_output=405)]

Async Chat

response = await agent.achat("What is 121 * 3?")
print(str(response))

=== Calling Function ===
Calling function: multiply with args: {
  "a": 121,
  "b": 3
}
Got output: 363
========================
121 * 3 is equal to 363.

Streaming Chat

Here, every LLM response is returned as a generator. You can stream every incremental step, or only the last response.

response = agent.stream_chat(
    "What is 121 * 2? Once you have the answer, use that number to write a story about a group of mice."
)

response_gen = response.response_gen

for token in response_gen:
    print(token, end="")

=== Calling Function ===
Calling function: multiply with args: {
  "a": 121,
  "b": 2
}
Got output: 242
========================
121 * 2 is equal to 242.

Once upon a time, in a small village, there was a group of mice who lived happily in a cozy little burrow. The leader of the group was a wise and courageous mouse named Milo. Milo was known for his intelligence and his ability to solve problems.

One day, Milo gathered all the mice together and announced that they needed to find a new home. Their current burrow had become overcrowded, and they needed more space to live comfortably. The mice were excited about the idea of exploring new territories.

With their tiny paws and keen senses, the mice set out on their journey. They traveled through fields, forests, and streams, searching for the perfect place to call home. Along the way, they encountered various challenges, such as crossing treacherous rivers and avoiding hungry predators.

After days of searching, the mice stumbled upon a hidden meadow surrounded by tall grass and blooming flowers. It was a peaceful and serene place, far away from the hustle and bustle of the village. The mice knew they had found their new home.

Using their collective strength and determination, the mice began building their new burrow. They dug tunnels and created intricate chambers, ensuring that each mouse had enough space to live comfortably. Milo, with his exceptional leadership skills, organized the mice into different teams, assigning tasks to each member.

As the mice settled into their new home, they realized that they had created a harmonious community. They worked together, sharing food, and looking out for one another. Milo's wisdom and guidance helped them overcome any obstacles they faced.

The mice flourished in their new meadow, living happily ever after. They grew in numbers and became known as the Meadow Mice, admired by other animals for their unity and resilience. Milo's legacy lived on, as he continued to lead and inspire the mice for generations to come.

And so, the story of the group of mice who found their new home after multiplying their efforts by 121 * 2 became a tale of courage, teamwork, and the power of determination.

Async Streaming Chat

response = await agent.astream_chat(
    "What is 121 + 8? Once you have the answer, use that number to write a story about a group of mice."
)

response_gen = response.response_gen

async for token in response.async_response_gen():
    print(token, end="")

=== Calling Function ===
Calling function: add with args: {
  "a": 121,
  "b": 8
}
Got output: 129
========================
121 + 8 is equal to 129.

Once upon a time, in a lush green forest, there was a group of mice who lived in harmony. They were known as the Forest Friends, and their leader was a wise and kind-hearted mouse named Oliver.

One sunny day, as the mice were going about their daily activities, they stumbled upon a mysterious object hidden beneath a pile of leaves. It was a magical acorn, shimmering with a golden glow. The acorn had the power to grant a wish to whoever possessed it.

Excited by the discovery, Oliver gathered all the mice together and shared the news. They decided to use the wish to make their forest home even more beautiful and abundant. With their hearts filled with hope, they held the magical acorn and made their wish.

As the mice closed their eyes and made their wish, a gentle breeze swept through the forest. When they opened their eyes, they couldn't believe what they saw. The forest had transformed into a magical wonderland, with vibrant flowers, sparkling streams, and towering trees that reached the sky.

The mice explored their enchanted forest, marveling at the beauty that surrounded them. The streams were filled with crystal-clear water, teeming with fish and other aquatic creatures. The trees bore fruits of all kinds, providing an abundance of food for the mice and other forest animals.

With their newfound paradise, the Forest Friends thrived. They lived in harmony with nature, sharing their blessings with other creatures. Oliver, as their wise leader, ensured that everyone had enough food and shelter. The mice worked together, building cozy burrows and gathering food for the winter.

Word of the magical forest spread far and wide, attracting animals from all corners of the land. The Forest Friends welcomed them with open arms, creating a diverse and vibrant community. The mice, with their kind hearts and generous spirits, became known as the Guardians of the Enchanted Forest.

As time passed, the Forest Friends continued to cherish their magical home. They lived in peace and harmony, always grateful for the gift they had received. Oliver, the wise leader, taught the younger mice the importance of unity and respect for nature.

And so, the story of the group of mice who discovered a magical acorn and transformed their forest home after adding their efforts by 121 + 8 became a tale of hope, gratitude, and the power of a shared dream. The Forest Friends lived happily ever after, forever grateful for the magic that had brought them together.

Agent with Personality

You can specify a system prompt to give the agent additional instruction or personality.

from llama_index.agent import OpenAIAgent
from llama_index.llms import OpenAI
from llama_index.prompts.system import SHAKESPEARE_WRITING_ASSISTANT

llm = OpenAI(model="gpt-3.5-turbo-0613")

agent = OpenAIAgent.from_tools(
    [multiply_tool, add_tool],
    llm=llm,
    verbose=True,
    system_prompt=SHAKESPEARE_WRITING_ASSISTANT,
)

response = agent.chat("Hi")
print(response)

Greetings, fair traveler! How may I assist thee on this fine day?

response = agent.chat("Tell me a story")
print(response)

Of course, dear friend! Allow me to weave a tale for thee in the style of Shakespeare. 

Once upon a time, in a land far away, there lived a noble knight named Sir William. He was known throughout the kingdom for his bravery and chivalry. One fateful day, as Sir William rode through the enchanted forest, he stumbled upon a hidden glade.

In the glade, he discovered a beautiful maiden named Lady Rosalind. She was fair of face and gentle of heart, and Sir William was instantly captivated by her beauty. They spent hours conversing, sharing stories, and laughing together.

As the days turned into weeks, Sir William and Lady Rosalind's bond grew stronger. They found solace in each other's company and discovered a love that was pure and true. However, their happiness was short-lived, for an evil sorcerer named Malachi had set his sights on Lady Rosalind.

Malachi, consumed by jealousy and darkness, sought to claim Lady Rosalind for himself. He devised a wicked plan to separate the two lovers and cast a spell upon Sir William, turning him into a statue of stone. Lady Rosalind, heartbroken and determined, vowed to find a way to break the curse and save her beloved.

With unwavering courage, Lady Rosalind embarked on a perilous journey to seek the help of a wise old wizard. She traveled through treacherous mountains, crossed raging rivers, and faced many trials along the way. Finally, after much hardship, she reached the wizard's humble abode.

The wizard, known as Merlin, listened to Lady Rosalind's tale of love and woe. He sympathized with her plight and agreed to aid her in breaking the curse. Together, they devised a plan to confront Malachi and restore Sir William to his human form.

On the eve of the full moon, Lady Rosalind and Merlin ventured into the heart of Malachi's lair. They faced countless obstacles and battled fierce creatures, but their determination never wavered. Finally, they reached the chamber where Sir William stood, frozen in stone.

With a wave of his staff and a powerful incantation, Merlin shattered the curse that held Sir William captive. As the first rays of dawn broke through the darkness, Sir William's eyes fluttered open, and he beheld Lady Rosalind standing before him.

Their love, stronger than ever, triumphed over the forces of evil. Sir William and Lady Rosalind returned to the kingdom, where they were hailed as heroes. They lived a long and joyous life together, their love serving as a beacon of hope for all who heard their tale.

And so, dear friend, ends the story of Sir William and Lady Rosalind, a tale of love, bravery, and the power of true devotion. May it inspire thee to seek love and adventure in thy own journey through life.