run-llama--llama_index/docs/examples/agent/openai_agent_retrieval.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "88216e84",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/agent/openai_agent_retrieval.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "99cea58c-48bc-4af6-8358-df9695659983",
   "metadata": {},
   "source": [
    "# Retrieval-Augmented Agents"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "673df1fe-eb6c-46ea-9a73-a96e7ae7942e",
   "metadata": {},
   "source": [
    "In this tutorial, we show you how to use our `FunctionAgent` or `ReActAgent` implementation with a tool retriever, \n",
    "to augment any existing agent and store/index an arbitrary number of tools. \n",
    "\n",
    "Our indexing/retrieval modules help to remove the complexity of having too many functions to fit in the prompt."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "54b7bc2e-606f-411a-9490-fcfab9236dfc",
   "metadata": {},
   "source": [
    "## Initial Setup "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "23e80e5b-aaee-4f23-b338-7ae62b08141f",
   "metadata": {},
   "source": [
    "Let's start by importing some simple building blocks.  \n",
    "\n",
    "The main thing we need is:\n",
    "1. the OpenAI API\n",
    "2. a place to keep conversation history \n",
    "3. a definition for tools that our agent can use."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4ce34f4d",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51441848",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d47283b-025e-4874-88ed-76245b22f82e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6fe08eb1-e638-4c00-9103-5c305bfacccf",
   "metadata": {},
   "source": [
    "Let's define some very simple calculator tools for our agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dd3c4a6-f3e0-46f9-ad3b-7ba57d1bc992",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.tools import FunctionTool\n",
    "\n",
    "\n",
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiply two integers and returns the result integer\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "def useless(a: int, b: int) -> int:\n",
    "    \"\"\"Toy useless function.\"\"\"\n",
    "    pass\n",
    "\n",
    "\n",
    "multiply_tool = FunctionTool.from_defaults(multiply, name=\"multiply\")\n",
    "add_tool = FunctionTool.from_defaults(add, name=\"add\")\n",
    "\n",
    "# toy-example of many tools\n",
    "useless_tools = [\n",
    "    FunctionTool.from_defaults(useless, name=f\"useless_{str(idx)}\")\n",
    "    for idx in range(28)\n",
    "]\n",
    "\n",
    "all_tools = [multiply_tool] + [add_tool] + useless_tools\n",
    "\n",
    "all_tools_map = {t.metadata.name: t for t in all_tools}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8170dd32-fa00-458c-aeb5-37292d0773c0",
   "metadata": {},
   "source": [
    "## Building an Object Index\n",
    "\n",
    "We have an `ObjectIndex` construct in LlamaIndex that allows the user to use our index data structures over arbitrary objects.\n",
    "The ObjectIndex will handle serialiation to/from the object, and use an underying index (e.g. VectorStoreIndex, SummaryIndex, KeywordTableIndex) as the storage mechanism. \n",
    "\n",
    "In this case, we have a large collection of Tool objects, and we'd want to define an ObjectIndex over these Tools.\n",
    "\n",
    "The index comes bundled with a retrieval mechanism, an `ObjectRetriever`. \n",
    "\n",
    "This can be passed in to our agent so that it can \n",
    "perform Tool retrieval during query-time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6704a755-7f05-43a3-8a56-f5f587ae4c40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define an \"object\" index over these tools\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.core.objects import ObjectIndex\n",
    "\n",
    "obj_index = ObjectIndex.from_objects(\n",
    "    all_tools,\n",
    "    index_cls=VectorStoreIndex,\n",
    "    # if we were using an external vector store, we could pass the stroage context and any other kwargs\n",
    "    # storage_context=storage_context,\n",
    "    # embed_model=embed_model,\n",
    "    # ...\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e13cc16f",
   "metadata": {},
   "source": [
    "To reload the index later, we can use the `from_objects_and_index` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56361dd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from llama_index.core import StorageContext, load_index_from_storage\n",
    "\n",
    "# saving and loading from disk\n",
    "# obj_index.index.storage_context.persist(persist_dir=\"obj_index_storage\")\n",
    "\n",
    "# reloading from disk\n",
    "# vector_index = load_index_from_storage(StorageContext.from_defaults(persist_dir=\"obj_index_storage\"))\n",
    "\n",
    "# or if using an external vector store, no need to persist, just reload the index\n",
    "# vector_index = VectorStoreIndex.from_vector_store(vector_store=vector_store, ...)\n",
    "\n",
    "# Then, we can reload the ObjectIndex\n",
    "# obj_index = ObjectIndex.from_objects_and_index(\n",
    "#     all_tools,\n",
    "#     index=vector_index,\n",
    "# )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "707d30b8-6405-4187-a9ed-6146dcc42167",
   "metadata": {},
   "source": [
    "## Agent w/ Tool Retrieval "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "798ca3fd-6711-4c0c-a853-d868dd14b484",
   "metadata": {},
   "source": [
    "Agents in LlamaIndex can be used with a `ToolRetriever` to retrieve tools during query-time.\n",
    "\n",
    "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. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38ab3938-1138-43ea-b085-f430b42f5377",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import FunctionAgent, ReActAgent\n",
    "from llama_index.core.workflow import Context\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    tool_retriever=obj_index.as_retriever(similarity_top_k=2),\n",
    "    llm=OpenAI(model=\"gpt-4o\"),\n",
    ")\n",
    "\n",
    "# context to hold the session/state\n",
    "ctx = Context(agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33ea069f-819b-4ec1-a93c-fcbaacb362a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The result of multiplying 212 by 122 is 25,864.\n",
      "[ToolCallResult(tool_name='multiply', tool_kwargs={'a': 212, 'b': 122}, tool_id='call_4Ygos3MpRH7Gj3R79HISRGyH', tool_output=ToolOutput(content='25864', tool_name='multiply', raw_input={'args': (), 'kwargs': {'a': 212, 'b': 122}}, raw_output=25864, is_error=False), return_direct=False)]\n"
     ]
    }
   ],
   "source": [
    "resp = await agent.run(\n",
    "    \"What's 212 multiplied by 122? Make sure to use Tools\", ctx=ctx\n",
    ")\n",
    "print(str(resp))\n",
    "print(resp.tool_calls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec423b90-59cd-40ef-b497-a3842b3e7b58",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The result of adding 212 to 122 is 334.\n",
      "[ToolCallResult(tool_name='add', tool_kwargs={'a': 212, 'b': 122}, tool_id='call_rXUfwQ477bcd6bxafQHgETaa', tool_output=ToolOutput(content='334', tool_name='add', raw_input={'args': (), 'kwargs': {'a': 212, 'b': 122}}, raw_output=334, is_error=False), return_direct=False)]\n"
     ]
    }
   ],
   "source": [
    "resp = await agent.run(\n",
    "    \"What's 212 added to 122 ? Make sure to use Tools\", ctx=ctx\n",
    ")\n",
    "print(str(resp))\n",
    "print(resp.tool_calls)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-caVs7DDe-py3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}