Retrieval-Augmented OpenAI Agent

In this tutorial, we show you how to use our FnRetrieverOpenAI implementation to build an agent on top of OpenAI’s function API and store/index an arbitrary number of tools. Our indexing/retrieval modules help to remove the complexity of having too many functions to fit in the prompt.

Initial Setup

Let’s start by importing some simple building blocks.

The main thing we need is:

1. the OpenAI API

2. a place to keep conversation history

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

If you’re opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

!pip install llama-index

import json
from typing import Sequence

from llama_index.tools import BaseTool, FunctionTool

/Users/suo/miniconda3/envs/llama/lib/python3.9/site-packages/deeplake/util/check_latest_version.py:32: UserWarning: A newer version of deeplake (3.6.7) is available. It's recommended that you update to the latest version using `pip install -U deeplake`.
  warnings.warn(

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

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


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


def useless(a: int, b: int) -> int:
    """Toy useless function."""
    pass


multiply_tool = FunctionTool.from_defaults(fn=multiply, name="multiply")
useless_tools = [
    FunctionTool.from_defaults(fn=useless, name=f"useless_{str(idx)}")
    for idx in range(28)
]
add_tool = FunctionTool.from_defaults(fn=add, name="add")

all_tools = [multiply_tool] + [add_tool] + useless_tools
all_tools_map = {t.metadata.name: t for t in all_tools}

Building an Object Index

We have an ObjectIndex construct in LlamaIndex that allows the user to use our index data structures over arbitrary objects. The ObjectIndex will handle serialiation to/from the object, and use an underying index (e.g. VectorStoreIndex, SummaryIndex, KeywordTableIndex) as the storage mechanism.

In this case, we have a large collection of Tool objects, and we’d want to define an ObjectIndex over these Tools.

The index comes bundled with a retrieval mechanism, an ObjectRetriever.

This can be passed in to our agent so that it can perform Tool retrieval during query-time.

# define an "object" index over these tools
from llama_index import VectorStoreIndex
from llama_index.objects import ObjectIndex, SimpleToolNodeMapping

tool_mapping = SimpleToolNodeMapping.from_objects(all_tools)
obj_index = ObjectIndex.from_objects(
    all_tools,
    tool_mapping,
    VectorStoreIndex,
)

Our FnRetrieverOpenAIAgent Implementation

We provide a FnRetrieverOpenAIAgent implementation in LlamaIndex, which can take in an ObjectRetriever over a set of BaseTool objects.

During query-time, we would first use the ObjectRetriever to retrieve a set of relevant Tools. These tools would then be passed into the agent; more specifically, their function signatures would be passed into the OpenAI Function calling API.

from llama_index.agent import FnRetrieverOpenAIAgent

agent = FnRetrieverOpenAIAgent.from_retriever(
    obj_index.as_retriever(), verbose=True
)

agent.chat("What's 212 multiplied by 122? Make sure to use Tools")

=== Calling Function ===
Calling function: multiply with args: {
  "a": 212,
  "b": 122
}
Got output: 25864
========================

Response(response='212 multiplied by 122 is 25,864.', source_nodes=[], metadata=None)

agent.chat("What's 212 added to 122 ? Make sure to use Tools")

=== Calling Function ===
Calling function: add with args: {
  "a": 212,
  "b": 122
}
Got output: 334
========================

Response(response='212 added to 122 is 334.', source_nodes=[], metadata=None)